diff --git a/.github/workflows/sycl-cpu-ci.yml b/.github/workflows/sycl-cpu-ci.yml new file mode 100644 index 000000000..a24210aef --- /dev/null +++ b/.github/workflows/sycl-cpu-ci.yml @@ -0,0 +1,79 @@ +name: Build and Test - SYCL CPU + +on: + pull_request: + types: [opened, edited, synchronize, labeled] + + workflow_dispatch: + +env: + GH_ACTIONS_TOOLCHAIN: .github/workflows/toolchains/gh-actions.cmake + ENV_PREFIX_PATH: "/opt/intel/oneapi/mkl/latest" + +jobs: + sycl_cpu_build: + name: SYCL CPU Build and Test + if: > + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && github.event.action == 'labeled' && github.event.label.name == 'SYCL') || + (github.event_name == 'pull_request' && github.event.action != 'labeled' && contains(github.event.pull_request.title, '[SYCL]')) + runs-on: ubuntu-latest + container: + image: intel/oneapi-basekit:2024.2.1-0-devel-ubuntu22.04 + timeout-minutes: 60 + + steps: + - uses: actions/checkout@v4 + + - name: Install HDF5 development headers + shell: bash + run: | + apt-get update + apt-get install -y libhdf5-dev + + - name: Setup Compiler + shell: bash + run: | + echo "set(CMAKE_C_COMPILER icx CACHE STRING \"\" FORCE)" >> ${GITHUB_WORKSPACE}/${GH_ACTIONS_TOOLCHAIN} + echo "set(CMAKE_CXX_COMPILER icpx CACHE STRING \"\" FORCE)" >> ${GITHUB_WORKSPACE}/${GH_ACTIONS_TOOLCHAIN} + + - name: Setup Build Type + shell: bash + run: echo "set(CMAKE_BUILD_TYPE Release CACHE BOOL \"\" FORCE)" >> + ${GITHUB_WORKSPACE}/${GH_ACTIONS_TOOLCHAIN} + + - name: Setup Enable SYCL + shell: bash + run: | + echo "set(GAUXC_ENABLE_SYCL ON CACHE BOOL \"\" FORCE)" >> ${GITHUB_WORKSPACE}/${GH_ACTIONS_TOOLCHAIN} + echo "set(GAUXC_SYCL_TARGET spir64_x86_64 CACHE STRING \"\" FORCE)" >> ${GITHUB_WORKSPACE}/${GH_ACTIONS_TOOLCHAIN} + echo "set(GAUXC_ENABLE_MPI OFF CACHE BOOL \"\" FORCE)" >> ${GITHUB_WORKSPACE}/${GH_ACTIONS_TOOLCHAIN} + + - name: Configure CMake + shell: bash + run: cmake -S $GITHUB_WORKSPACE -B ${{runner.workspace}}/build + -DCMAKE_INSTALL_PREFIX=${{runner.workspace}}/install + -DCMAKE_PREFIX_PATH=${ENV_PREFIX_PATH} + -DCMAKE_TOOLCHAIN_FILE=${GITHUB_WORKSPACE}/${GH_ACTIONS_TOOLCHAIN} + + - name: Build + shell: bash + timeout-minutes: 35 + run: cmake --build ${{runner.workspace}}/build -j2 + + - name: List SYCL devices + shell: bash + run: sycl-ls || true + + # The SYCL kernels run on the OpenCL CPU device here, which is orders of + # magnitude slower than a GPU: the XC integrator sections in particular are + # minutes apiece on a many-core host and far worse on a 2-core runner. Cap + # the step so a hang is reported against the test rather than burning the + # job's whole budget, and give ctest its own per-test timeout so a single + # wedged case cannot stall the rest. + - name: Test + shell: bash + timeout-minutes: 20 + env: + ONEAPI_DEVICE_SELECTOR: opencl:cpu + run: ctest --test-dir ${{runner.workspace}}/build --output-on-failure --timeout 900 diff --git a/CMakeLists.txt b/CMakeLists.txt index e5459eb51..f040d7a7c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,6 +31,7 @@ option( GAUXC_ENABLE_C "Enable C API" ON ) option( GAUXC_ENABLE_HOST "Enable Host Integrator" ON ) option( GAUXC_ENABLE_CUDA "Enable CUDA Bindings" OFF ) option( GAUXC_ENABLE_HIP "Enable HIP Bindings" OFF ) +option( GAUXC_ENABLE_SYCL "Enable SYCL Bindings" OFF ) option( GAUXC_ENABLE_MPI "Enable MPI Bindings" ON ) option( GAUXC_ENABLE_OPENMP "Enable OpenMP Compilation" ON ) option( GAUXC_ENABLE_TESTS "Enable Unit Tests" ON ) @@ -41,10 +42,25 @@ option( GAUXC_BLAS_PREFER_ILP64 "Prefer ILP64 for host BLAS" OFF ) option( GAUXC_LINK_CUDA_STATIC "Link GauXC with static CUDA libs" OFF ) include(CMakeDependentOption) + +# MAGMA has no released SYCL backend: upstream master carries only the CUDA and +# HIP backends, and the SYCL work lives on an unmerged branch. Rather than +# silently dropping the request, fail loudly if the user asked for both. +if( GAUXC_ENABLE_SYCL AND GAUXC_ENABLE_MAGMA ) + message( FATAL_ERROR + "GAUXC_ENABLE_MAGMA=ON is not supported with GAUXC_ENABLE_SYCL=ON: MAGMA " + "provides no SYCL backend. Batched linear algebra for the SYCL backend is " + "served by oneMKL instead; configure with -DGAUXC_ENABLE_MAGMA=OFF." ) +endif() + cmake_dependent_option( GAUXC_ENABLE_MAGMA "Enable MAGMA Linear Algebra" ON "GAUXC_ENABLE_CUDA OR GAUXC_ENABLE_HIP" OFF ) +cmake_dependent_option( GAUXC_ENABLE_ONECCL + "Enable oneCCL Collectives" OFF + "GAUXC_ENABLE_SYCL" OFF +) cmake_dependent_option( GAUXC_ENABLE_NCCL "Enable NCCL Collectives" OFF "GAUXC_ENABLE_CUDA" OFF @@ -59,12 +75,14 @@ set( GAUXC_HAS_C FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_HOST FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_CUDA FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_HIP FALSE CACHE BOOL "" FORCE ) +set( GAUXC_HAS_SYCL FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_MPI FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_OPENMP FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_GAU2GRID FALSE CACHE STRING "" FORCE ) set( GAUXC_HAS_HDF5 FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_MAGMA FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_NCCL FALSE CACHE BOOL "" FORCE ) +set( GAUXC_HAS_ONECCL FALSE CACHE BOOL "" FORCE ) set( GAUXC_HAS_CUTLASS FALSE CACHE BOOL "" FORCE ) set( GAUXC_BLAS_IS_LP64 FALSE CACHE BOOL "" FORCE ) @@ -73,12 +91,14 @@ mark_as_advanced( FORCE GAUXC_HAS_HOST GAUXC_HAS_CUDA GAUXC_HAS_HIP + GAUXC_HAS_SYCL GAUXC_HAS_MPI GAUXC_HAS_OPENMP GAUXC_HAS_GAU2GRID GAUXC_HAS_HDF5 GAUXC_HAS_MAGMA GAUXC_HAS_NCCL + GAUXC_HAS_ONECCL GAUXC_HAS_CUTLASS GAUXC_BLAS_IS_LP64 ) @@ -104,8 +124,44 @@ if( GAUXC_ENABLE_HIP ) set( GAUXC_HAS_HIP TRUE CACHE BOOL "GauXC has HIP and will build HIP bindings" FORCE ) endif() +if( GAUXC_ENABLE_SYCL ) + # SYCL is compiled through the CXX compiler; no separate language to enable + set( GAUXC_HAS_SYCL TRUE CACHE BOOL "GauXC has SYCL and will build SYCL bindings" FORCE ) + + # e.g. intel_gpu_pvc | spir64_x86_64 | nvidia_gpu_sm_80 | amd_gpu_gfx90a + # Declared here (rather than in gauxc-sycl) so that it is visible to the + # ExchCXX pass-through, which is configured before the gauxc target exists. + # + # Defaults to ahead-of-time for PVC: that is how these builds are deployed, + # and a JIT build pays the whole device-image compile on the first kernel + # launch of every run, which is easily minutes for the collocation and + # ExchCXX images. Set to an empty string explicitly for a JIT build. + set( GAUXC_SYCL_TARGET "intel_gpu_pvc" CACHE STRING + "Alias for -fsycl-targets (see Users Manual); empty for a JIT build" ) + + # SYCL device code shares the CXX FP model with the host, and ExchCXX's SYCL + # bindings require IEEE semantics for agreement with the reference values + include( CheckCXXCompilerFlag ) + check_cxx_compiler_flag( "-fp-model=precise" GAUXC_CXX_HAS_FP_MODEL_PRECISE ) + if( GAUXC_CXX_HAS_FP_MODEL_PRECISE ) + set( CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fp-model=precise" ) + endif() + + # Claim SYCL::SYCL before ExchCXX is fetched. ExchCXX ships its own FindSYCL + # which attaches -fsycl as an INTERFACE_COMPILE_OPTION; were that definition + # to win, every GauXC translation unit would run the SYCL frontend, including + # the large host-only sources where it buys nothing. GauXC's own FindSYCL + # keeps the target link-only and stamps -fsycl per source instead. + find_package( SYCL REQUIRED ) +endif() + +# Device backends are mutually exclusive +if( (GAUXC_HAS_CUDA AND GAUXC_HAS_SYCL) OR (GAUXC_HAS_HIP AND GAUXC_HAS_SYCL) ) + message( FATAL_ERROR "SYCL bindings are mutually exclusive with CUDA / HIP" ) +endif() + # Decided if we're compiling device bindings -if( GAUXC_HAS_CUDA OR GAUXC_HAS_HIP ) +if( GAUXC_HAS_CUDA OR GAUXC_HAS_HIP OR GAUXC_HAS_SYCL ) set( GAUXC_HAS_DEVICE TRUE CACHE BOOL "Enable Device Code" ) else() set( GAUXC_HAS_DEVICE FALSE CACHE BOOL "Enable Device Code" ) diff --git a/README.md b/README.md index 3bbba6738..e34164efa 100644 --- a/README.md +++ b/README.md @@ -202,12 +202,15 @@ target_link_libraries( my_target PUBLIC gauxc::gauxc ) | `GAUXC_ENABLE_HOST` | Enable HOST integrators | `ON` | | `GAUXC_ENABLE_CUDA` | Enable CUDA integrators | `OFF` | | `GAUXC_ENABLE_HIP` | Enable HIP integrators | `OFF` | +| `GAUXC_ENABLE_SYCL` | Enable SYCL integrators | `OFF` | | `GAUXC_ENABLE_MAGMA` | Enable MAGMA for batched BLAS (No effect if no GPU) | `ON` | | `GAUXC_ENABLE_CUTLASS` | Enable CUTLASS for batched BLAS (No effect if no CUDA) | `OFF` | | `GAUXC_ENABLE_NCCL` | Enable NCCL bindings for topology aware GPU reductions | `OFF` | +| `GAUXC_ENABLE_ONECCL` | Enable oneCCL bindings for GPU reductions (SYCL only) | `OFF` | | `GAUXC_ENABLE_MPI` | Enable MPI Bindings | `ON` | | `GAUXC_ENABLE_OPENMP` | Enable OpenMP Bindings | `ON` | | `CMAKE_CUDA_ARCHITECTURES` | CUDA architechtures (e.g. 70 for Volta, 80 for Ampere) | -- | +| `GAUXC_SYCL_TARGET` | SYCL AoT target (e.g. `intel_gpu_pvc`); JIT if unset | -- | | `BLAS_LIBRARIES` | Full BLAS linker. | -- | | `MAGMA_ROOT_DIR` | Install prefix for MAGMA. | -- | diff --git a/cmake/gauxc-config.cmake.in b/cmake/gauxc-config.cmake.in index 50230f9a6..0771f8feb 100644 --- a/cmake/gauxc-config.cmake.in +++ b/cmake/gauxc-config.cmake.in @@ -14,8 +14,10 @@ set( GAUXC_HAS_C @GAUXC_HAS_C@ ) set( GAUXC_HAS_HOST @GAUXC_HAS_HOST@ ) set( GAUXC_HAS_CUDA @GAUXC_HAS_CUDA@ ) set( GAUXC_HAS_HIP @GAUXC_HAS_HIP@ ) +set( GAUXC_HAS_SYCL @GAUXC_HAS_SYCL@ ) set( GAUXC_HAS_MAGMA @GAUXC_HAS_MAGMA@ ) set( GAUXC_HAS_NCCL @GAUXC_HAS_NCCL@ ) +set( GAUXC_HAS_ONECCL @GAUXC_HAS_ONECCL@ ) set( GAUXC_HAS_CUTLASS @GAUXC_HAS_CUTLASS@ ) set( GAUXC_HAS_MPI @GAUXC_HAS_MPI@ ) set( GAUXC_HAS_OPENMP @GAUXC_HAS_OPENMP@ ) @@ -56,6 +58,25 @@ if( GAUXC_HAS_CUDA ) endif() endif() +if( GAUXC_HAS_SYCL ) + find_dependency( SYCL ) + # MKLConfig defaults to ILP64 and keeps separate cache entries for the host + # and SYCL domains; both must match the interface GauXC was built against or + # the duplicate dgemm_ binds to the wrong copy. + if( GAUXC_BLAS_IS_LP64 ) + set( _gauxc_mkl_interface intel_lp64 ) + else() + set( _gauxc_mkl_interface intel_ilp64 ) + endif() + set( MKL_INTERFACE_FULL ${_gauxc_mkl_interface} CACHE STRING "" FORCE ) + set( MKL_SYCL_INTERFACE_FULL ${_gauxc_mkl_interface} CACHE STRING "" FORCE ) + unset( _gauxc_mkl_interface ) + find_dependency( MKL CONFIG ) + if( GAUXC_HAS_ONECCL ) + find_dependency( oneCCL ) + endif() +endif() + if( GAUXC_HAS_MPI ) find_dependency( MPI ) endif() diff --git a/cmake/gauxc-exchcxx.cmake b/cmake/gauxc-exchcxx.cmake index 412df9b30..ccae1d39d 100644 --- a/cmake/gauxc-exchcxx.cmake +++ b/cmake/gauxc-exchcxx.cmake @@ -9,8 +9,13 @@ if( NOT ${ExchCXX_FOUND} ) set( EXCHCXX_ENABLE_CUDA ${GAUXC_HAS_CUDA} CACHE BOOL "" ) set( EXCHCXX_ENABLE_HIP ${GAUXC_HAS_HIP} CACHE BOOL "" ) + set( EXCHCXX_ENABLE_SYCL ${GAUXC_HAS_SYCL} CACHE BOOL "" ) set( EXCHCXX_ENABLE_TESTS OFF CACHE BOOL "" ) + if( GAUXC_HAS_SYCL AND GAUXC_SYCL_TARGET ) + set( EXCHCXX_SYCL_TARGET ${GAUXC_SYCL_TARGET} CACHE STRING "" ) + endif() + FetchContent_Declare( exchcxx GIT_REPOSITORY ${GAUXC_EXCHCXX_REPOSITORY} @@ -30,6 +35,10 @@ else() message( FATAL_ERROR "GauXC HIP BINDINGS REQUIRE ExchCXX HIP Bindings" ) endif() + if( ${GAUXC_HAS_SYCL} AND NOT ${EXCHCXX_ENABLE_SYCL} ) + message( FATAL_ERROR "GauXC SYCL BINDINGS REQUIRE ExchCXX SYCL Bindings" ) + endif() + endif() diff --git a/cmake/gauxc-sycl.cmake b/cmake/gauxc-sycl.cmake new file mode 100644 index 000000000..dcb19f4f0 --- /dev/null +++ b/cmake/gauxc-sycl.cmake @@ -0,0 +1,184 @@ +# +# GauXC Copyright (c) 2020-2024, The Regents of the University of California, +# through Lawrence Berkeley National Laboratory (subject to receipt of +# any required approvals from the U.S. Dept. of Energy). +# +# (c) 2024-2025, Microsoft Corporation +# +# All rights reserved. +# +# See LICENSE.txt for details +# +find_package( SYCL REQUIRED ) + +target_link_libraries( gauxc PUBLIC SYCL::SYCL ) +target_link_options( gauxc PUBLIC -fsycl ) + +# oneMKL supplies the device BLAS; the host BLAS is BLAS::BLAS, found +# separately. Host and device must agree on the integer interface or the +# duplicate dgemm_ binds to the wrong copy. Re-pinned from the host driver +# once BLAS_IS_LP64 is discovered. +function( gauxc_sycl_pin_mkl_interface _is_lp64 ) + if( _is_lp64 ) + set( _iface intel_lp64 ) + else() + set( _iface intel_ilp64 ) + endif() + set( MKL_INTERFACE_FULL ${_iface} CACHE STRING "" FORCE ) + set( MKL_SYCL_INTERFACE_FULL ${_iface} CACHE STRING "" FORCE ) +endfunction() + +# Provisional; re-pinned from the discovered value in the host driver. +if( GAUXC_BLAS_PREFER_ILP64 ) + gauxc_sycl_pin_mkl_interface( FALSE ) +else() + gauxc_sycl_pin_mkl_interface( TRUE ) +endif() + +if( NOT TARGET MKL::MKL_SYCL::BLAS ) + find_package( MKL CONFIG REQUIRED ) +endif() + +if( NOT TARGET MKL::MKL_SYCL::BLAS ) + message( FATAL_ERROR + "oneMKL was found but does not provide MKL::MKL_SYCL::BLAS, which the " + "GauXC SYCL backend requires for its device BLAS." ) +endif() + +# MKLConfig also puts -fsycl in INTERFACE_COMPILE_OPTIONS; strip it for the +# same reason as SYCL::SYCL above. Wrapped in a generator expression, hence +# the regex match. +get_target_property( _mkl_blas_copts MKL::MKL_SYCL::BLAS INTERFACE_COMPILE_OPTIONS ) +if( _mkl_blas_copts ) + set( _mkl_blas_copts_keep ) + foreach( _opt ${_mkl_blas_copts} ) + if( NOT _opt MATCHES "(^|[:;])-fsycl($|[;>])" ) + list( APPEND _mkl_blas_copts_keep ${_opt} ) + endif() + endforeach() + set_target_properties( MKL::MKL_SYCL::BLAS PROPERTIES + INTERFACE_COMPILE_OPTIONS "${_mkl_blas_copts_keep}" ) +endif() + +target_link_libraries( gauxc PUBLIC $ ) + +if( GAUXC_ENABLE_ONECCL ) + find_package( oneCCL REQUIRED ) + target_link_libraries( gauxc PUBLIC oneCCL ) + set( GAUXC_HAS_ONECCL TRUE CACHE BOOL "GauXC has oneCCL" FORCE ) +endif() + + +# AoT target aliases, mirroring ExchCXX's mapping +set( _GAUXC_SYCL_ALLOWED + intel_gpu_pvc + spir64_x86_64 + nvidia_gpu_sm_80 + nvidia_gpu_sm_90 + amd_gpu_gfx90a + amd_gpu_gfx942 +) + +if( DEFINED GAUXC_SYCL_TARGET AND NOT GAUXC_SYCL_TARGET STREQUAL "" ) + list( FIND _GAUXC_SYCL_ALLOWED "${GAUXC_SYCL_TARGET}" _gauxc_sycl_idx ) + if( _gauxc_sycl_idx EQUAL -1 ) + message( FATAL_ERROR + "Invalid GAUXC_SYCL_TARGET='${GAUXC_SYCL_TARGET}'. " + "Allowed values: ${_GAUXC_SYCL_ALLOWED}" ) + endif() + + unset( _gauxc_sycl_compile_opts ) + unset( _gauxc_sycl_link_opts ) + + if( GAUXC_SYCL_TARGET STREQUAL "intel_gpu_pvc" ) + # SOURCE COMPILE_OPTIONS does not honor SHELL:, so keep tokens separate. + list( APPEND _gauxc_sycl_compile_opts + -fsycl-targets=spir64_gen + -Xsycl-target-backend=spir64_gen + "-device pvc" + ) + list( APPEND _gauxc_sycl_link_opts + "SHELL:-ftarget-register-alloc-mode=pvc:large" + "SHELL:-fsycl-targets=spir64_gen" + "SHELL:-Xsycl-target-backend \"-device pvc\"" + ) + elseif( GAUXC_SYCL_TARGET STREQUAL "spir64_x86_64" ) + list( APPEND _gauxc_sycl_compile_opts -fsycl-targets=spir64_x86_64 ) + list( APPEND _gauxc_sycl_link_opts "SHELL:-fsycl-targets=spir64_x86_64" ) + elseif( GAUXC_SYCL_TARGET MATCHES "^nvidia_gpu_" OR + GAUXC_SYCL_TARGET MATCHES "^amd_gpu_" ) + list( APPEND _gauxc_sycl_compile_opts -fsycl-targets=${GAUXC_SYCL_TARGET} ) + list( APPEND _gauxc_sycl_link_opts "SHELL:-fsycl-targets=${GAUXC_SYCL_TARGET}" ) + endif() + + if( NOT _gauxc_sycl_compile_opts ) + message( FATAL_ERROR + "GAUXC_SYCL_TARGET='${GAUXC_SYCL_TARGET}' is allowed but has no AoT flag " + "mapping in ${CMAKE_CURRENT_LIST_FILE}; refusing to fall back to a JIT " + "build that would be reported as AoT." ) + endif() + + # Compile options are applied per source; link options are PUBLIC because + # gauxc is a static archive and device linking happens in the consumer. + set( GAUXC_SYCL_AOT_COMPILE_OPTIONS ${_gauxc_sycl_compile_opts} ) + target_link_options( gauxc PUBLIC ${_gauxc_sycl_link_opts} ) + + message( STATUS "GauXC SYCL AoT enabled for target: ${GAUXC_SYCL_TARGET}" ) +endif() + +# Mark sources as carrying SYCL device code: only these get -fsycl and the +# device-code flags. +function( gauxc_sycl_device_sources ) + # TARGET_DIRECTORY is required: SOURCE properties are directory-scoped and + # these sources are added from subdirectories, gauxc is defined in src/. + foreach( _src ${ARGN} ) + get_filename_component( _abs ${_src} ABSOLUTE ) + set_property( SOURCE ${_abs} TARGET_DIRECTORY gauxc + APPEND PROPERTY COMPILE_OPTIONS + ${SYCL_DEVICE_COMPILE_OPTIONS} + ${GAUXC_SYCL_DEVICE_COMPILE_OPTIONS} + ${GAUXC_SYCL_AOT_COMPILE_OPTIONS} ) + endforeach() +endfunction() + + +include( CheckCXXCompilerFlag ) +check_cxx_compiler_flag( "-fno-sycl-id-queries-fit-in-int" GAUXC_SYCL_ID_QUERIES_FIT_IN_INT ) +check_cxx_compiler_flag( "-fsycl-device-code-split=per_kernel" GAUXC_SYCL_DEVICE_CODE_SPLIT_PER_KERNEL ) +check_cxx_compiler_flag( "-Xsycl-target-frontend \"-fp-model=precise\"" GAUXC_HAVE_SYCL_TARGET_FRONTEND_FP_MODEL_PRECISE ) + +include( CheckLinkerFlag ) +check_linker_flag( CXX "-flink-huge-device-code" GAUXC_SYCL_LINK_HUGE_DEVICE_CODE ) +check_linker_flag( CXX "-fsycl-max-parallel-link-jobs=4" GAUXC_SYCL_MAX_PARALLEL_LINK_JOBS ) + +# GauXC indexes into task batches which routinely exceed INT_MAX +if( GAUXC_SYCL_ID_QUERIES_FIT_IN_INT ) + list( APPEND GAUXC_SYCL_DEVICE_COMPILE_OPTIONS -fno-sycl-id-queries-fit-in-int ) +endif() + +# ~50 collocation kernels alone; splitting keeps device image size manageable +if( GAUXC_SYCL_DEVICE_CODE_SPLIT_PER_KERNEL ) + list( APPEND GAUXC_SYCL_DEVICE_COMPILE_OPTIONS -fsycl-device-code-split=per_kernel ) + target_link_options( gauxc PUBLIC + $<$:-fsycl-device-code-split=per_kernel> + ) +endif() + +# Device-side FP model; host side is set project-wide in the top-level file. +if( GAUXC_HAVE_SYCL_TARGET_FRONTEND_FP_MODEL_PRECISE ) + # Tokens kept separate: SOURCE properties do not honor SHELL:. + list( APPEND GAUXC_SYCL_DEVICE_COMPILE_OPTIONS + -Xsycl-target-frontend -fp-model=precise ) +endif() + +if( GAUXC_SYCL_LINK_HUGE_DEVICE_CODE ) + target_link_options( gauxc PUBLIC + $<$:-flink-huge-device-code> + ) +endif() + +if( GAUXC_SYCL_MAX_PARALLEL_LINK_JOBS ) + target_link_options( gauxc PUBLIC + $<$:-fsycl-max-parallel-link-jobs=4> + ) +endif() diff --git a/cmake/modules/FindSYCL.cmake b/cmake/modules/FindSYCL.cmake new file mode 100644 index 000000000..f4d4623db --- /dev/null +++ b/cmake/modules/FindSYCL.cmake @@ -0,0 +1,30 @@ +# +# GauXC Copyright (c) 2020-2024, The Regents of the University of California, +# through Lawrence Berkeley National Laboratory (subject to receipt of +# any required approvals from the U.S. Dept. of Energy). +# +# (c) 2024-2025, Microsoft Corporation +# +# All rights reserved. +# +# See LICENSE.txt for details +# +include( FindPackageHandleStandardArgs ) +include( CheckCXXCompilerFlag ) +check_cxx_compiler_flag("-fsycl" CXX_HAS_FSYCL) + +find_package_handle_standard_args( SYCL + REQUIRED_VARS CXX_HAS_FSYCL +) + +if( SYCL_FOUND AND NOT TARGET SYCL::SYCL ) + add_library( SYCL::SYCL INTERFACE IMPORTED ) + set_target_properties( SYCL::SYCL PROPERTIES + INTERFACE_COMPILE_OPTIONS "$<$:-fsycl>" + INTERFACE_LINK_OPTIONS "-fsycl" + ) +endif() + +# Flags every translation unit holding SYCL device code must be compiled with. +# GauXC applies these per source rather than target-wide; see gauxc-sycl.cmake. +set( SYCL_DEVICE_COMPILE_OPTIONS -fsycl ) diff --git a/include/gauxc/c/gauxc_config.h.in b/include/gauxc/c/gauxc_config.h.in index 992d1b7c2..c106a3943 100644 --- a/include/gauxc/c/gauxc_config.h.in +++ b/include/gauxc/c/gauxc_config.h.in @@ -15,6 +15,7 @@ #cmakedefine GAUXC_HAS_HOST #cmakedefine GAUXC_HAS_CUDA #cmakedefine GAUXC_HAS_HIP +#cmakedefine GAUXC_HAS_SYCL #cmakedefine GAUXC_HAS_MPI #cmakedefine GAUXC_HAS_MAGMA #cmakedefine GAUXC_HAS_NCCL diff --git a/include/gauxc/exceptions.hpp b/include/gauxc/exceptions.hpp index efd60fbeb..6850d2e24 100644 --- a/include/gauxc/exceptions.hpp +++ b/include/gauxc/exceptions.hpp @@ -28,6 +28,11 @@ class hip_exception; class hipblas_exception; #endif +#ifdef GAUXC_HAS_SYCL +class sycl_exception; +class onemkl_exception; +#endif + #ifdef GAUXC_HAS_MAGMA class magma_exception; #endif diff --git a/include/gauxc/gauxc_config.hpp.in b/include/gauxc/gauxc_config.hpp.in index 94d8f11ae..afbb493f1 100644 --- a/include/gauxc/gauxc_config.hpp.in +++ b/include/gauxc/gauxc_config.hpp.in @@ -15,9 +15,11 @@ #cmakedefine GAUXC_HAS_HOST #cmakedefine GAUXC_HAS_CUDA #cmakedefine GAUXC_HAS_HIP +#cmakedefine GAUXC_HAS_SYCL #cmakedefine GAUXC_HAS_MPI #cmakedefine GAUXC_HAS_MAGMA #cmakedefine GAUXC_HAS_NCCL +#cmakedefine GAUXC_HAS_ONECCL #cmakedefine GAUXC_HAS_CUTLASS #cmakedefine GAUXC_HAS_GAU2GRID #cmakedefine GAUXC_HAS_HDF5 diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 619689239..7de3acf7d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -104,6 +104,10 @@ target_link_libraries( gauxc PUBLIC ExchCXX::ExchCXX IntegratorXX::IntegratorXX ) + +if( GAUXC_HAS_SYCL ) + include( gauxc-sycl ) +endif() if( GAUXC_HAS_OPENMP AND TARGET OpenMP::OpenMP_CXX ) target_link_libraries( gauxc PUBLIC OpenMP::OpenMP_CXX ) else() @@ -123,6 +127,14 @@ add_subdirectory( load_balancer ) add_subdirectory( reduction_driver ) add_subdirectory( external ) +# oneMKL serves both the host and device BLAS for the SYCL backend, so its +# integer interface has to track the one the host BLAS resolved to. Deferred +# to here because GAUXC_BLAS_IS_LP64 is only known once the host driver has +# been configured above. +if( GAUXC_HAS_SYCL ) + gauxc_sycl_pin_mkl_interface( ${GAUXC_BLAS_IS_LP64} ) +endif() + add_library( gauxc::gauxc ALIAS gauxc ) @@ -162,9 +174,11 @@ set(export_properties "GAUXC_HAS_DEVICE" "GAUXC_HAS_CUDA" "GAUXC_HAS_HIP" + "GAUXC_HAS_SYCL" "GAUXC_HAS_MAGMA" "GAUXC_HAS_CUTLASS" "GAUXC_HAS_NCCL" + "GAUXC_HAS_ONECCL" "GAUXC_HAS_MPI" "GAUXC_HAS_OPENMP" "GAUXC_HAS_HDF5" @@ -183,9 +197,11 @@ set_target_properties(gauxc "GAUXC_HAS_DEVICE" ${GAUXC_HAS_DEVICE} "GAUXC_HAS_CUDA" ${GAUXC_HAS_CUDA} "GAUXC_HAS_HIP" ${GAUXC_HAS_HIP} + "GAUXC_HAS_SYCL" ${GAUXC_HAS_SYCL} "GAUXC_HAS_MAGMA" ${GAUXC_HAS_MAGMA} "GAUXC_HAS_CUTLASS" ${GAUXC_HAS_CUTLASS} "GAUXC_HAS_NCCL" ${GAUXC_HAS_NCCL} + "GAUXC_HAS_ONECCL" ${GAUXC_HAS_ONECCL} "GAUXC_HAS_MPI" ${GAUXC_HAS_MPI} "GAUXC_HAS_OPENMP" ${GAUXC_HAS_OPENMP} "GAUXC_HAS_HDF5" ${GAUXC_HAS_HDF5} @@ -243,6 +259,7 @@ configure_package_config_file( install( FILES ${PROJECT_SOURCE_DIR}/cmake/modules/FindMAGMA.cmake ${PROJECT_SOURCE_DIR}/cmake/modules/FindNCCL.cmake + ${PROJECT_SOURCE_DIR}/cmake/modules/FindSYCL.cmake ${CMAKE_CURRENT_BINARY_DIR}/gauxc-config.cmake ${CMAKE_CURRENT_BINARY_DIR}/gauxc-config-version.cmake DESTINATION ${INSTALL_CONFIGDIR} diff --git a/src/exceptions/onemkl_exception.hpp b/src/exceptions/onemkl_exception.hpp new file mode 100644 index 000000000..38c637caf --- /dev/null +++ b/src/exceptions/onemkl_exception.hpp @@ -0,0 +1,81 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include +#include +#include + +#ifdef GAUXC_HAS_SYCL +#include +// Only the exception type is needed here. The umbrella pulls +// in the C LAPACK prototypes, whose parameter named "csl" collides with the +// csl macro from buffer_adaptor.hpp +#include + +namespace GauXC { + +/** + * @brief A class to handle excecptions arising from oneMKL operations + */ +class onemkl_exception : public std::exception { + + std::string file_; ///< File which contains the code that threw the exception + int line_; ///< Line number of file_ that threw exception + std::string msg_prefix_; ///< General descriptor of task which threw exception + std::string err_str_; ///< oneMKL error string pertaining to the thrown exception + std::string what_msg_; + + /** + * @brief Get a descriptive message pertaining to the thrown oneMKL error + * + * @returns a descritive message pertaining to the oneMKL error represented + * by the internal state of the exception object. + */ + const char* what() const noexcept override { + return what_msg_.c_str(); + } + +public: + + /** + * @brief Construct a onemkl_exception object + * + * @param[in] file File which contains the code that threw the exception + * @param[in] line Line number of file that threw exception + * @param[in] msg General descriptor of task which threw exception + * @param[in] err oneMKL exception pertaining to the thrown exception + */ + onemkl_exception( std::string file, int line, std::string msg, + const std::exception& err ) : + file_(file), line_(line), msg_prefix_(msg), err_str_(err.what()) { + std::stringstream ss; + ss << "oneMKL Exception (" << msg_prefix_ << ")" << std::endl + << " Error \"" << err_str_ << "\"" << std::endl + << " File " << file_ << std::endl + << " Line " << line_ << std::endl; + what_msg_ = ss.str(); + } + +}; // class onemkl_exception + +} // namespace GauXC + +// Macro to wrap oneMKL error handling +#define GAUXC_ONEMKL_ERROR( MSG, EXPR ) \ + try { EXPR; } \ + catch( const oneapi::mkl::exception& e ) { \ + throw onemkl_exception( __FILE__, __LINE__, MSG, e ); \ + } + +#endif diff --git a/src/exceptions/sycl_exception.hpp b/src/exceptions/sycl_exception.hpp new file mode 100644 index 000000000..db34314a6 --- /dev/null +++ b/src/exceptions/sycl_exception.hpp @@ -0,0 +1,77 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include +#include +#include + +#ifdef GAUXC_HAS_SYCL +#include + +namespace GauXC { + +/** + * @brief A class to handle excecptions arising from SYCL operations + */ +class sycl_exception : public std::exception { + + std::string file_; ///< File which contains the code that threw the exception + int line_; ///< Line number of file_ that threw exception + std::string msg_prefix_; ///< General descriptor of task which threw exception + std::string err_str_; ///< SYCL error string pertaining to the thrown exception + std::string what_msg_; + + /** + * @brief Get a descriptive message pertaining to the thrown SYCL error + * + * @returns a descritive message pertaining to the SYCL error represented by + * the internal state of the exception object. + */ + const char* what() const noexcept override { + return what_msg_.c_str(); + } + +public: + + /** + * @brief Construct a sycl_exception object + * + * @param[in] file File which contains the code that threw the exception + * @param[in] line Line number of file that threw exception + * @param[in] msg General descriptor of task which threw exception + * @param[in] err SYCL exception pertaining to the thrown exception + */ + sycl_exception( std::string file, int line, std::string msg, + const ::sycl::exception& err ) : + file_(file), line_(line), msg_prefix_(msg), err_str_(err.what()) { + std::stringstream ss; + ss << "SYCL Exception (" << msg_prefix_ << ")" << std::endl + << " Error \"" << err_str_ << "\"" << std::endl + << " File " << file_ << std::endl + << " Line " << line_ << std::endl; + what_msg_ = ss.str(); + } + +}; // class sycl_exception + +} // namespace GauXC + +// Macro to wrap SYCL error handling +#define GAUXC_SYCL_ERROR( MSG, EXPR ) \ + try { EXPR; } \ + catch( const ::sycl::exception& e ) { \ + throw sycl_exception( __FILE__, __LINE__, MSG, e ); \ + } + +#endif diff --git a/src/load_balancer/device/CMakeLists.txt b/src/load_balancer/device/CMakeLists.txt index 00b5dee2b..75dcdb973 100644 --- a/src/load_balancer/device/CMakeLists.txt +++ b/src/load_balancer/device/CMakeLists.txt @@ -20,3 +20,7 @@ endif() if( GAUXC_HAS_HIP ) add_subdirectory( hip ) endif() + +if( GAUXC_HAS_SYCL ) + add_subdirectory( sycl ) +endif() diff --git a/src/load_balancer/device/load_balancer_device_factory.cxx b/src/load_balancer/device/load_balancer_device_factory.cxx index e481f4f24..2264aaa89 100644 --- a/src/load_balancer/device/load_balancer_device_factory.cxx +++ b/src/load_balancer/device/load_balancer_device_factory.cxx @@ -20,6 +20,10 @@ #include "hip/replicated_hip_load_balancer.hpp" #endif +#ifdef GAUXC_HAS_SYCL +#include "sycl/replicated_sycl_load_balancer.hpp" +#endif + namespace GauXC { std::shared_ptr LoadBalancerDeviceFactory::get_shared_instance( diff --git a/src/load_balancer/device/sycl/CMakeLists.txt b/src/load_balancer/device/sycl/CMakeLists.txt new file mode 100644 index 000000000..6844340b5 --- /dev/null +++ b/src/load_balancer/device/sycl/CMakeLists.txt @@ -0,0 +1,15 @@ +# +# GauXC Copyright (c) 2020-2024, The Regents of the University of California, +# through Lawrence Berkeley National Laboratory (subject to receipt of +# any required approvals from the U.S. Dept. of Energy). +# +# (c) 2024-2025, Microsoft Corporation +# +# All rights reserved. +# +# See LICENSE.txt for details +# +target_sources( gauxc PRIVATE sycl_collision_detection.cxx + replicated_sycl_load_balancer.cxx ) +gauxc_sycl_device_sources( sycl_collision_detection.cxx + replicated_sycl_load_balancer.cxx ) diff --git a/src/load_balancer/device/sycl/replicated_sycl_load_balancer.cxx b/src/load_balancer/device/sycl/replicated_sycl_load_balancer.cxx new file mode 100644 index 000000000..38e457497 --- /dev/null +++ b/src/load_balancer/device/sycl/replicated_sycl_load_balancer.cxx @@ -0,0 +1,349 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "replicated_sycl_load_balancer.hpp" +#include +#include +#include "device_specific/sycl_util.hpp" +#include "device/sycl/sycl_backend.hpp" + +#include "sycl_collision_detection.hpp" +#include + +using namespace GauXC::load_balancer::sycl; + +namespace GauXC { +namespace detail { + +// TODO: Make this use a pinned allocator +template +using pinned_vector = std::vector; + +// Helper data struction to keep inputs to collision detection kernels organized +struct CollisionDetectionSyclData { + // Inputs + double* low_points_device; + double* high_points_device; + double* centers_device; + double* radii_device; + size_t* shell_sizes_device; + // Outputs + int32_t position_list_length; + int32_t* position_list_device; + int32_t* counts_device; + size_t* nbe_list_device; + // Intermediates + int32_t* collisions_device; + size_t temp_storage_bytes; + void * temp_storage_device; +}; + + +DeviceReplicatedLoadBalancer::DeviceReplicatedLoadBalancer( const DeviceReplicatedLoadBalancer& ) = default; +DeviceReplicatedLoadBalancer::DeviceReplicatedLoadBalancer( DeviceReplicatedLoadBalancer&& ) noexcept = default; + +DeviceReplicatedLoadBalancer::~DeviceReplicatedLoadBalancer() noexcept = default; + +std::unique_ptr DeviceReplicatedLoadBalancer::clone() const { + return std::make_unique(*this); +} + +std::vector inline copy_shell_list( + const size_t idx, + const std::vector& counts, + const pinned_vector &position_list +) { + int32_t start = 0; + if ( idx != 0 ) start += counts[idx-1]; + int32_t end = counts[idx]; + + std::vector shell_list(end - start); + std::copy(position_list.begin() + start, position_list.begin() + end, shell_list.begin()); + + return shell_list; +} + + +std::vector< XCTask > DeviceReplicatedLoadBalancer::create_local_tasks_() const { + + const int32_t n_deriv = 1; + const size_t atBatchSz = 256; + + int32_t world_rank = runtime_.comm_rank(); + int32_t world_size = runtime_.comm_size(); + + std::vector< XCTask > local_work; + std::vector global_workload( world_size, 0 ); + + const auto natoms = this->mol_->natoms(); + const size_t nspheres = (*this->basis_).size(); + const size_t num_atom_batch = util::div_ceil(natoms, atBatchSz); + const size_t max_nbatches = mg_->max_nbatches() * atBatchSz; + const size_t LD_bit = util::div_ceil(nspheres, 32); + + CollisionDetectionSyclData data; + + std::optional fallback_stream; + ::sycl::queue* master_stream_ptr = nullptr; + if( auto* p = dynamic_cast(&runtime_) ) { + auto* backend = dynamic_cast( p->device_backend() ); + if( !backend ) GAUXC_BAD_BACKEND_CAST(); + master_stream_ptr = &backend->master_stream->queue; + } else { + fallback_stream.emplace( util::sycl_default_device() ); + master_stream_ptr = &fallback_stream->queue; + } + ::sycl::queue& master_stream = *master_stream_ptr; + + std::vector< XCTask > temp_tasks; temp_tasks.reserve( max_nbatches ); + std::vector> low_points; low_points.reserve( max_nbatches ); + std::vector> high_points; high_points.reserve( max_nbatches ); + std::vector> centers; centers.reserve(nspheres); + std::vector radii; radii.reserve(nspheres); + std::vector shell_sizes; shell_sizes.reserve(nspheres); + // These two vectors are populated by cuda memcopies on their data pointer + // So maybe we should be resizing them instead of just adding capacity? + std::vector pos_list_idx; pos_list_idx.reserve(max_nbatches); + std::vector nbe_vec; nbe_vec.reserve(max_nbatches); + + // The postion list is the largest struction so I am using pinned memory for the improved bandwidth + pinned_vector position_list; + + data.temp_storage_bytes = compute_scratch(max_nbatches, data.counts_device); + data.temp_storage_device = data.temp_storage_bytes + ? util::sycl_malloc(data.temp_storage_bytes, master_stream) // char is 1 byte + : nullptr; + + data.low_points_device = util::sycl_malloc(max_nbatches * 3, master_stream); + data.high_points_device = util::sycl_malloc(max_nbatches * 3, master_stream); + data.collisions_device = util::sycl_malloc(LD_bit * max_nbatches, master_stream); + data.nbe_list_device = util::sycl_malloc(max_nbatches, master_stream); + data.counts_device = util::sycl_malloc(max_nbatches, master_stream); + + data.centers_device = util::sycl_malloc(nspheres * 3, master_stream); + data.radii_device = util::sycl_malloc(nspheres, master_stream); + data.shell_sizes_device = util::sycl_malloc(nspheres, master_stream); + + for(auto& shell : (*this->basis_)) { + centers.push_back(shell.O()); + radii.push_back(shell.cutoff_radius()); + shell_sizes.push_back(shell.size()); + } + + util::sycl_copy(nspheres * 3, data.centers_device, centers[0].data(), master_stream, "Centers HtoD"); + util::sycl_copy(nspheres, data.radii_device, radii.data(), master_stream, "Radii HtoD"); + util::sycl_copy(nspheres, data.shell_sizes_device, shell_sizes.data(), master_stream, "ShellSize HtoD"); + + // For batching of multiple atom screening + for (size_t atom_batch = 0; atom_batch < num_atom_batch; ++atom_batch) { + //--------------------------------------------------------------------- + // production step + int32_t iCurrent = atom_batch * atBatchSz; + for ( size_t atom_idx = 0; atom_idx < atBatchSz && atom_batch * atBatchSz + atom_idx < natoms; ++atom_idx ) { + + const auto atom = (*this->mol_)[atom_batch * atBatchSz + atom_idx]; + const std::array center = { atom.x, atom.y, atom.z }; + + auto& batcher = this->mg_->get_grid(atom.Z).batcher(); + batcher.quadrature().recenter( center ); + const size_t nbatches = batcher.nbatches(); + + for( size_t ibatch = 0; ibatch < nbatches; ++ibatch ) { + // Generate the batch (non-negligible cost) + auto [ npts, pts_b, pts_en, w_b, w_en ] = (batcher.begin() + ibatch).range(); + auto [lo, up] = IntegratorXX::detail::get_box_bounds_points(pts_b, pts_en); + + if( npts == 0 ) continue; + + // Partially copy task data + XCTask task; + task.iParent = iCurrent; + task.npts = npts; + task.dist_nearest = this->molmeta_->dist_nearest()[iCurrent]; + temp_tasks.push_back( std::move( task ) ); + low_points.push_back( std::move( lo ) ); + high_points.push_back( std::move( up ) ); + } + iCurrent++; + } + + //--------------------------------------------------------------------- + // Device collision detection step + const size_t ncubes = low_points.size(); + util::sycl_copy(ncubes * 3, data.low_points_device, low_points[0].data(), master_stream, "Low points HtoD"); + util::sycl_copy(ncubes * 3, data.high_points_device, high_points[0].data(), master_stream, "High points HtoD"); + + collision_detection( + ncubes, nspheres, LD_bit, + data.low_points_device, data.high_points_device, + data.centers_device, data.radii_device, + data.temp_storage_bytes, data.temp_storage_device, + data.collisions_device, data.counts_device, + master_stream + ); + + // Copy total number of collisions back to host to allocate result array + int32_t total_collisions; + util::sycl_copy(1, &total_collisions, data.counts_device + ncubes - 1, master_stream, "Total collisions DtoH"); + data.position_list_device = util::sycl_malloc(total_collisions, master_stream); + + compute_position_list( + ncubes, nspheres, LD_bit, + data.shell_sizes_device, + data.collisions_device, + data.counts_device, + data.position_list_device, + data.nbe_list_device, + master_stream + ); + + position_list.reserve(total_collisions); + + util::sycl_device_sync(master_stream); + // Copy results back to host + util::sycl_copy(total_collisions, position_list.data(), data.position_list_device, master_stream, "Position List DtoH"); + util::sycl_copy(ncubes, pos_list_idx.data(), data.counts_device, master_stream, "Position List Idx DtoH"); + util::sycl_copy(ncubes, nbe_vec.data(), data.nbe_list_device, master_stream, "NBE counts DtoH"); + util::sycl_free(master_stream, data.position_list_device); + + low_points.clear(); + high_points.clear(); + + //--------------------------------------------------------------------- + // Assign batches to MPI ranks + size_t idx = 0; + for ( size_t atom_idx = 0; atom_idx < atBatchSz && atom_batch * atBatchSz + atom_idx < natoms; ++atom_idx ) { + + const auto atom = (*this->mol_)[atom_batch * atBatchSz + atom_idx]; + const std::array center = { atom.x, atom.y, atom.z }; + + auto& batcher = this->mg_->get_grid(atom.Z).batcher(); + batcher.quadrature().recenter( center ); + const size_t nbatches = batcher.nbatches(); + + for( size_t ibatch = 0; ibatch < nbatches; ++ibatch ) { + auto [ npts, pts_b, pts_en, w_b, w_en ] = (batcher.begin() + ibatch).range(); + XCTask task = std::move( temp_tasks.at( idx ) ); + task.bfn_screening.nbe = nbe_vec[idx]; + + // Update npts + task.npts = npts; + + // Get rank with minimum work + auto min_rank_it = + std::min_element( global_workload.begin(), global_workload.end() ); + int64_t min_rank = std::distance( global_workload.begin(), min_rank_it ); + + global_workload[ min_rank ] += task.cost( n_deriv, natoms ); + + if( world_rank == min_rank ) { + auto shell_list = std::move( copy_shell_list(idx, pos_list_idx, position_list) ); + // Course grain screening + if( shell_list.size() ) { + task.bfn_screening.shell_list = shell_list; + + // Get local copy of points weights + std::vector> points(pts_b, pts_en); + std::vector weights(w_b, w_en); + + task.points = std::move(points); + task.weights = std::move(weights); + local_work.push_back( std::move(task) ); + } + } + idx++; + } + } + temp_tasks.clear(); + + } + + // Lexicographic ordering of tasks + auto task_order = []( const auto& a, const auto& b ) { + + // Sort by iParent first + if( a.iParent < b.iParent ) return true; + else if( a.iParent > b.iParent ) return false; + + // Equal iParent: lex sort on shell list + else return a.bfn_screening.shell_list < b.bfn_screening.shell_list; + + }; + + std::sort( local_work.begin(), local_work.end(), + task_order ); + + + // Get unique tasks + auto task_equiv = []( const auto& a, const auto& b ) { + return a.equiv_with(b); + }; + + auto local_work_unique = local_work; + auto last_unique = + std::unique( local_work_unique.begin(), + local_work_unique.end(), + task_equiv ); + local_work_unique.erase( last_unique, local_work_unique.end() ); + + + // Merge tasks + for( auto&& t : local_work_unique ) { + t.points.clear(); + t.weights.clear(); + } + + auto cur_lw_begin = local_work.begin(); + auto cur_uniq_it = local_work_unique.begin(); + + for( auto lw_it = local_work.begin(); lw_it != local_work.end(); ++lw_it ) + if( not task_equiv( *lw_it, *cur_uniq_it ) ) { + + if( cur_uniq_it == local_work_unique.end() ) + GAUXC_GENERIC_EXCEPTION("Messed up in unique"); + + cur_uniq_it->merge_with( cur_lw_begin, lw_it ); + + cur_lw_begin = lw_it; + cur_uniq_it++; + + } + + // Merge the last set of batches + for( ; cur_lw_begin != local_work.end(); ++cur_lw_begin ) + cur_uniq_it->merge_with( *cur_lw_begin ); + cur_uniq_it++; + + + local_work = std::move(local_work_unique); + + // Free all device memory + util::sycl_free(master_stream, data.low_points_device); + util::sycl_free(master_stream, data.high_points_device); + util::sycl_free(master_stream, data.centers_device); + util::sycl_free(master_stream, data.radii_device); + util::sycl_free(master_stream, data.shell_sizes_device); + util::sycl_free(master_stream, data.collisions_device); + util::sycl_free(master_stream, data.nbe_list_device); + util::sycl_free(master_stream, data.counts_device); + util::sycl_free(master_stream, data.temp_storage_device); + + return local_work; +} + + + + + + + +} +} diff --git a/src/load_balancer/device/sycl/replicated_sycl_load_balancer.hpp b/src/load_balancer/device/sycl/replicated_sycl_load_balancer.hpp new file mode 100644 index 000000000..585edde3f --- /dev/null +++ b/src/load_balancer/device/sycl/replicated_sycl_load_balancer.hpp @@ -0,0 +1,43 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include "load_balancer_impl.hpp" + +namespace GauXC { +namespace detail { + +class DeviceReplicatedLoadBalancer : public LoadBalancerImpl { + +protected: + + using basis_type = BasisSet; + std::vector< XCTask > create_local_tasks_() const override; + +public: + + DeviceReplicatedLoadBalancer() = delete; + template + DeviceReplicatedLoadBalancer( Args&&... args ): + LoadBalancerImpl( std::forward(args)... ) { } + + DeviceReplicatedLoadBalancer( const DeviceReplicatedLoadBalancer& ); + DeviceReplicatedLoadBalancer( DeviceReplicatedLoadBalancer&& ) noexcept; + + virtual ~DeviceReplicatedLoadBalancer() noexcept; + + std::unique_ptr clone() const override; + +}; + +} +} diff --git a/src/load_balancer/device/sycl/sycl_collision_detection.cxx b/src/load_balancer/device/sycl/sycl_collision_detection.cxx new file mode 100644 index 000000000..e3e96359d --- /dev/null +++ b/src/load_balancer/device/sycl/sycl_collision_detection.cxx @@ -0,0 +1,263 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include +#include +#include + +#include "sycl_collision_detection.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "xc_integrator/local_work_driver/device/sycl/kernels/sycl_launch.hpp" + +namespace GauXC { +namespace load_balancer { +namespace sycl { + +// NOTE: `sycl` here names GauXC::load_balancer::sycl (this namespace), which +// also shadows GauXC::sycl (device constants, brought in via `using namespace` +// below) and the global ::sycl namespace. Every reference to the SYCL runtime +// itself must therefore be spelled with a leading `::sycl::` in this file. +using namespace GauXC::sycl; + +inline +int cube_sphere_intersect( + const double3 lo, + const double3 up, + const double3 center, + const double rad +) { + + double dist = rad * rad; + + if( center.x < lo.x ) { + const double r_lo = center.x - lo.x; + const double dist_lo = r_lo * r_lo; + dist -= dist_lo; + } else if( center.x > up.x ) { + const double r_up = center.x - up.x; + const double dist_up = r_up * r_up; + dist -= dist_up; + } + + if( dist < 0. ) return false; + + if( center.y < lo.y ) { + const double r_lo = center.y - lo.y; + const double dist_lo = r_lo * r_lo; + dist -= dist_lo; + } else if( center.y > up.y ) { + const double r_up = center.y - up.y; + const double dist_up = r_up * r_up; + dist -= dist_up; + } + + if( dist < 0. ) return false; + + + if( center.z < lo.z ) { + const double r_lo = center.z - lo.z; + const double dist_lo = r_lo * r_lo; + dist -= dist_lo; + } else if( center.z > up.z ) { + const double r_up = center.z - up.z; + const double dist_up = r_up * r_up; + dist -= dist_up; + } + + return dist > 0.; + +} + + +void collision_detection_gpu( + size_t ncubes, + size_t nspheres, + size_t LD_bit, + const double* low_points, + const double* high_points, + const double* centers, + const double* radii, + int32_t* collisions, + int32_t* counts +) { + auto it = GauXC::sycl::this_item(); + const size_t nspheres_block = (nspheres + 31) / 32; + for (int i = it.get_global_id(2); i < ncubes; i += it.get_local_range(2) * it.get_group_range(2)) { + counts[i] = 0; + double3 low_point; + double3 high_point; + low_point.x = low_points[3*i+0]; + low_point.y = low_points[3*i+1]; + low_point.z = low_points[3*i+2]; + + high_point.x = high_points[3*i+0]; + high_point.y = high_points[3*i+1]; + high_point.z = high_points[3*i+2]; + + + for (int j_block = 0; j_block < nspheres_block; j_block++) { + int temp_collisions = 0; + for (int j_inner = 0; j_inner < 32; j_inner++) { + int j = j_block * 32 + j_inner; + if (j < nspheres) { + double3 center; + double radius; + center.x = centers[3*j+0]; + center.y = centers[3*j+1]; + center.z = centers[3*j+2]; + + radius = radii[j]; + temp_collisions |= (cube_sphere_intersect(low_point, high_point, center, radius) ? 1 << (j_inner) : 0); + } + } + collisions[i * LD_bit + j_block] = temp_collisions; + counts[i] += ::sycl::popcount(temp_collisions); + } + } +} + + +static constexpr int32_t buffer_size = 8; +static constexpr int32_t element_size = 32; +static constexpr int32_t buffer_size_bits = buffer_size * element_size; + +// This kernel converts the bitvector produced by the collision detection kernel above into a position list. +// For simplicity, the collision detection kernel stores its output as a bitvector. However, the `shell_list` +// of the task is a list of the qualifying indexes, so we must convert the bitvector to a position list. +// +// We take this chance to compute the nbe value from the shell sizes since the data is already being read in +void bitvector_to_position_list( + size_t ncubes, + size_t nspheres, + size_t LD_bit, + const int32_t* collisions, + const int32_t* counts, + const size_t* shell_size, + int32_t* position_list, + size_t* nbe_list +) { + auto it = GauXC::sycl::this_item(); + auto& collisions_buffer = + GauXC::sycl::local_mem< int32_t[warp_size][warp_size][buffer_size] >(); + + // We are converting a large number of small bitvectors into position lists. For this reason, I am assigning a single thread to each bitvector + // This avoids having to do popcounts and warp wide reductions, but hurts the memory access pattern + + // All threads in a warp must be active to do shared memory loads, so we seperate out the threadId.x + for (int i_base = it.get_local_id(1) * it.get_local_range(2) + it.get_group(2) * it.get_local_range(2) * it.get_local_range(1); + i_base < ncubes; + i_base += it.get_local_range(2) * it.get_local_range(1) * it.get_group_range(2)) { + const int i = i_base + it.get_local_id(2); + int32_t* out = position_list; + if (i != 0 && i < ncubes) { + out += counts[i-1]; + } + + int current = 0; + size_t nbe = 0; + size_t nsphere_blocks = (nspheres + buffer_size_bits - 1) / buffer_size_bits; + for (int j_block = 0; j_block < nsphere_blocks; j_block++) { + // Each thread has a buffer of length BUFFER_SIZE. All the threads in the warp work to + // load this data in a coalesced way (at least as much as possible) + for (int buffer_loop = 0; buffer_loop < warp_size; buffer_loop += warp_size/buffer_size) { + const int t_id_x = it.get_local_id(2) % buffer_size; + const int buffer_thread = it.get_local_id(2) / buffer_size; + const int buffer_idx = buffer_thread + buffer_loop; + if (j_block * buffer_size_bits + t_id_x * element_size < nspheres && i_base + buffer_idx < ncubes) { + collisions_buffer[it.get_local_id(1)][buffer_idx][t_id_x] = collisions[(i_base + buffer_idx) * LD_bit + j_block * buffer_size + t_id_x]; + } + } + + ::sycl::group_barrier( it.get_sub_group() ); + if (i < ncubes) { // Once the data has been loaded, we exclude the threads not corresponding to a bitvector + // We have loaded in BUFFER_SIZE_BITS elements to be processed by each warp + for (int j_inner = 0; j_inner < buffer_size_bits && j_block * buffer_size_bits + j_inner < nspheres; j_inner++) { + const int j = buffer_size_bits * j_block + j_inner; + const int j_int = j_inner / element_size; + const int j_bit = j_inner % element_size; + if( collisions_buffer[it.get_local_id(1)][it.get_local_id(2)][j_int] & (1 << (j_bit)) ) { + out[current++] = j; + nbe += shell_size[j]; + } + } + } + ::sycl::group_barrier( it.get_sub_group() ); + } + if (i < ncubes) { + nbe_list[i] = nbe; + } + } +} + +size_t compute_scratch( size_t ncubes, int32_t* counts_device ) { + // Computes amount of memory that will be required to do the inclusive sum + return 0; +} + +void collision_detection( size_t ncubes, + size_t nspheres, + size_t LD_bit, + const double* low_points_device, + const double* high_points_device, + const double* centers_device, + const double* radii_device, + size_t temp_storage_bytes, + void * temp_storage_device, + int32_t* collisions_device, + int32_t* counts_device, + ::sycl::queue& stream) { + + GauXC::sycl::launch_dim threads( max_threads_per_thread_block ); + GauXC::sycl::launch_dim blocks( util::div_ceil( ncubes, threads.x ) ); + + GauXC::sycl::launch_kernel( stream, blocks, threads, "collision_detection_gpu", + [=](){ + collision_detection_gpu( + ncubes, nspheres, LD_bit, + low_points_device, high_points_device, centers_device, radii_device, + collisions_device, counts_device + ); + }); + + // Run inclusive prefix sum + oneapi::dpl::inclusive_scan( + oneapi::dpl::execution::make_device_policy(stream), + counts_device, counts_device + ncubes, counts_device ); + +} + +void compute_position_list(size_t ncubes, + size_t nspheres, + size_t LD_bit, + const size_t* shell_sizes_device, + const int32_t* collisions_device, + const int32_t* counts_device, + int32_t* position_list_device, + size_t* nbe_list_device, + ::sycl::queue& stream) { + GauXC::sycl::launch_dim threads( warp_size, warp_size ); + GauXC::sycl::launch_dim blocks( util::div_ceil( ncubes, threads.x * threads.y ) ); + + // convert from bitvector to position list + GauXC::sycl::launch_kernel( stream, blocks, threads, "bitvector_to_position_list", + [=](){ + bitvector_to_position_list( + ncubes, nspheres, LD_bit, + collisions_device, counts_device, shell_sizes_device, + position_list_device, nbe_list_device + ); + }); +} + +} +} +} diff --git a/src/load_balancer/device/sycl/sycl_collision_detection.hpp b/src/load_balancer/device/sycl/sycl_collision_detection.hpp new file mode 100644 index 000000000..128d8b67e --- /dev/null +++ b/src/load_balancer/device/sycl/sycl_collision_detection.hpp @@ -0,0 +1,46 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include + +namespace GauXC { +namespace load_balancer { +namespace sycl { + +size_t compute_scratch( size_t ncubes, int32_t* counts_device ); + +void collision_detection( size_t ncubes, + size_t nspheres, + size_t LD_bit, + const double* low_points_device, + const double* high_points_device, + const double* centers_device, + const double* radii_device, + size_t temp_storage_bytes, + void * temp_storage_device, + int32_t* collisions_device, + int32_t* counts_device, + ::sycl::queue& stream); + +void compute_position_list(size_t ncubes, + size_t nspheres, + size_t LD_bit, + const size_t* shell_sizes_device, + const int32_t* collisions_device, + const int32_t* counts_device, + int32_t* position_list_device, + size_t* nbe_list_device, + ::sycl::queue& stream); + +} +} +} diff --git a/src/reduction_driver/device/CMakeLists.txt b/src/reduction_driver/device/CMakeLists.txt index 95d6575dd..373ba7cb1 100644 --- a/src/reduction_driver/device/CMakeLists.txt +++ b/src/reduction_driver/device/CMakeLists.txt @@ -21,3 +21,8 @@ if( GAUXC_ENABLE_NCCL ) set(GAUXC_HAS_NCCL TRUE CACHE BOOL "GauXC has NCCL" FORCE) endif() +if( GAUXC_ENABLE_ONECCL ) + target_sources( gauxc PRIVATE oneccl_reduction_driver.cxx ) + gauxc_sycl_device_sources( oneccl_reduction_driver.cxx ) + message( STATUS "oneCCL Has Been Enabled" ) +endif() diff --git a/src/reduction_driver/device/oneccl_reduction_driver.cxx b/src/reduction_driver/device/oneccl_reduction_driver.cxx new file mode 100644 index 000000000..48174eb0d --- /dev/null +++ b/src/reduction_driver/device/oneccl_reduction_driver.cxx @@ -0,0 +1,138 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "oneccl_reduction_driver.hpp" +#include +#include +#include +#include +#include + +#include "device/sycl/sycl_backend.hpp" +#include "device/device_queue.hpp" + +namespace GauXC { + +ccl::datatype get_ccl_datatype( std::type_index idx ) { + + static std::map map { + {std::type_index(typeid(double)), ccl::datatype::float64}, + {std::type_index(typeid(float)), ccl::datatype::float32} + }; + + return map.at(idx); + +} + +ccl::reduction get_ccl_op( ReductionOp op ) { + + static std::map< ReductionOp, ccl::reduction > map { + { ReductionOp::Sum, ccl::reduction::sum } + }; + + return map.at(op); + +} + +SYCLBackend* get_sycl_backend_from_runtime( const RuntimeEnvironment& rt ) { + auto* backend = + dynamic_cast(detail::as_device_runtime(rt).device_backend()); + if( !backend ) GAUXC_GENERIC_EXCEPTION("OneCCL Reduction Driver Requires A SYCL Backend"); + return backend; +} + +util::sycl_queue* get_sycl_queue_from_optional_args( std::any& args, + SYCLBackend* backend ) { + + if( args.has_value() ) { + if( auto ptr = std::any_cast( &args ) ) + if( auto passed_stream = ptr->queue_as_ptr() ) { + return passed_stream; + } + } + + device_queue backend_queue = backend->queue(); + auto* q = backend_queue.queue_as_ptr(); + if( !q ) GAUXC_GENERIC_EXCEPTION("SYCL Backend Did Not Return A Valid Queue"); + return q; +} + + +OneCCLReductionDriver::OneCCLReductionDriver(const RuntimeEnvironment& rt) : + DeviceReductionDriver(rt) { + + auto* backend = get_sycl_backend_from_runtime(rt); + auto context = backend->master_stream->queue.get_context(); + oneccl_comm_ = std::make_shared( rt.comm(), + backend->device, context ); + ccl_streams_ = + std::make_shared>(); + +} + + +OneCCLReductionDriver::~OneCCLReductionDriver() noexcept = default; +OneCCLReductionDriver::OneCCLReductionDriver(const OneCCLReductionDriver&) = default; + + +ccl::stream& OneCCLReductionDriver::ccl_stream_for( ::sycl::queue& q ) { + auto it = ccl_streams_->find(q); + if( it == ccl_streams_->end() ) + it = ccl_streams_->emplace( q, ccl::create_stream(q) ).first; + return it->second; +} + + +void OneCCLReductionDriver::allreduce_typeerased( const void* src, void* dest, + size_t size, ReductionOp op, std::type_index idx, std::any optional_args ) { + + auto* backend = get_sycl_backend_from_runtime( runtime_ ); + auto* queue = get_sycl_queue_from_optional_args( optional_args, backend ); + + auto synchronize = [&]() { queue->synchronize(); }; + + synchronize(); + try { + auto& stream = ccl_stream_for( queue->queue ); + ccl::allreduce( src, dest, size, get_ccl_datatype(idx), + get_ccl_op(op), *oneccl_comm_, stream ).wait(); + } catch( const ccl::exception& e ) { + GAUXC_GENERIC_EXCEPTION( std::string("OneCCL FAILED: ") + e.what() ); + } + synchronize(); + +} +void OneCCLReductionDriver::allreduce_inplace_typeerased( void* data, size_t size, + ReductionOp op, std::type_index idx, std::any optional_args) { + + auto* backend = get_sycl_backend_from_runtime( runtime_ ); + auto* queue = get_sycl_queue_from_optional_args( optional_args, backend ); + + auto synchronize = [&]() { queue->synchronize(); }; + + synchronize(); + try { + auto& stream = ccl_stream_for( queue->queue ); + ccl::allreduce( data, data, size, get_ccl_datatype(idx), + get_ccl_op(op), *oneccl_comm_, stream ).wait(); + } catch( const ccl::exception& e ) { + GAUXC_GENERIC_EXCEPTION( std::string("OneCCL FAILED: ") + e.what() ); + } + synchronize(); + +} + +std::unique_ptr OneCCLReductionDriver::clone() { + return std::make_unique(*this); +} + + +} diff --git a/src/reduction_driver/device/oneccl_reduction_driver.hpp b/src/reduction_driver/device/oneccl_reduction_driver.hpp new file mode 100644 index 000000000..e32a98def --- /dev/null +++ b/src/reduction_driver/device/oneccl_reduction_driver.hpp @@ -0,0 +1,91 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "device_reduction_driver.hpp" +#include +#include +#include + +namespace GauXC { + +namespace util { +struct oneccl_comm { + + std::shared_ptr kvs; + ccl::communicator comm; + + inline oneccl_comm( MPI_Comm mpi_comm, ::sycl::device& device, + ::sycl::context& context ) : + comm( make_comm(mpi_comm, device, context, kvs) ) { } + + // oneCCL communicators are move-only (they wrap a PIMPL, cf. ccl::communicator + // in communicator.hpp) + oneccl_comm( const oneccl_comm& ) = delete; + oneccl_comm( oneccl_comm&& ) noexcept = default; + + inline operator ccl::communicator&() { return comm; } + +private: + + // Builds the communicator in a helper so the member initializer list above + // can construct `comm` directly (ccl::communicator has no default ctor). + static inline ccl::communicator make_comm( MPI_Comm mpi_comm, + const ::sycl::device& device, const ::sycl::context& context, + std::shared_ptr& kvs_out ) { + + int32_t world_rank, world_size; + MPI_Comm_rank( mpi_comm, &world_rank ); + MPI_Comm_size( mpi_comm, &world_size ); + + ccl::kvs::address_type addr; + if( world_rank == 0 ) { + kvs_out = ccl::create_main_kvs(); + addr = kvs_out->get_address(); + } + MPI_Bcast( addr.data(), addr.size(), MPI_BYTE, 0, mpi_comm ); + if( world_rank != 0 ) kvs_out = ccl::create_kvs( addr ); + + // create_communicator wraps the device in ccl::pair_class + // internally, so it needs oneCCL's own device/context types. DeviceType is + // deduced from a non-const lvalue reference, hence the named local. + ccl::device ccl_dev = ccl::create_device( ::sycl::device(device) ); + return ccl::create_communicator( world_size, world_rank, ccl_dev, + ccl::create_context( ::sycl::context(context) ), kvs_out ); + } + +}; +} + + +struct OneCCLReductionDriver : public DeviceReductionDriver { + + std::shared_ptr oneccl_comm_; + + // ccl::stream wrappers around the SYCL queues this driver has been handed, + // keyed by queue. Cached so that a collective does not rebuild one per call; + // shared so that copies of the driver (clone()) reuse the same wrappers. + std::shared_ptr> ccl_streams_; + + OneCCLReductionDriver(const RuntimeEnvironment& rt); + virtual ~OneCCLReductionDriver() noexcept; + OneCCLReductionDriver(const OneCCLReductionDriver& ); + + ccl::stream& ccl_stream_for( ::sycl::queue& q ); + + void allreduce_typeerased( const void*, void*, size_t, ReductionOp, std::type_index, std::any) override; + void allreduce_inplace_typeerased( void*, size_t, ReductionOp, std::type_index, std::any ) override; + + std::unique_ptr clone() override; + +}; + +} diff --git a/src/reduction_driver/reduction_driver_factory.cxx b/src/reduction_driver/reduction_driver_factory.cxx index 39fe2f10b..f929682b2 100644 --- a/src/reduction_driver/reduction_driver_factory.cxx +++ b/src/reduction_driver/reduction_driver_factory.cxx @@ -16,6 +16,10 @@ #include "device/nccl_reduction_driver.hpp" #endif +#ifdef GAUXC_HAS_ONECCL +#include "device/oneccl_reduction_driver.hpp" +#endif + #include #include @@ -41,6 +45,11 @@ std::shared_ptr ReductionDriverFactory::get_shared_instance( ptr = std::make_unique(rt); #endif + #ifdef GAUXC_HAS_ONECCL + if( kernel_name == "ONECCL" ) + ptr = std::make_unique(rt); + #endif + if( !ptr ) GAUXC_GENERIC_EXCEPTION("Unknown Reduction Driver " + kernel_name); return std::make_shared(std::move(ptr)); diff --git a/src/runtime_environment/device/CMakeLists.txt b/src/runtime_environment/device/CMakeLists.txt index 151a58894..a918c3fe8 100644 --- a/src/runtime_environment/device/CMakeLists.txt +++ b/src/runtime_environment/device/CMakeLists.txt @@ -18,3 +18,7 @@ if(GAUXC_ENABLE_HIP) add_subdirectory( hip ) endif() +if(GAUXC_ENABLE_SYCL) + add_subdirectory( sycl ) +endif() + diff --git a/src/runtime_environment/device/sycl/CMakeLists.txt b/src/runtime_environment/device/sycl/CMakeLists.txt new file mode 100644 index 000000000..8e4938aa7 --- /dev/null +++ b/src/runtime_environment/device/sycl/CMakeLists.txt @@ -0,0 +1,13 @@ +# +# GauXC Copyright (c) 2020-2024, The Regents of the University of California, +# through Lawrence Berkeley National Laboratory (subject to receipt of +# any required approvals from the U.S. Dept. of Energy). +# +# (c) 2024-2025, Microsoft Corporation +# +# All rights reserved. +# +# See LICENSE.txt for details +# +target_sources( gauxc PRIVATE sycl_backend.cxx ) +gauxc_sycl_device_sources( sycl_backend.cxx ) diff --git a/src/runtime_environment/device/sycl/sycl_backend.cxx b/src/runtime_environment/device/sycl/sycl_backend.cxx new file mode 100644 index 000000000..ebcc9b583 --- /dev/null +++ b/src/runtime_environment/device/sycl/sycl_backend.cxx @@ -0,0 +1,178 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "sycl_backend.hpp" + +namespace GauXC { + +SYCLBackend::SYCLBackend() { + + // Create SYCL queue. oneMKL is queue-driven, so the "BLAS handle" is a + // binding to the same queue rather than a separate object to be tied to it + device = util::sycl_default_device(); + + master_stream = std::make_shared< util::sycl_queue >( device ); + master_handle = std::make_shared< util::onemkl_handle >( master_stream ); + +#ifdef GAUXC_HAS_MAGMA + #error "MAGMA is not supported with the SYCL backend" +#endif + +} + +SYCLBackend::~SYCLBackend() noexcept = default; + +SYCLBackend::device_buffer_t SYCLBackend::allocate_device_buffer(int64_t sz) { + void* ptr = nullptr; + GAUXC_SYCL_ERROR( "SYCL Malloc Failed", + ptr = ::sycl::malloc_device( sz, master_stream->queue ) ); + if( !ptr ) GAUXC_GENERIC_EXCEPTION("SYCL Malloc Returned NULL"); + return device_buffer_t{ptr,sz}; +} + +size_t SYCLBackend::get_available_mem() { + // free_memory requires ZES_ENABLE_SYSMAN=1 on Level Zero; fall back to the + // total global memory when the query is unavailable + if( device.has( ::sycl::aspect::ext_intel_free_memory ) ) + return device.get_info< ::sycl::ext::intel::info::device::free_memory >(); + else + return device.get_info< ::sycl::info::device::global_mem_size >(); +} + +void SYCLBackend::free_device_buffer( void* ptr ) { + // cudaFree implicitly synchronizes the device, so the shared code is free + // to release a buffer without first waiting on the work that reads it. + // sycl::free carries no such guarantee -- freeing a USM allocation while a + // kernel still references it is undefined behaviour -- so the outstanding + // work is drained first to preserve the same contract. + master_stream->synchronize(); + GAUXC_SYCL_ERROR( "Free Failed", ::sycl::free( ptr, master_stream->queue ) ); +} + +void SYCLBackend::master_queue_synchronize() { + master_stream->synchronize(); +} + +device_queue SYCLBackend::queue() { + return device_queue(master_stream); +} + +void SYCLBackend::create_blas_queue_pool(int32_t ns) { + blas_streams.resize(ns); + blas_handles.resize(ns); + for( auto i = 0; i < ns; ++i ) { + blas_streams[i] = std::make_shared( device ); + blas_handles[i] = std::make_shared( blas_streams[i] ); + } +} + +void SYCLBackend::sync_master_with_blas_pool() { + const auto n_streams = blas_streams.size(); + std::vector blas_events( n_streams ); + for( size_t iS = 0; iS < n_streams; ++iS ) + blas_events[iS].record( *blas_streams[iS] ); + + for( auto& event : blas_events ) master_stream->wait(event); +} + +void SYCLBackend::sync_blas_pool_with_master() { + util::sycl_event master_event; + master_event.record( *master_stream ); + for( auto& stream : blas_streams ) stream->wait( master_event ); +} + +size_t SYCLBackend::blas_pool_size(){ return blas_streams.size(); } + +device_queue SYCLBackend::blas_pool_queue(int32_t i) { + return device_queue( blas_streams.at(i) ); +} + +device_blas_handle SYCLBackend::blas_pool_handle(int32_t i) { + return device_blas_handle( blas_handles.at(i) ); +} +device_blas_handle SYCLBackend::master_blas_handle() { + return device_blas_handle( master_handle ); +} + +void SYCLBackend::copy_async_( size_t sz, const void* src, void* dest, + std::string msg ) { + // Despite the name, this must complete before returning. + // + // cudaMemcpyAsync is only truly asynchronous when the host side is pinned + // memory; for a device-to-pageable-host transfer CUDA makes it implicitly + // synchronizing (the driver stages through an internal buffer and blocks). + // The shared code relies on that: retrieve_exc_vxc_integrands issues these + // copies into ordinary host allocations and its caller reads them straight + // away, with no synchronization in between. + // + // SYCL's queue::memcpy offers no such guarantee -- it returns an event and + // nothing more -- so without waiting here the host races the in-flight copy + // and reads garbage (observed as NaN in VXC, then heap corruption). + GAUXC_SYCL_ERROR( "SYCL Memcpy Async Failed ["+msg+"]", + master_stream->queue.memcpy( dest, src, sz ).wait_and_throw() ); +} + +void SYCLBackend::set_zero_(size_t sz, void* data, std::string msg ) { + GAUXC_SYCL_ERROR( "SYCL Memset Failed ["+msg+"]", + master_stream->queue.memset( data, 0, sz ).wait_and_throw() ); +} + +void SYCLBackend::set_zero_async_master_queue_(size_t sz, void* data, std::string msg ) { + GAUXC_SYCL_ERROR( "SYCL Memset Failed ["+msg+"]", + master_stream->queue.memset( data, 0, sz ) ); +} + +void SYCLBackend::copy_async_2d_( size_t M, size_t N, const void* A, size_t LDA, + void* B, size_t LDB, std::string msg ) { + + // A strided copy of N rows, M bytes each. + // + // queue::ext_oneapi_memcpy2d is deliberately avoided: its helper kernel is + // not emitted into an ahead-of-time device image, so an AoT build faults + // inside the runtime's kernel-info lookup. The same entry point is also + // implicated in a known Level Zero defect on this platform + // (CMPLRLLVM-76543 / GSD-11506). One enqueue per row keeps the copy on the + // ordinary USM memcpy path, which both AoT and JIT handle. + // + // The copy is completed before returning, for the same reason as + // copy_async_: callers inherit CUDA's implicitly-synchronizing semantics + // for a device-to-pageable-host transfer. + if( !M or !N ) return; + + auto* dst = static_cast(B); + const auto* src = static_cast(A); + + GAUXC_SYCL_ERROR( "SYCL 2D Memcpy Async Failed ["+msg+"]", + [&]() { + // Contiguous on both sides: a single memcpy covers the whole region. + if( LDA == M and LDB == M ) { + master_stream->queue.memcpy( dst, src, M*N ); + } else { + for( size_t row = 0; row < N; ++row ) + master_stream->queue.memcpy( dst + row*LDB, src + row*LDA, M ); + } + master_stream->queue.wait_and_throw(); + }() + ); +} + + +void SYCLBackend::check_error_(std::string msg) { + // SYCL reports kernel failures asynchronously through the queue's exception + // handler. Draining it here is the closest analogue to cudaGetLastError + GAUXC_SYCL_ERROR( "SYCL Failed ["+msg+"]", + master_stream->queue.throw_asynchronous() ); +} + +std::unique_ptr make_device_backend() { + return std::make_unique(); +} +} diff --git a/src/runtime_environment/device/sycl/sycl_backend.hpp b/src/runtime_environment/device/sycl/sycl_backend.hpp new file mode 100644 index 000000000..1f5321a92 --- /dev/null +++ b/src/runtime_environment/device/sycl/sycl_backend.hpp @@ -0,0 +1,59 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "device/device_backend.hpp" +#include +#include "device_specific/sycl_util.hpp" +#include "device_specific/onemkl_util.hpp" + +namespace GauXC { + +struct SYCLBackend : public DeviceBackend { + + device_buffer_t allocate_device_buffer(int64_t sz) override final; + size_t get_available_mem() override final; + void free_device_buffer( void* ptr ) override final; + void master_queue_synchronize() override final; + void create_blas_queue_pool(int32_t) override final; + void sync_master_with_blas_pool() override final; + void sync_blas_pool_with_master() override final; + size_t blas_pool_size() override final; + + device_queue queue() override final; + device_queue blas_pool_queue(int32_t) override final; + device_blas_handle blas_pool_handle(int32_t) override final; + device_blas_handle master_blas_handle() override final; + + void copy_async_( size_t sz, const void* src, void* dest, + std::string msg ) override final; + void set_zero_( size_t sz, void* data, std::string msg) override final; + void set_zero_async_master_queue_( size_t sz, void* data, std::string msg) override final; + + void copy_async_2d_( size_t M, size_t N, const void* A, size_t LDA, + void* B, size_t LDB, std::string msg ) override final; + + void check_error_(std::string msg) override final; + + SYCLBackend(); + ~SYCLBackend() noexcept; + + // Execution management + ::sycl::device device; + + std::shared_ptr master_stream = nullptr; + std::shared_ptr master_handle = nullptr; + + std::vector> blas_streams; + std::vector> blas_handles; +}; + +} diff --git a/src/runtime_environment/device_specific/onemkl_util.hpp b/src/runtime_environment/device_specific/onemkl_util.hpp new file mode 100644 index 000000000..78eb08e4d --- /dev/null +++ b/src/runtime_environment/device_specific/onemkl_util.hpp @@ -0,0 +1,52 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include "exceptions/onemkl_exception.hpp" +#include "sycl_util.hpp" + +#ifdef GAUXC_HAS_SYCL +#include + +namespace GauXC { +namespace util { + +/// oneMKL takes the SYCL queue directly rather than an opaque handle, so the +/// GauXC BLAS "handle" is a thin non-owning binding to the queue the BLAS +/// calls are to be issued on. +struct onemkl_handle { + + std::shared_ptr queue_ptr; + + inline onemkl_handle() = delete; + inline onemkl_handle( std::shared_ptr q ) : + queue_ptr(std::move(q)) { } + + inline ~onemkl_handle() noexcept = default; + + onemkl_handle( const onemkl_handle& ) = delete; + onemkl_handle( onemkl_handle&& ) noexcept = default; + + inline ::sycl::queue& queue() { return queue_ptr->queue; } + inline operator ::sycl::queue&() { return queue_ptr->queue; } + +}; + + +inline static ::sycl::queue& get_queue( onemkl_handle& handle ) { + return handle.queue(); +} + +} +} + +#endif diff --git a/src/runtime_environment/device_specific/sycl_device_constants.hpp b/src/runtime_environment/device_specific/sycl_device_constants.hpp new file mode 100644 index 000000000..a02a63a0b --- /dev/null +++ b/src/runtime_environment/device_specific/sycl_device_constants.hpp @@ -0,0 +1,27 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include + +namespace GauXC { +namespace sycl { + +/// Sub-group size requested by all GauXC SYCL kernels. Matches +/// detail::shell_nprim_max so that a shell's primitives map 1:1 onto the +/// lanes of a sub-group. +static constexpr uint32_t warp_size = 32; +static constexpr uint32_t max_threads_per_thread_block = 1024; +static constexpr uint32_t max_warps_per_thread_block = + max_threads_per_thread_block / warp_size; + +} +} diff --git a/src/runtime_environment/device_specific/sycl_util.hpp b/src/runtime_environment/device_specific/sycl_util.hpp new file mode 100644 index 000000000..e15870c83 --- /dev/null +++ b/src/runtime_environment/device_specific/sycl_util.hpp @@ -0,0 +1,238 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include "exceptions/sycl_exception.hpp" + +#ifdef GAUXC_HAS_SYCL +#include +#include + +namespace GauXC { +namespace util { + +struct sycl_queue; +struct sycl_event; + +/// Async exception handler which promotes SYCL errors to GauXC exceptions +inline void sycl_async_error_handler( ::sycl::exception_list exceptions ) { + for( auto& e : exceptions ) { + try { std::rethrow_exception(e); } + catch( const ::sycl::exception& err ) { + throw sycl_exception( __FILE__, __LINE__, "SYCL Async Error", err ); + } + } +} + +/// Device selection honoring ZE_AFFINITY_MASK / ONEAPI_DEVICE_SELECTOR. Falls +/// back to the default selector when no GPU is visible so that CPU / FPGA +/// backends remain usable. +inline ::sycl::device sycl_default_device() { + try { + return ::sycl::device( ::sycl::gpu_selector_v ); + } catch( const ::sycl::exception& ) { + return ::sycl::device( ::sycl::default_selector_v ); + } +} + +/// RAII wrapper for an in-order SYCL queue. +struct sycl_queue { + + ::sycl::queue queue; + + inline sycl_queue() : sycl_queue( sycl_default_device() ) { } + + /// Bind to the platform's default context for `dev`. + inline sycl_queue( const ::sycl::device& dev ) : + queue( dev, sycl_async_error_handler, + ::sycl::property_list{ ::sycl::property::queue::in_order{} } ) { } + + inline sycl_queue( const ::sycl::context& ctx, const ::sycl::device& dev ) : + queue( ctx, dev, sycl_async_error_handler, + ::sycl::property_list{ ::sycl::property::queue::in_order{} } ) { } + + inline ~sycl_queue() noexcept = default; + + sycl_queue( const sycl_queue& ) = delete; + sycl_queue( sycl_queue&& ) noexcept = default; + + inline operator ::sycl::queue&() { return queue; } + inline operator const ::sycl::queue&() const { return queue; } + + inline ::sycl::queue* operator->() { return &queue; } + + /// Enqueue a barrier which blocks this queue until "event" has completed + inline void wait( ::sycl::event event ) { + GAUXC_SYCL_ERROR( "QUEUE WAIT FAILED", + queue.ext_oneapi_submit_barrier({event}) ); + } + + inline void synchronize() { + GAUXC_SYCL_ERROR( "SYCL Queue Synchronize Failed", queue.wait_and_throw() ); + } + +}; + + +/// RAII wrapper for a SYCL event. +struct sycl_event { + + ::sycl::event event; + + inline sycl_event() = default; + inline ~sycl_event() noexcept = default; + + sycl_event( const sycl_event& ) = default; + sycl_event( sycl_event&& ) noexcept = default; + + inline operator ::sycl::event() const { return event; } + + inline void record( ::sycl::queue& queue ) { + GAUXC_SYCL_ERROR( "Event Record Failed", + event = queue.ext_oneapi_submit_barrier() ); + } + + inline void record( sycl_queue& queue ) { record( queue.queue ); } + +}; + + + + + +template +inline T* sycl_malloc( size_t n, ::sycl::queue& queue ) { + + T* ptr = nullptr; + GAUXC_SYCL_ERROR( "SYCL Malloc Failed", + ptr = ::sycl::malloc_device( n, queue ) ); + if( !ptr ) GAUXC_GENERIC_EXCEPTION("SYCL Malloc Returned NULL"); + + return ptr; +} + +template +inline T* sycl_malloc_host( size_t n, ::sycl::queue& queue ) { + + T* ptr = nullptr; + GAUXC_SYCL_ERROR( "SYCL Malloc Host Failed", + ptr = ::sycl::malloc_host( n, queue ) ); + if( !ptr ) GAUXC_GENERIC_EXCEPTION("SYCL Malloc Host Returned NULL"); + + return ptr; +} + +/// Shared (host + device accessible) USM allocation. Needed for parameter +/// arrays that a host-side dispatch routine reads (e.g. oneMKL's variable-group +/// gemm_batch, which enumerates group boundaries before enqueueing), as +/// opposed to device-only scratch that only device kernels touch. +template +inline T* sycl_malloc_shared( size_t n, ::sycl::queue& queue ) { + + T* ptr = nullptr; + GAUXC_SYCL_ERROR( "SYCL Malloc Shared Failed", + ptr = ::sycl::malloc_shared( n, queue ) ); + if( !ptr ) GAUXC_GENERIC_EXCEPTION("SYCL Malloc Shared Returned NULL"); + + return ptr; +} + + + + + + +template +inline void sycl_free( ::sycl::queue& queue, T*& ptr ) { + // sycl::free does not synchronize the queue, and releasing a USM allocation + // that enqueued work still references is undefined behaviour. Drain the + // queue first. + GAUXC_SYCL_ERROR( "SYCL Free Sync Failed", queue.wait_and_throw() ); + GAUXC_SYCL_ERROR( "SYCL Free Failed", ::sycl::free( (void*)ptr, queue ) ); + ptr = nullptr; +} + +template +inline void sycl_free( ::sycl::queue& queue, T*& ptr, Args&&... args ) { + sycl_free(queue, ptr); + sycl_free(queue, std::forward(args)...); +} + + + + +template +inline void sycl_copy( size_t len, T* dest, const T* src, ::sycl::queue& queue, + std::string m = "") { + GAUXC_SYCL_ERROR( "SYCL Memcpy Failed ["+m+"]", + queue.memcpy( dest, src, len * sizeof(T) ).wait_and_throw() ); +} + +template +inline void sycl_copy_async( size_t len, T* dest, const T* src, + ::sycl::queue& queue, std::string m = "" ) { + GAUXC_SYCL_ERROR( "SYCL Memcpy Async Failed ["+m+"]", + queue.memcpy( dest, src, len * sizeof(T) ) ); +} + + +template +inline void sycl_copy_2d( T* dest, size_t dest_pitch, const T* src, + size_t src_pitch, size_t width, size_t height, + ::sycl::queue& queue, std::string m = "" ) { + GAUXC_SYCL_ERROR( "SYCL 2D Memcpy Failed ["+m+"]", + queue.ext_oneapi_memcpy2d( dest, dest_pitch, src, src_pitch, width, + height ).wait_and_throw() ); +} + +template +inline void sycl_copy_2d_async( T* dest, size_t dest_pitch, const T* src, + size_t src_pitch, size_t width, size_t height, + ::sycl::queue& queue, std::string m = "" ) { + GAUXC_SYCL_ERROR( "SYCL 2D Memcpy Async Failed ["+m+"]", + queue.ext_oneapi_memcpy2d( dest, dest_pitch, src, src_pitch, width, height ) ); +} + + +template +inline void sycl_set_zero( size_t len, T* ptr, ::sycl::queue& queue, + std::string m = "" ) { + GAUXC_SYCL_ERROR( "SYCL Memset Failed ["+m+"]", + queue.memset( ptr, 0, len * sizeof(T) ).wait_and_throw() ); +} + +template +inline void sycl_set_zero_async( size_t len, T* ptr, ::sycl::queue& queue, + std::string m = "" ) { + GAUXC_SYCL_ERROR( "SYCL Memset Async Failed ["+m+"]", + queue.memset( ptr, 0, len * sizeof(T) ) ); +} + + + +inline void sycl_device_sync( ::sycl::queue& queue ) { + GAUXC_SYCL_ERROR( "SYCL Device Sync Failed", queue.wait_and_throw() ); +} + + + +/// Largest work-group a kernel may be launched with on the queue's device. +/// SYCL has no per-kernel query prior to bundle specialization, so this is the +/// device-wide limit clamped by the GauXC block-size convention. +inline size_t sycl_kernel_max_threads_per_block( const ::sycl::queue& queue ) { + return queue.get_device().get_info< ::sycl::info::device::max_work_group_size >(); +} + +} +} + +#endif diff --git a/src/runtime_environment/device_specific/sycl_vector_types.hpp b/src/runtime_environment/device_specific/sycl_vector_types.hpp new file mode 100644 index 000000000..58ea5faa7 --- /dev/null +++ b/src/runtime_environment/device_specific/sycl_vector_types.hpp @@ -0,0 +1,26 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include + +namespace GauXC { + +/// Packed vector aggregates matching the layout / alignment of the CUDA and +/// HIP builtin vector types. sycl::vec is deliberately not used here: it pads +/// 3-component vectors out to 4 components, which would break the +/// reinterpret_cast reads of contiguous coordinate triples that the kernels +/// share with the CUDA backend. +struct alignas(16) double2 { double x, y; }; +struct alignas(8) double3 { double x, y, z; }; +struct alignas(4) int3 { int32_t x, y, z; }; + +} diff --git a/src/xc_integrator/local_work_driver/device/CMakeLists.txt b/src/xc_integrator/local_work_driver/device/CMakeLists.txt index 0911dff23..ec2addef6 100644 --- a/src/xc_integrator/local_work_driver/device/CMakeLists.txt +++ b/src/xc_integrator/local_work_driver/device/CMakeLists.txt @@ -35,3 +35,16 @@ endif() if(GAUXC_HAS_HIP) add_subdirectory( hip ) endif() + +if(GAUXC_HAS_SYCL) + # oneMKL's batched GEMM is the SYCL analogue of MAGMA's vbatched GEMM for + # the CUDA/HIP backends; unlike MAGMA it is not optional when SYCL is + # enabled, so these sources are added unconditionally alongside scheme1_base + # rather than gated behind their own GAUXC_ENABLE_* option. + target_sources( gauxc PRIVATE + scheme1_onemkl_base.cxx + scheme1_onemkl_data_base.cxx + ) + + add_subdirectory( sycl ) +endif() diff --git a/src/xc_integrator/local_work_driver/device/scheme1_data_base.cxx b/src/xc_integrator/local_work_driver/device/scheme1_data_base.cxx index 7818a5a83..ce19b353b 100644 --- a/src/xc_integrator/local_work_driver/device/scheme1_data_base.cxx +++ b/src/xc_integrator/local_work_driver/device/scheme1_data_base.cxx @@ -11,6 +11,10 @@ */ #include "scheme1_data_base.hpp" #include "buffer_adaptor.hpp" +#ifdef GAUXC_HAS_SYCL +// CUDA and HIP supply double2 through their runtime headers; SYCL does not +#include "device_specific/sycl_vector_types.hpp" +#endif namespace GauXC { diff --git a/src/xc_integrator/local_work_driver/device/scheme1_onemkl_base.cxx b/src/xc_integrator/local_work_driver/device/scheme1_onemkl_base.cxx new file mode 100644 index 000000000..8b461ae77 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/scheme1_onemkl_base.cxx @@ -0,0 +1,202 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "scheme1_onemkl_base.hpp" +#include "device/common/pack_submat.hpp" +#include "device/common/inc_potential.hpp" +#include "device/common/device_blas.hpp" +#include "device/sycl/sycl_backend.hpp" + +// Only the BLAS domain is needed. The umbrella also pulls in +// the C LAPACK prototypes, whose parameter named "csl" collides with the csl +// macro from buffer_adaptor.hpp +#include + +#ifdef GAUXC_HAS_CUDA +#define GAUXC_ENABLE_EXX +#endif + +namespace GauXC { + +namespace { + // onemkl_data stores transpose flags as plain char (see scheme1_onemkl_base.hpp + // for why); reinterpret to the real oneMKL type at the call site. + inline oneapi::mkl::transpose* as_transpose( char* p ) { + return reinterpret_cast(p); + } + inline const oneapi::mkl::transpose* as_transpose( const char* p ) { + return reinterpret_cast(p); + } +} + +void AoSScheme1OneMKLBase::eval_xmat( double fac, XCDeviceData* _data, bool do_grad, density_id den ){ + + if( do_grad ) GAUXC_GENERIC_EXCEPTION("oneMKL + X Gradient NYI"); + + auto* data = dynamic_cast(_data); + if( !data ) GAUXC_BAD_LWD_DATA_CAST(); + + if( not data->device_backend_ ) GAUXC_UNINITIALIZED_DEVICE_BACKEND(); + + auto tasks = data->host_device_tasks; + const auto ntasks = tasks.size(); + + // Pack density matrix + const auto nbf = data->global_dims.nbf; + const auto submat_block_size = data->get_submat_chunk_size( nbf, 0 ); + auto static_stack = data->static_stack; + auto aos_stack = data->aos_stack; + auto* backend = dynamic_cast(data->device_backend_); + if( !backend ) GAUXC_BAD_BACKEND_CAST(); + ::sycl::queue& master_queue = backend->master_stream->queue; + auto& onemkl_stack = data->onemkl_stack; + double* dmat_ptr = nullptr; + switch( den ) { + case DEN_S: + dmat_ptr = static_stack.dmat_s_device; + break; + case DEN_Z: + dmat_ptr = static_stack.dmat_z_device; + break; + case DEN_Y: + dmat_ptr = static_stack.dmat_y_device; + break; + case DEN_X: + dmat_ptr = static_stack.dmat_x_device; + break; + default: + GAUXC_GENERIC_EXCEPTION( "eval_xmat called with invalid density specifier" ); + } + sym_pack_submat( ntasks, aos_stack.device_tasks, dmat_ptr, + nbf, submat_block_size, data->device_backend_->queue() ); + + // Update dmat pointers on onemkl_stack. xdmat_array_shared is USM-shared, + // so this can be written directly from the host instead of routed through + // device_backend_->copy_async as MAGMA's device-only array requires. + for( decltype(tasks.size()) i = 0; i < ntasks; i++ ) { + auto& task = tasks[i]; + if( task.bfn_screening.ncut > 1 ) { + onemkl_stack.xdmat_array_shared[i] = task.nbe_scr; + } else { + onemkl_stack.xdmat_array_shared[i] = dmat_ptr + task.bfn_screening.ibf_begin*(nbf+1); + } + } + + // fac is only known at this call site (e.g. 2.0 for RKS), unlike beta which + // is always 0 here, so alpha is refilled on every call; groupsize, m/n/k, + // leading dimensions, and the (constant) transpose flags were already + // filled once in pack_and_send_xmat. + for( decltype(tasks.size()) i = 0; i < ntasks; i++ ) + onemkl_stack.xmat_alpha_array_shared[i] = fac; + + // MAGMA's vbatched GEMM takes one m/n/k per matrix. oneMKL's group-form + // gemm_batch instead takes one m/n/k/transpose per GROUP plus a groupsize + // array; "every matrix has its own size" is expressed by making each task + // its own group of size 1 (group_count == ntasks, groupsize[i] == 1). + GAUXC_ONEMKL_ERROR( "ONEMKL DGEMM_BATCH (XMAT) FAILED", + oneapi::mkl::blas::column_major::gemm_batch( master_queue, + as_transpose(onemkl_stack.xmat_transa_array_shared), + as_transpose(onemkl_stack.xmat_transb_array_shared), + onemkl_stack.xmat_m_array_shared, onemkl_stack.xmat_n_array_shared, + onemkl_stack.xmat_k_array_shared, onemkl_stack.xmat_alpha_array_shared, + (const double**)onemkl_stack.bf_array_shared, onemkl_stack.ld_bf_array_shared, + (const double**)onemkl_stack.xdmat_array_shared, onemkl_stack.ld_xdmat_array_shared, + onemkl_stack.xmat_beta_array_shared, onemkl_stack.zmat_array_shared, onemkl_stack.ld_zmat_array_shared, + ntasks, onemkl_stack.groupsize_array_shared ).wait() ); + +} + +void AoSScheme1OneMKLBase::eval_exx_fmat( XCDeviceData* _data ) { +#ifndef GAUXC_ENABLE_EXX + GAUXC_GENERIC_EXCEPTION("EXX + non-CUDA NYI"); +#else + auto* data = dynamic_cast(_data); + if( !data ) GAUXC_BAD_LWD_DATA_CAST(); + + if( not data->device_backend_ ) GAUXC_UNINITIALIZED_DEVICE_BACKEND(); + + auto tasks = data->host_device_tasks; + const auto ntasks = tasks.size(); + + // Pack the density matrix into (bfn, cou) shape + const auto nbf = data->global_dims.nbf; + const auto submat_block_size = data->get_submat_chunk_size( nbf, 0 ); + auto static_stack = data->static_stack; + auto aos_stack = data->aos_stack; + asym_pack_submat( ntasks, aos_stack.device_tasks, static_stack.dmat_s_device, + nbf, submat_block_size, data->device_backend_->queue() ); + + auto* backend = dynamic_cast(data->device_backend_); + if( !backend ) GAUXC_BAD_BACKEND_CAST(); + ::sycl::queue& master_queue = backend->master_stream->queue; + auto& onemkl_stack = data->onemkl_stack; + GAUXC_ONEMKL_ERROR( "ONEMKL DGEMM_BATCH (FMAT) FAILED", + oneapi::mkl::blas::column_major::gemm_batch( master_queue, + as_transpose(onemkl_stack.fmat_transa_array_shared), + as_transpose(onemkl_stack.fmat_transb_array_shared), + onemkl_stack.fmat_m_array_shared, onemkl_stack.fmat_n_array_shared, + onemkl_stack.fmat_k_array_shared, onemkl_stack.fmat_alpha_array_shared, + (const double**)onemkl_stack.bf_array_shared, onemkl_stack.ld_bf_array_shared, + (const double**)onemkl_stack.fdmat_array_shared, onemkl_stack.ld_fdmat_array_shared, + onemkl_stack.fmat_beta_array_shared, onemkl_stack.fmat_array_shared, onemkl_stack.ld_fmat_array_shared, + ntasks, onemkl_stack.groupsize_array_shared ).wait() ); +#endif +} + +// inc_vxc is intentionally NOT overridden here: there is no batched syr2k in +// oneMKL (only syrk_batch, see usm_decls.hpp), and the correct decomposition +// of syr2k C := alpha*(A^T B + B^T A) + beta*C into two GEMMs is exactly what +// AoSScheme1Base::inc_vxc already does via inc_potential_impl's round-robin +// do_syr2k lambda, which itself calls the device_blas syr2k() entry point +// already specialized for oneMKL in kernels/onemkl_extensions.cxx (over the +// BLAS stream pool). Substituting anything else here would be a numerically +// different, and still unbatched (no group syr2k exists to batch it with), +// implementation for no benefit, so the base class implementation stands. + +void AoSScheme1OneMKLBase::inc_exx_k( XCDeviceData* _data){ +#ifndef GAUXC_ENABLE_EXX + GAUXC_GENERIC_EXCEPTION("EXX + non-CUDA NYI"); +#else + auto* data = dynamic_cast(_data); + if( !data ) GAUXC_BAD_LWD_DATA_CAST(); + + if( not data->device_backend_ ) GAUXC_UNINITIALIZED_DEVICE_BACKEND(); + + auto& tasks = data->host_device_tasks; + const auto ntasks = tasks.size(); + + auto* backend = dynamic_cast(data->device_backend_); + if( !backend ) GAUXC_BAD_BACKEND_CAST(); + ::sycl::queue& master_queue = backend->master_stream->queue; + auto& onemkl_stack = data->onemkl_stack; + GAUXC_ONEMKL_ERROR( "ONEMKL DGEMM_BATCH (EXX_K) FAILED", + oneapi::mkl::blas::column_major::gemm_batch( master_queue, + as_transpose(onemkl_stack.kmat_transa_array_shared), + as_transpose(onemkl_stack.kmat_transb_array_shared), + onemkl_stack.fmat_k_array_shared, onemkl_stack.fmat_n_array_shared, + onemkl_stack.fmat_m_array_shared, onemkl_stack.fmat_alpha_array_shared, + (const double**)onemkl_stack.bf_array_shared, onemkl_stack.ld_bf_array_shared, + (const double**)onemkl_stack.gmat_array_shared, onemkl_stack.ld_fmat_array_shared, + onemkl_stack.fmat_beta_array_shared, onemkl_stack.kmat_array_shared, onemkl_stack.ld_fdmat_array_shared, + ntasks, onemkl_stack.groupsize_array_shared ).wait() ); + + // Increment EXX_K + const auto nbf = data->global_dims.nbf; + const auto submat_block_size = data->get_submat_chunk_size( nbf, 0 ); + auto static_stack = data->static_stack; + auto aos_stack = data->aos_stack; + asym_task_inc_potential( ntasks, aos_stack.device_tasks, + static_stack.exx_k_device, nbf, submat_block_size, + data->device_backend_->queue() ); +#endif +} + +} diff --git a/src/xc_integrator/local_work_driver/device/scheme1_onemkl_base.hpp b/src/xc_integrator/local_work_driver/device/scheme1_onemkl_base.hpp new file mode 100644 index 000000000..8fe05f6c0 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/scheme1_onemkl_base.hpp @@ -0,0 +1,153 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "scheme1_base.hpp" + +namespace GauXC { + +struct AoSScheme1OneMKLBase : public AoSScheme1Base { + + void eval_xmat( double fac, XCDeviceData*, bool do_grad, density_id den ) override final; + void eval_exx_fmat( XCDeviceData* ) override final; + void inc_exx_k( XCDeviceData* ) override final; + + struct Data; + + virtual ~AoSScheme1OneMKLBase() = default; +}; + +struct AoSScheme1OneMKLBase::Data : public AoSScheme1Base::Data { + + using base_type = AoSScheme1Base::Data; + using base_type::host_task_type; + using base_type::device_buffer_t; + + // oneMKL only exposes a variable-size batched GEMM through the "group" form + // of gemm_batch: it takes one m/n/k/transpose/alpha/beta PER GROUP plus a + // groupsize array, not one m/n/k per matrix the way MAGMA's vbatched + // routines do. To express "every matrix in the batch has its own size" + // (exactly what MAGMA's per-matrix arrays give for free), each task is + // modeled as its own group of size 1: group_count == ntasks and + // groupsize[i] == 1 for all i. That promotes transa/transb/alpha/beta from + // MAGMA's single scalars-for-the-whole-batch to per-task arrays here. + // + // gemm_batch's group form is a host-side dispatch call: it walks + // transa/transb/m/n/k/alpha/beta/lda/ldb/ldc/groupsize -- and the arrays of + // matrix pointers themselves (a/b/c) -- on the host to build the batched + // command lists before anything is enqueued to the device. The reference + // oneMKL SYCL example (oneMKL's examples_sycl.tgz, sycl/blas/source/gemm_batch.cpp) + // allocates every one of those parameter arrays, INCLUDING the + // pointer-to-pointer matrix arrays a/b/c, with sycl::malloc_shared for + // exactly this reason. Device-only (malloc_device) memory, which is what + // the rest of this dynamic stack uses via buffer_adaptor, is not + // host-readable and would make that host-side walk undefined behavior. + // Every array passed directly to gemm_batch is therefore USM-shared here + // (suffix "_shared"); only the matrix DATA the pointers point to (task.bf, + // task.zmat, static_stack.dmat_s_device, ...) stays device-resident exactly + // as it already is elsewhere in GauXC, since only the batched kernel itself + // (running on-device) ever dereferences that data. + struct onemkl_data { + // Arrays of matrix pointers passed as gemm_batch's a/b/c. Shared USM: the + // pointer VALUES must be host-readable for gemm_batch's dispatch, even + // though each pointer refers to device-resident matrix data. + double** xdmat_array_shared = nullptr; + double** fdmat_array_shared = nullptr; + double** vmat_array_shared = nullptr; + double** kmat_array_shared = nullptr; + double** zmat_array_shared = nullptr; + double** fmat_array_shared = nullptr; + double** gmat_array_shared = nullptr; + double** bf_array_shared = nullptr; + + // Per-group (== per-task) sizes and leading dimensions. + int32_t* xmat_m_array_shared = nullptr; + int32_t* xmat_n_array_shared = nullptr; + int32_t* xmat_k_array_shared = nullptr; + int32_t* ld_xdmat_array_shared = nullptr; + + int32_t* fmat_m_array_shared = nullptr; + int32_t* fmat_n_array_shared = nullptr; + int32_t* fmat_k_array_shared = nullptr; + int32_t* ld_fdmat_array_shared = nullptr; + + int32_t* ld_vmat_array_shared = nullptr; + int32_t* ld_zmat_array_shared = nullptr; + int32_t* ld_fmat_array_shared = nullptr; + int32_t* ld_bf_array_shared = nullptr; + + // groupsize[i] == 1 for all i (each task is its own group); allocated and + // filled once since it never changes across calls or task batches. + int32_t* groupsize_array_shared = nullptr; + + // Constant per-group transpose flags. eval_xmat/eval_exx_fmat/inc_exx_k + // each use one fixed (transa,transb) pair across the whole batch, but the + // group API still requires a per-group array, so one array per call site + // is filled once and reused on every call. + // Stored as plain `char` (the underlying type of oneapi::mkl::transpose, + // an enum class : char) rather than the oneMKL enum itself, so this + // header -- like scheme1_magma_base.hpp -- stays free of any + // backend-specific #include; the .cxx reinterpret_casts these to + // oneapi::mkl::transpose* at the gemm_batch call site. + char* xmat_transa_array_shared = nullptr; + char* xmat_transb_array_shared = nullptr; + char* fmat_transa_array_shared = nullptr; + char* fmat_transb_array_shared = nullptr; + char* kmat_transa_array_shared = nullptr; + char* kmat_transb_array_shared = nullptr; + + // alpha varies per call (eval_xmat's fac is a call-site argument), so it + // is refilled in eval_xmat itself on every call; beta is always 0 for + // these calls but still needs a per-group array to satisfy the API, so + // it is filled once in pack_and_send, like groupsize and the transposes. + double* xmat_alpha_array_shared = nullptr; + double* xmat_beta_array_shared = nullptr; + double* fmat_alpha_array_shared = nullptr; + double* fmat_beta_array_shared = nullptr; + + inline void reset(){ std::memset(this,0,sizeof(onemkl_data)); } + }; + + onemkl_data onemkl_stack; + + template + Data( Args&&... args ) : base_type( std::forward(args)... ) { } + + virtual ~Data(); + + // get_mem_req/get_static_mem_requirement are deliberately NOT overridden: + // unlike MAGMA's arrays, every onemkl_stack field is USM-shared memory + // allocated directly with sycl::malloc_shared in allocate_dynamic_stack + // (see .cxx), not carved out of the buffer_adaptor dynamic stack that + // get_mem_req/get_static_mem_requirement account for. It therefore does not + // affect, and must not be added to, the task-batch-size accounting those + // two methods perform against the fixed device buffer. + void reset_allocations() override final; + device_buffer_t allocate_dynamic_stack( integrator_term_tracker terms, + host_task_iterator begin, host_task_iterator end, device_buffer_t buf ) + override final; + void pack_and_send( integrator_term_tracker terms, + host_task_iterator begin, host_task_iterator end, + const BasisSetMap& basis_map ) override final; + + + void pack_and_send_xmat( host_task_iterator, host_task_iterator ); + void pack_and_send_fmat( host_task_iterator, host_task_iterator ); + +private: + // All of onemkl_data's fields are USM-shared allocations made outside the + // device-only buffer_adaptor dynamic stack (see allocate_dynamic_stack), + // so they need their own explicit sycl::free at teardown; done from the + // .cxx (the only place a SYCL queue is available to free them). + void free_shared_arrays(); +}; + +} diff --git a/src/xc_integrator/local_work_driver/device/scheme1_onemkl_data_base.cxx b/src/xc_integrator/local_work_driver/device/scheme1_onemkl_data_base.cxx new file mode 100644 index 000000000..3835fa23c --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/scheme1_onemkl_data_base.cxx @@ -0,0 +1,256 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "scheme1_onemkl_base.hpp" +#include "buffer_adaptor.hpp" +#include "device/sycl/sycl_backend.hpp" +#include "device_specific/sycl_util.hpp" +#include + +// Only the BLAS domain is needed for oneapi::mkl::transpose -- see +// scheme1_onemkl_base.cxx for why the umbrella is avoided. +#include + +namespace GauXC { + +AoSScheme1OneMKLBase::Data::~Data() { + free_shared_arrays(); +} + +// onemkl_stack holds no device-only allocations of its own (unlike MAGMA's +// magma_stack, whose double**/int32_t* arrays live in the buffer_adaptor +// dynamic stack) -- every field here is USM-shared memory allocated and +// freed independently in allocate_dynamic_stack/free_shared_arrays. There is +// therefore nothing for reset_allocations to hand back to the dynamic stack; +// it only needs to release the shared allocations from the previous task +// batch before the next one is sized, exactly mirroring what +// base_type::reset_allocations() does for the (device-only) base stacks. +void AoSScheme1OneMKLBase::Data::reset_allocations() { + base_type::reset_allocations(); + free_shared_arrays(); + onemkl_stack.reset(); +} + +void AoSScheme1OneMKLBase::Data::free_shared_arrays() { + auto* backend = dynamic_cast(device_backend_); + if( !backend ) return; // Not yet bound to a backend, or already torn down + auto& queue = backend->master_stream->queue; + auto& s = onemkl_stack; + + auto free_if = [&]( auto*& ptr ) { + if( ptr ) { util::sycl_free( queue, ptr ); } + }; + + free_if(s.xdmat_array_shared); free_if(s.fdmat_array_shared); + free_if(s.vmat_array_shared); free_if(s.kmat_array_shared); + free_if(s.zmat_array_shared); free_if(s.fmat_array_shared); + free_if(s.gmat_array_shared); free_if(s.bf_array_shared); + + free_if(s.xmat_m_array_shared); free_if(s.xmat_n_array_shared); + free_if(s.xmat_k_array_shared); free_if(s.ld_xdmat_array_shared); + + free_if(s.fmat_m_array_shared); free_if(s.fmat_n_array_shared); + free_if(s.fmat_k_array_shared); free_if(s.ld_fdmat_array_shared); + + free_if(s.ld_vmat_array_shared); free_if(s.ld_zmat_array_shared); + free_if(s.ld_fmat_array_shared); free_if(s.ld_bf_array_shared); + + free_if(s.groupsize_array_shared); + + free_if(s.xmat_transa_array_shared); free_if(s.xmat_transb_array_shared); + free_if(s.fmat_transa_array_shared); free_if(s.fmat_transb_array_shared); + free_if(s.kmat_transa_array_shared); free_if(s.kmat_transb_array_shared); + + free_if(s.xmat_alpha_array_shared); free_if(s.xmat_beta_array_shared); + free_if(s.fmat_alpha_array_shared); free_if(s.fmat_beta_array_shared); +} + +// allocate_dynamic_stack is otherwise the point at which per-batch device +// memory is carved out of the fixed device_buffer_t handed down from +// XCDeviceStackData::generate_buffers. onemkl_stack's arrays are excluded +// from that accounting (see the .hpp) because the group-form gemm_batch +// requires them to be host-readable USM-shared memory, which buffer_adaptor +// -- built to slice a single device-only allocation -- cannot provide. +// Instead they are allocated directly here via sycl::malloc_shared, sized to +// the same ntask this task batch was already sized to by the base class. +AoSScheme1OneMKLBase::Data::device_buffer_t + AoSScheme1OneMKLBase::Data::allocate_dynamic_stack( + integrator_term_tracker terms, + host_task_iterator task_begin, host_task_iterator task_end, + device_buffer_t buf ){ + + // Allocate base info on the stack + buf = base_type::allocate_dynamic_stack( terms, task_begin, task_end, + buf ); + + required_term_storage reqt(terms); + + auto* backend = dynamic_cast(device_backend_); + if( !backend ) GAUXC_BAD_BACKEND_CAST(); + auto& queue = backend->master_stream->queue; + + // A single operation may need more than one task batch to fit the fixed + // device dynamic-stack budget, meaning allocate_dynamic_stack can be called + // more than once between two reset_allocations() calls. Unlike MAGMA's + // buffer_adaptor slices (harmlessly re-carved from the same fixed device + // buffer every call), onemkl_stack's fields are independent + // sycl::malloc_shared allocations, so the previous batch's arrays must be + // freed here -- not only in reset_allocations() -- or every batch after the + // first in a multi-batch operation leaks. + free_shared_arrays(); + + const auto ntask = std::distance( task_begin, task_end ); + auto& s = onemkl_stack; + + // bf_array_shared and groupsize_array_shared are shared between the xmat + // and fmat call sites (same bf pointers, same "every task is its own + // group of 1" groupsize semantics), so -- unlike MAGMA's buffer_adaptor + // slices, where allocating twice just wastes stack space -- they must only + // be sycl::malloc_shared'd once or the first allocation leaks. + if(reqt.task_xmat or reqt.task_fmat) { + s.bf_array_shared = util::sycl_malloc_shared(ntask, queue); + s.ld_bf_array_shared = util::sycl_malloc_shared(ntask, queue); + s.groupsize_array_shared = util::sycl_malloc_shared(ntask, queue); + } + + if(reqt.task_xmat) { + s.xdmat_array_shared = util::sycl_malloc_shared(ntask, queue); + s.zmat_array_shared = util::sycl_malloc_shared(ntask, queue); + + s.xmat_m_array_shared = util::sycl_malloc_shared(ntask, queue); + s.xmat_n_array_shared = util::sycl_malloc_shared(ntask, queue); + s.xmat_k_array_shared = util::sycl_malloc_shared(ntask, queue); + s.ld_xdmat_array_shared = util::sycl_malloc_shared(ntask, queue); + s.ld_zmat_array_shared = util::sycl_malloc_shared(ntask, queue); + + s.xmat_transa_array_shared = util::sycl_malloc_shared(ntask, queue); + s.xmat_transb_array_shared = util::sycl_malloc_shared(ntask, queue); + s.xmat_alpha_array_shared = util::sycl_malloc_shared(ntask, queue); + s.xmat_beta_array_shared = util::sycl_malloc_shared(ntask, queue); + } + + if(reqt.task_fmat) { + s.fdmat_array_shared = util::sycl_malloc_shared(ntask, queue); + s.fmat_array_shared = util::sycl_malloc_shared(ntask, queue); + s.gmat_array_shared = util::sycl_malloc_shared(ntask, queue); + s.kmat_array_shared = util::sycl_malloc_shared(ntask, queue); + + s.fmat_m_array_shared = util::sycl_malloc_shared(ntask, queue); + s.fmat_n_array_shared = util::sycl_malloc_shared(ntask, queue); + s.fmat_k_array_shared = util::sycl_malloc_shared(ntask, queue); + s.ld_fdmat_array_shared = util::sycl_malloc_shared(ntask, queue); + s.ld_fmat_array_shared = util::sycl_malloc_shared(ntask, queue); + + s.fmat_transa_array_shared = util::sycl_malloc_shared(ntask, queue); + s.fmat_transb_array_shared = util::sycl_malloc_shared(ntask, queue); + s.kmat_transa_array_shared = util::sycl_malloc_shared(ntask, queue); + s.kmat_transb_array_shared = util::sycl_malloc_shared(ntask, queue); + s.fmat_alpha_array_shared = util::sycl_malloc_shared(ntask, queue); + s.fmat_beta_array_shared = util::sycl_malloc_shared(ntask, queue); + } + + // Nothing was taken from the device-only dynamic stack, so hand it back + // untouched for further derived-class allocations. + return buf; +} + +void AoSScheme1OneMKLBase::Data::pack_and_send( + integrator_term_tracker terms, + host_task_iterator task_begin, host_task_iterator task_end, + const BasisSetMap& basis_map ) { + + + base_type::pack_and_send( terms, task_begin, task_end, basis_map ); + required_term_storage reqt(terms); + + if(reqt.task_xmat) pack_and_send_xmat(task_begin, task_end); + if(reqt.task_fmat) pack_and_send_fmat(task_begin, task_end); + +} + + +void AoSScheme1OneMKLBase::Data::pack_and_send_xmat( + host_task_iterator task_begin, host_task_iterator task_end +) { + + const auto ntask = std::distance( task_begin, task_end ); + auto& s = onemkl_stack; + + const auto nbf = global_dims.nbf; + + // Every field written here is USM-shared, so it can be filled directly + // from the host without a device_backend_->copy_async round trip -- unlike + // MAGMA's device-only magma_stack, which requires exactly that. + // + // xdmat_array_shared is intentionally NOT filled here: the density pointer + // depends on which of DEN_S/DEN_Z/DEN_Y/DEN_X eval_xmat is called for, so + // eval_xmat fills it (and xmat_alpha_array_shared, since alpha == fac is + // also only known at that call site) itself, on every call. + for( std::decay_t i = 0; i < ntask; ++i ) { + auto& task = host_device_tasks[i]; + s.zmat_array_shared[i] = task.zmat; s.ld_zmat_array_shared[i] = task.npts; + s.bf_array_shared[i] = task.bf; s.ld_bf_array_shared[i] = task.npts; + + if( task.bfn_screening.ncut > 1 ) { + s.ld_xdmat_array_shared[i] = task.bfn_screening.nbe; + } else { + s.ld_xdmat_array_shared[i] = nbf; + } + + s.xmat_m_array_shared[i] = task.npts; + s.xmat_n_array_shared[i] = task.bfn_screening.nbe; + s.xmat_k_array_shared[i] = task.bfn_screening.nbe; + + s.groupsize_array_shared[i] = 1; + s.xmat_transa_array_shared[i] = static_cast(oneapi::mkl::transpose::nontrans); + s.xmat_transb_array_shared[i] = static_cast(oneapi::mkl::transpose::nontrans); + s.xmat_beta_array_shared[i] = 0.; + } + +} + + + +void AoSScheme1OneMKLBase::Data::pack_and_send_fmat( + host_task_iterator task_begin, host_task_iterator task_end +) { + + const auto ntask = std::distance( task_begin, task_end ); + auto& s = onemkl_stack; + + // host_device_tasks should be populated by parent impl called at top + for( std::decay_t i = 0; i < ntask; ++i ) { + auto& task = host_device_tasks[i]; + s.fmat_array_shared[i] = task.fmat; s.ld_fmat_array_shared[i] = task.npts; + s.gmat_array_shared[i] = task.gmat; + s.bf_array_shared[i] = task.bf; s.ld_bf_array_shared[i] = task.npts; + s.kmat_array_shared[i] = task.nbe_scr; + + s.fdmat_array_shared[i] = task.nbe_scr; + s.ld_fdmat_array_shared[i] = task.bfn_screening.nbe; + + s.fmat_m_array_shared[i] = task.npts; + s.fmat_n_array_shared[i] = task.cou_screening.nbe; + s.fmat_k_array_shared[i] = task.bfn_screening.nbe; + + s.groupsize_array_shared[i] = 1; + s.fmat_transa_array_shared[i] = static_cast(oneapi::mkl::transpose::nontrans); + s.fmat_transb_array_shared[i] = static_cast(oneapi::mkl::transpose::nontrans); + s.fmat_beta_array_shared[i] = 0.; + + // inc_exx_k reuses these same groupsize/beta arrays with its own + // (transposed) transa/transb pair. + s.kmat_transa_array_shared[i] = static_cast(oneapi::mkl::transpose::trans); + s.kmat_transb_array_shared[i] = static_cast(oneapi::mkl::transpose::nontrans); + } + +} +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/CMakeLists.txt b/src/xc_integrator/local_work_driver/device/sycl/CMakeLists.txt new file mode 100644 index 000000000..a30a72df0 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/CMakeLists.txt @@ -0,0 +1,34 @@ +# +# GauXC Copyright (c) 2020-2024, The Regents of the University of California, +# through Lawrence Berkeley National Laboratory (subject to receipt of +# any required approvals from the U.S. Dept. of Energy). +# +# (c) 2024-2025, Microsoft Corporation +# +# All rights reserved. +# +# See LICENSE.txt for details +# +set( GAUXC_SYCL_LWD_SOURCES + sycl_aos_scheme1_data.cxx + sycl_aos_scheme1.cxx + sycl_aos_scheme1_weights.cxx + + xc_functional_eval_wrapper.cxx + + kernels/collocation_device.cxx + kernels/grid_to_center.cxx + kernels/sycl_ssf_1d.cxx + kernels/sycl_ssf_2d.cxx + kernels/pack_submat.cxx + kernels/onemkl_extensions.cxx + kernels/uvvars.cxx + kernels/zmat_vxc.cxx + kernels/zmat_fxc.cxx + kernels/sycl_inc_potential.cxx + kernels/symmetrize_mat.cxx + kernels/increment_exc_grad.cxx +) + +target_sources(gauxc PRIVATE ${GAUXC_SYCL_LWD_SOURCES}) +gauxc_sycl_device_sources( ${GAUXC_SYCL_LWD_SOURCES} ) diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_angular_cartesian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_angular_cartesian.hpp new file mode 100644 index 000000000..ecb2b233e --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_angular_cartesian.hpp @@ -0,0 +1,418 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include +#include + +#ifndef GPGAUEVAL_INLINE +# define GPGAUEVAL_INLINE inline +#endif + +namespace GauXC { + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_0( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = bf; + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_0_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = bf_x; + + eval_y[npts * 0] = bf_y; + + eval_z[npts * 0] = bf_z; + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_1( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = bf*x; + eval[npts * 1] = bf*y; + eval[npts * 2] = bf*z; + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_1_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = bf + bf_x*x; + eval_x[npts * 1] = bf_x*y; + eval_x[npts * 2] = bf_x*z; + + eval_y[npts * 0] = bf_y*x; + eval_y[npts * 1] = bf + bf_y*y; + eval_y[npts * 2] = bf_y*z; + + eval_z[npts * 0] = bf_z*x; + eval_z[npts * 1] = bf_z*y; + eval_z[npts * 2] = bf + bf_z*z; + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_2( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = bf*x*x; + eval[npts * 1] = bf*x*y; + eval[npts * 2] = bf*x*z; + eval[npts * 3] = bf*y*y; + eval[npts * 4] = bf*y*z; + eval[npts * 5] = bf*z*z; + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_2_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = x*(2*bf + bf_x*x); + eval_x[npts * 1] = y*(bf + bf_x*x); + eval_x[npts * 2] = z*(bf + bf_x*x); + eval_x[npts * 3] = bf_x*y*y; + eval_x[npts * 4] = bf_x*y*z; + eval_x[npts * 5] = bf_x*z*z; + + eval_y[npts * 0] = bf_y*x*x; + eval_y[npts * 1] = x*(bf + bf_y*y); + eval_y[npts * 2] = bf_y*x*z; + eval_y[npts * 3] = y*(2*bf + bf_y*y); + eval_y[npts * 4] = z*(bf + bf_y*y); + eval_y[npts * 5] = bf_y*z*z; + + eval_z[npts * 0] = bf_z*x*x; + eval_z[npts * 1] = bf_z*x*y; + eval_z[npts * 2] = x*(bf + bf_z*z); + eval_z[npts * 3] = bf_z*y*y; + eval_z[npts * 4] = y*(bf + bf_z*z); + eval_z[npts * 5] = z*(2*bf + bf_z*z); + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_3( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = bf*x*x*x; + eval[npts * 1] = bf*x*x*y; + eval[npts * 2] = bf*x*x*z; + eval[npts * 3] = bf*x*y*y; + eval[npts * 4] = bf*x*y*z; + eval[npts * 5] = bf*x*z*z; + eval[npts * 6] = bf*y*y*y; + eval[npts * 7] = bf*y*y*z; + eval[npts * 8] = bf*y*z*z; + eval[npts * 9] = bf*z*z*z; + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_3_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = x*x*(3*bf + bf_x*x); + eval_x[npts * 1] = x*y*(2*bf + bf_x*x); + eval_x[npts * 2] = x*z*(2*bf + bf_x*x); + eval_x[npts * 3] = y*y*(bf + bf_x*x); + eval_x[npts * 4] = y*z*(bf + bf_x*x); + eval_x[npts * 5] = z*z*(bf + bf_x*x); + eval_x[npts * 6] = bf_x*y*y*y; + eval_x[npts * 7] = bf_x*y*y*z; + eval_x[npts * 8] = bf_x*y*z*z; + eval_x[npts * 9] = bf_x*z*z*z; + + eval_y[npts * 0] = bf_y*x*x*x; + eval_y[npts * 1] = x*x*(bf + bf_y*y); + eval_y[npts * 2] = bf_y*x*x*z; + eval_y[npts * 3] = x*y*(2*bf + bf_y*y); + eval_y[npts * 4] = x*z*(bf + bf_y*y); + eval_y[npts * 5] = bf_y*x*z*z; + eval_y[npts * 6] = y*y*(3*bf + bf_y*y); + eval_y[npts * 7] = y*z*(2*bf + bf_y*y); + eval_y[npts * 8] = z*z*(bf + bf_y*y); + eval_y[npts * 9] = bf_y*z*z*z; + + eval_z[npts * 0] = bf_z*x*x*x; + eval_z[npts * 1] = bf_z*x*x*y; + eval_z[npts * 2] = x*x*(bf + bf_z*z); + eval_z[npts * 3] = bf_z*x*y*y; + eval_z[npts * 4] = x*y*(bf + bf_z*z); + eval_z[npts * 5] = x*z*(2*bf + bf_z*z); + eval_z[npts * 6] = bf_z*y*y*y; + eval_z[npts * 7] = y*y*(bf + bf_z*z); + eval_z[npts * 8] = y*z*(2*bf + bf_z*z); + eval_z[npts * 9] = z*z*(3*bf + bf_z*z); + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_4( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = bf*x*x*x*x; + eval[npts * 1] = bf*x*x*x*y; + eval[npts * 2] = bf*x*x*x*z; + eval[npts * 3] = bf*x*x*y*y; + eval[npts * 4] = bf*x*x*y*z; + eval[npts * 5] = bf*x*x*z*z; + eval[npts * 6] = bf*x*y*y*y; + eval[npts * 7] = bf*x*y*y*z; + eval[npts * 8] = bf*x*y*z*z; + eval[npts * 9] = bf*x*z*z*z; + eval[npts * 10] = bf*y*y*y*y; + eval[npts * 11] = bf*y*y*y*z; + eval[npts * 12] = bf*y*y*z*z; + eval[npts * 13] = bf*y*z*z*z; + eval[npts * 14] = bf*z*z*z*z; + +} + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_4_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = x*x*x*(4*bf + bf_x*x); + eval_x[npts * 1] = x*x*y*(3*bf + bf_x*x); + eval_x[npts * 2] = x*x*z*(3*bf + bf_x*x); + eval_x[npts * 3] = x*y*y*(2*bf + bf_x*x); + eval_x[npts * 4] = x*y*z*(2*bf + bf_x*x); + eval_x[npts * 5] = x*z*z*(2*bf + bf_x*x); + eval_x[npts * 6] = y*y*y*(bf + bf_x*x); + eval_x[npts * 7] = y*y*z*(bf + bf_x*x); + eval_x[npts * 8] = y*z*z*(bf + bf_x*x); + eval_x[npts * 9] = z*z*z*(bf + bf_x*x); + eval_x[npts * 10] = bf_x*y*y*y*y; + eval_x[npts * 11] = bf_x*y*y*y*z; + eval_x[npts * 12] = bf_x*y*y*z*z; + eval_x[npts * 13] = bf_x*y*z*z*z; + eval_x[npts * 14] = bf_x*z*z*z*z; + + eval_y[npts * 0] = bf_y*x*x*x*x; + eval_y[npts * 1] = x*x*x*(bf + bf_y*y); + eval_y[npts * 2] = bf_y*x*x*x*z; + eval_y[npts * 3] = x*x*y*(2*bf + bf_y*y); + eval_y[npts * 4] = x*x*z*(bf + bf_y*y); + eval_y[npts * 5] = bf_y*x*x*z*z; + eval_y[npts * 6] = x*y*y*(3*bf + bf_y*y); + eval_y[npts * 7] = x*y*z*(2*bf + bf_y*y); + eval_y[npts * 8] = x*z*z*(bf + bf_y*y); + eval_y[npts * 9] = bf_y*x*z*z*z; + eval_y[npts * 10] = y*y*y*(4*bf + bf_y*y); + eval_y[npts * 11] = y*y*z*(3*bf + bf_y*y); + eval_y[npts * 12] = y*z*z*(2*bf + bf_y*y); + eval_y[npts * 13] = z*z*z*(bf + bf_y*y); + eval_y[npts * 14] = bf_y*z*z*z*z; + + eval_z[npts * 0] = bf_z*x*x*x*x; + eval_z[npts * 1] = bf_z*x*x*x*y; + eval_z[npts * 2] = x*x*x*(bf + bf_z*z); + eval_z[npts * 3] = bf_z*x*x*y*y; + eval_z[npts * 4] = x*x*y*(bf + bf_z*z); + eval_z[npts * 5] = x*x*z*(2*bf + bf_z*z); + eval_z[npts * 6] = bf_z*x*y*y*y; + eval_z[npts * 7] = x*y*y*(bf + bf_z*z); + eval_z[npts * 8] = x*y*z*(2*bf + bf_z*z); + eval_z[npts * 9] = x*z*z*(3*bf + bf_z*z); + eval_z[npts * 10] = bf_z*y*y*y*y; + eval_z[npts * 11] = y*y*y*(bf + bf_z*z); + eval_z[npts * 12] = y*y*z*(2*bf + bf_z*z); + eval_z[npts * 13] = y*z*z*(3*bf + bf_z*z); + eval_z[npts * 14] = z*z*z*(4*bf + bf_z*z); + +} + + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular( + const int32_t npts, + const int32_t l, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + if( l == 0 ) { + + collocation_cartesian_angular_0( npts, bf, x, y, z, eval ); + + } else if( l == 1 ) { + + collocation_cartesian_angular_1( npts, bf, x, y, z, eval ); + + } else if( l == 2 ) { + + collocation_cartesian_angular_2( npts, bf, x, y, z, eval ); + + } else if( l == 3 ) { + + collocation_cartesian_angular_3( npts, bf, x, y, z, eval ); + + } else if( l == 4 ) { + + collocation_cartesian_angular_4( npts, bf, x, y, z, eval ); + + } else { + assert( false && "L < L_MAX" ); + } + +} // collocation_cartesian_angular + + +template +GPGAUEVAL_INLINE void collocation_cartesian_angular_deriv1( + const int32_t npts, + const int32_t l, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + + if( l == 0 ) { + + collocation_cartesian_angular_0( npts, bf, x, y, z, eval ); + collocation_cartesian_angular_0_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 1 ) { + + collocation_cartesian_angular_1( npts, bf, x, y, z, eval ); + collocation_cartesian_angular_1_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 2 ) { + + collocation_cartesian_angular_2( npts, bf, x, y, z, eval ); + collocation_cartesian_angular_2_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 3 ) { + + collocation_cartesian_angular_3( npts, bf, x, y, z, eval ); + collocation_cartesian_angular_3_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 4 ) { + + collocation_cartesian_angular_4( npts, bf, x, y, z, eval ); + collocation_cartesian_angular_4_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else { + assert( false && "L < L_MAX" ); + } + +} // collocation_cartesian_angular_deriv1 + + +} // namespace GauXC + diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_angular_spherical_unnorm.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_angular_spherical_unnorm.hpp new file mode 100644 index 000000000..75d7be4a7 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_angular_spherical_unnorm.hpp @@ -0,0 +1,378 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include +#include + +#ifndef GPGAUEVAL_INLINE +# define GPGAUEVAL_INLINE inline +#endif + +namespace GauXC { + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_0( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = bf; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_0_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = bf_x; + + eval_y[npts * 0] = bf_y; + + eval_z[npts * 0] = bf_z; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_1( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = bf*y; + eval[npts * 1] = bf*z; + eval[npts * 2] = bf*x; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_1_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = bf_x*y; + eval_x[npts * 1] = bf_x*z; + eval_x[npts * 2] = bf + bf_x*x; + + eval_y[npts * 0] = bf + bf_y*y; + eval_y[npts * 1] = bf_y*z; + eval_y[npts * 2] = bf_y*x; + + eval_z[npts * 0] = bf_z*y; + eval_z[npts * 1] = bf + bf_z*z; + eval_z[npts * 2] = bf_z*x; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_2( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = sqrt_3*bf*x*y; + eval[npts * 1] = sqrt_3*bf*y*z; + eval[npts * 2] = bf*(-x*x - y*y + 2*z*z)/2; + eval[npts * 3] = sqrt_3*bf*x*z; + eval[npts * 4] = sqrt_3*bf*(x*x - y*y)/2; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_2_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = sqrt_3*y*(bf + bf_x*x); + eval_x[npts * 1] = sqrt_3*bf_x*y*z; + eval_x[npts * 2] = -bf*x - bf_x*(x*x + y*y - 2*z*z)/2; + eval_x[npts * 3] = sqrt_3*z*(bf + bf_x*x); + eval_x[npts * 4] = sqrt_3*(bf*x + bf_x*(x*x - y*y)/2); + + eval_y[npts * 0] = sqrt_3*x*(bf + bf_y*y); + eval_y[npts * 1] = sqrt_3*z*(bf + bf_y*y); + eval_y[npts * 2] = -bf*y - bf_y*(x*x + y*y - 2*z*z)/2; + eval_y[npts * 3] = sqrt_3*bf_y*x*z; + eval_y[npts * 4] = sqrt_3*(-bf*y + bf_y*(x*x - y*y)/2); + + eval_z[npts * 0] = sqrt_3*bf_z*x*y; + eval_z[npts * 1] = sqrt_3*y*(bf + bf_z*z); + eval_z[npts * 2] = 2*bf*z - bf_z*(x*x + y*y - 2*z*z)/2; + eval_z[npts * 3] = sqrt_3*x*(bf + bf_z*z); + eval_z[npts * 4] = sqrt_3*bf_z*(x*x - y*y)/2; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_3( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = sqrt_10*bf*y*(3*x*x - y*y)/4; + eval[npts * 1] = sqrt_15*bf*x*y*z; + eval[npts * 2] = sqrt_6*bf*y*(-x*x - y*y + 4*z*z)/4; + eval[npts * 3] = bf*z*(-3*x*x - 3*y*y + 2*z*z)/2; + eval[npts * 4] = sqrt_6*bf*x*(-x*x - y*y + 4*z*z)/4; + eval[npts * 5] = sqrt_15*bf*z*(x*x - y*y)/2; + eval[npts * 6] = sqrt_10*bf*x*(x*x - 3*y*y)/4; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_3_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = sqrt_10*y*(6*bf*x + bf_x*(3*x*x - y*y))/4; + eval_x[npts * 1] = sqrt_15*y*z*(bf + bf_x*x); + eval_x[npts * 2] = sqrt_6*y*(-2*bf*x - bf_x*(x*x + y*y - 4*z*z))/4; + eval_x[npts * 3] = z*(-6*bf*x - bf_x*(3*x*x + 3*y*y - 2*z*z))/2; + eval_x[npts * 4] = sqrt_6*(-bf*(3*x*x + y*y - 4*z*z) - bf_x*x*(x*x + y*y - 4*z*z))/4; + eval_x[npts * 5] = sqrt_15*z*(2*bf*x + bf_x*(x*x - y*y))/2; + eval_x[npts * 6] = sqrt_10*(3*bf*(x*x - y*y) + bf_x*x*(x*x - 3*y*y))/4; + + eval_y[npts * 0] = sqrt_10*(-3*bf*(-x*x + y*y) + bf_y*y*(3*x*x - y*y))/4; + eval_y[npts * 1] = sqrt_15*x*z*(bf + bf_y*y); + eval_y[npts * 2] = sqrt_6*(-bf*(x*x + 3*y*y - 4*z*z) - bf_y*y*(x*x + y*y - 4*z*z))/4; + eval_y[npts * 3] = z*(-6*bf*y - bf_y*(3*x*x + 3*y*y - 2*z*z))/2; + eval_y[npts * 4] = sqrt_6*x*(-2*bf*y - bf_y*(x*x + y*y - 4*z*z))/4; + eval_y[npts * 5] = sqrt_15*z*(-2*bf*y + bf_y*(x*x - y*y))/2; + eval_y[npts * 6] = sqrt_10*x*(-6*bf*y + bf_y*(x*x - 3*y*y))/4; + + eval_z[npts * 0] = sqrt_10*bf_z*y*(3*x*x - y*y)/4; + eval_z[npts * 1] = sqrt_15*x*y*(bf + bf_z*z); + eval_z[npts * 2] = sqrt_6*y*(8*bf*z - bf_z*(x*x + y*y - 4*z*z))/4; + eval_z[npts * 3] = -3*bf*(x*x + y*y - 2*z*z)/2 - bf_z*z*(3*x*x + 3*y*y - 2*z*z)/2; + eval_z[npts * 4] = sqrt_6*x*(8*bf*z - bf_z*(x*x + y*y - 4*z*z))/4; + eval_z[npts * 5] = sqrt_15*(bf + bf_z*z)*(x*x - y*y)/2; + eval_z[npts * 6] = sqrt_10*bf_z*x*(x*x - 3*y*y)/4; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_4( + int32_t npts, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + eval[npts * 0] = sqrt_35*bf*x*y*(x*x - y*y)/2; + eval[npts * 1] = sqrt_70*bf*y*z*(3*x*x - y*y)/4; + eval[npts * 2] = sqrt_5*bf*x*y*(-x*x - y*y + 6*z*z)/2; + eval[npts * 3] = sqrt_10*bf*y*z*(-3*x*x - 3*y*y + 4*z*z)/4; + eval[npts * 4] = bf*(3*x*x*x*x + 6*x*x*y*y - 24*x*x*z*z + 3*y*y*y*y - 24*y*y*z*z + 8*z*z*z*z)/8; + eval[npts * 5] = sqrt_10*bf*x*z*(-3*x*x - 3*y*y + 4*z*z)/4; + eval[npts * 6] = sqrt_5*bf*(-x*x*x*x + 6*x*x*z*z + y*y*y*y - 6*y*y*z*z)/4; + eval[npts * 7] = sqrt_70*bf*x*z*(x*x - 3*y*y)/4; + eval[npts * 8] = sqrt_35*bf*(x*x*x*x - 6*x*x*y*y + y*y*y*y)/8; + +} + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_4_deriv1( + const int32_t npts, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + eval_x[npts * 0] = sqrt_35*y*(bf*(3*x*x - y*y) + bf_x*x*(x*x - y*y))/2; + eval_x[npts * 1] = sqrt_70*y*z*(6*bf*x + bf_x*(3*x*x - y*y))/4; + eval_x[npts * 2] = sqrt_5*y*(-bf*(3*x*x + y*y - 6*z*z) - bf_x*x*(x*x + y*y - 6*z*z))/2; + eval_x[npts * 3] = sqrt_10*y*z*(-6*bf*x - bf_x*(3*x*x + 3*y*y - 4*z*z))/4; + eval_x[npts * 4] = 3*bf*x*(x*x + y*y - 4*z*z)/2 + bf_x*(3*x*x*x*x + 6*x*x*y*y - 24*x*x*z*z + 3*y*y*y*y - 24*y*y*z*z + 8*z*z*z*z)/8; + eval_x[npts * 5] = sqrt_10*z*(-bf*(9*x*x + 3*y*y - 4*z*z) - bf_x*x*(3*x*x + 3*y*y - 4*z*z))/4; + eval_x[npts * 6] = sqrt_5*(-bf*x*(x*x - 3*z*z) - bf_x*(x*x*x*x - 6*x*x*z*z - y*y*y*y + 6*y*y*z*z)/4); + eval_x[npts * 7] = sqrt_70*z*(3*bf*(x*x - y*y) + bf_x*x*(x*x - 3*y*y))/4; + eval_x[npts * 8] = sqrt_35*(4*bf*x*(x*x - 3*y*y) + bf_x*(x*x*x*x - 6*x*x*y*y + y*y*y*y))/8; + + eval_y[npts * 0] = sqrt_35*x*(-bf*(-x*x + 3*y*y) + bf_y*y*(x*x - y*y))/2; + eval_y[npts * 1] = sqrt_70*z*(-3*bf*(-x*x + y*y) + bf_y*y*(3*x*x - y*y))/4; + eval_y[npts * 2] = sqrt_5*x*(-bf*(x*x + 3*y*y - 6*z*z) - bf_y*y*(x*x + y*y - 6*z*z))/2; + eval_y[npts * 3] = sqrt_10*z*(-bf*(3*x*x + 9*y*y - 4*z*z) - bf_y*y*(3*x*x + 3*y*y - 4*z*z))/4; + eval_y[npts * 4] = 3*bf*y*(x*x + y*y - 4*z*z)/2 + bf_y*(3*x*x*x*x + 6*x*x*y*y - 24*x*x*z*z + 3*y*y*y*y - 24*y*y*z*z + 8*z*z*z*z)/8; + eval_y[npts * 5] = sqrt_10*x*z*(-6*bf*y - bf_y*(3*x*x + 3*y*y - 4*z*z))/4; + eval_y[npts * 6] = sqrt_5*(bf*y*(y*y - 3*z*z) - bf_y*(x*x*x*x - 6*x*x*z*z - y*y*y*y + 6*y*y*z*z)/4); + eval_y[npts * 7] = sqrt_70*x*z*(-6*bf*y + bf_y*(x*x - 3*y*y))/4; + eval_y[npts * 8] = sqrt_35*(-4*bf*y*(3*x*x - y*y) + bf_y*(x*x*x*x - 6*x*x*y*y + y*y*y*y))/8; + + eval_z[npts * 0] = sqrt_35*bf_z*x*y*(x*x - y*y)/2; + eval_z[npts * 1] = sqrt_70*y*(bf + bf_z*z)*(3*x*x - y*y)/4; + eval_z[npts * 2] = sqrt_5*x*y*(12*bf*z - bf_z*(x*x + y*y - 6*z*z))/2; + eval_z[npts * 3] = sqrt_10*y*(3*bf*(-x*x - y*y + 4*z*z) - bf_z*z*(3*x*x + 3*y*y - 4*z*z))/4; + eval_z[npts * 4] = -2*bf*z*(3*x*x + 3*y*y - 2*z*z) + bf_z*(3*x*x*x*x + 6*x*x*y*y - 24*x*x*z*z + 3*y*y*y*y - 24*y*y*z*z + 8*z*z*z*z)/8; + eval_z[npts * 5] = sqrt_10*x*(3*bf*(-x*x - y*y + 4*z*z) - bf_z*z*(3*x*x + 3*y*y - 4*z*z))/4; + eval_z[npts * 6] = sqrt_5*(12*bf*z*(x*x - y*y) - bf_z*(x*x*x*x - 6*x*x*z*z - y*y*y*y + 6*y*y*z*z))/4; + eval_z[npts * 7] = sqrt_70*x*(bf + bf_z*z)*(x*x - 3*y*y)/4; + eval_z[npts * 8] = sqrt_35*bf_z*(x*x*x*x - 6*x*x*y*y + y*y*y*y)/8; + +} + + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular( + const int32_t npts, + const int32_t l, + const T bf, + const T x, + const T y, + const T z, + T* __restrict__ eval +) { + + if( l == 0 ) { + + collocation_spherical_unnorm_angular_0( npts, bf, x, y, z, eval ); + + } else if( l == 1 ) { + + collocation_spherical_unnorm_angular_1( npts, bf, x, y, z, eval ); + + } else if( l == 2 ) { + + collocation_spherical_unnorm_angular_2( npts, bf, x, y, z, eval ); + + } else if( l == 3 ) { + + collocation_spherical_unnorm_angular_3( npts, bf, x, y, z, eval ); + + } else if( l == 4 ) { + + collocation_spherical_unnorm_angular_4( npts, bf, x, y, z, eval ); + + } else { + assert( false && "L < L_MAX" ); + } + +} // collocation_spherical_unnorm_angular + + +template +GPGAUEVAL_INLINE void collocation_spherical_unnorm_angular_deriv1( + const int32_t npts, + const int32_t l, + const T bf, + const T bf_x, + const T bf_y, + const T bf_z, + const T x, + const T y, + const T z, + T* __restrict__ eval, + T* __restrict__ eval_x, + T* __restrict__ eval_y, + T* __restrict__ eval_z +) { + + + if( l == 0 ) { + + collocation_spherical_unnorm_angular_0( npts, bf, x, y, z, eval ); + collocation_spherical_unnorm_angular_0_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 1 ) { + + collocation_spherical_unnorm_angular_1( npts, bf, x, y, z, eval ); + collocation_spherical_unnorm_angular_1_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 2 ) { + + collocation_spherical_unnorm_angular_2( npts, bf, x, y, z, eval ); + collocation_spherical_unnorm_angular_2_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 3 ) { + + collocation_spherical_unnorm_angular_3( npts, bf, x, y, z, eval ); + collocation_spherical_unnorm_angular_3_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else if( l == 4 ) { + + collocation_spherical_unnorm_angular_4( npts, bf, x, y, z, eval ); + collocation_spherical_unnorm_angular_4_deriv1( npts, bf, bf_x, bf_y, bf_z, x, y, z, eval_x, eval_y, eval_z ); + + } else { + assert( false && "L < L_MAX" ); + } + +} // collocation_spherical_unnorm_angular_deriv1 + + +} // namespace GauXC + diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_device_constants.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_device_constants.hpp new file mode 100644 index 000000000..216a6326a --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_device_constants.hpp @@ -0,0 +1,24 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +namespace GauXC { + + constexpr double sqrt_3 = 1.7320508075688772; + constexpr double sqrt_5 = 2.23606797749979; + constexpr double sqrt_15 = 3.872983346207417; + constexpr double sqrt_10 = 3.1622776601683795; + constexpr double sqrt_6 = 2.449489742783178; + constexpr double sqrt_35 = 5.916079783099616; + constexpr double sqrt_70 = 8.366600265340756; + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_radial.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_radial.hpp new file mode 100644 index 000000000..3140678ba --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_radial.hpp @@ -0,0 +1,105 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include + +#include + + +namespace GauXC { + +inline void collocation_device_radial_eval( + const Shell& shell, + const double* pt, + double* x, + double* y, + double* z, + double* eval_device +) { + + const auto* O = shell.O_data(); + const auto* alpha = shell.alpha_data(); + const auto* coeff = shell.coeff_data(); + + const double xc = pt[0] - O[0]; + const double yc = pt[1] - O[1]; + const double zc = pt[2] - O[2]; + *x = xc; + *y = yc; + *z = zc; + + const double rsq = xc*xc + yc*yc + zc*zc; + + const uint32_t nprim = shell.nprim(); + double tmp = 0.; + for( uint32_t i = 0; i < nprim; ++i ) + tmp += coeff[i] * ::sycl::exp( - alpha[i] * rsq ); + + *eval_device = tmp; + +} + + + +inline void collocation_device_radial_eval_deriv1( + const Shell& shell, + const double* pt, + double* x, + double* y, + double* z, + double* eval_device, + double* deval_device_x, + double* deval_device_y, + double* deval_device_z +) { + + const auto* O = shell.O_data(); + const auto* alpha = shell.alpha_data(); + const auto* coeff = shell.coeff_data(); + + const double xc = pt[0] - O[0]; + const double yc = pt[1] - O[1]; + const double zc = pt[2] - O[2]; + *x = xc; + *y = yc; + *z = zc; + + const double rsq = xc*xc + yc*yc + zc*zc; + + const uint32_t nprim = shell.nprim(); + double tmp = 0.; + double tmp_x = 0., tmp_y = 0., tmp_z = 0.; + for( uint32_t i = 0; i < nprim; ++i ) { + + const double a = alpha[i]; + const double e = coeff[i] * ::sycl::exp( - a * rsq ); + + const double ae = 2. * a * e; + + tmp += e; + tmp_x -= ae * xc; + tmp_y -= ae * yc; + tmp_z -= ae * zc; + + } + + *eval_device = tmp; + *deval_device_x = tmp_x; + *deval_device_y = tmp_y; + *deval_device_z = tmp_z; + +} + +} // namespace GauXC + + diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0.hpp new file mode 100644 index 000000000..d1369fcbc --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0.hpp @@ -0,0 +1,140 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_0_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_0_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_cartesian_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_gradient.hpp new file mode 100644 index 000000000..08b2f5137 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_gradient.hpp @@ -0,0 +1,163 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_0_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_0_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_cartesian_gradient_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_hessian.hpp new file mode 100644 index 000000000..7fc9a3513 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_hessian.hpp @@ -0,0 +1,193 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_0_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_0_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_cartesian_hessian_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = radial_eval_alpha_squared*(x*x); + const auto x1 = radial_eval_alpha_squared*x; + const auto x2 = radial_eval_alpha_squared*(y*y); + const auto x3 = radial_eval_alpha_squared*(z*z); + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = radial_eval_alpha + x0; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x1*y; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x1*z; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = radial_eval_alpha + x2; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*y*z; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = radial_eval_alpha + x3; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_lapgrad.hpp new file mode 100644 index 000000000..463b08bdf --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_lapgrad.hpp @@ -0,0 +1,215 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_0_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_0_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_cartesian_lapgrad_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = radial_eval_alpha_squared*x0; + const auto x2 = radial_eval_alpha_squared*x; + const auto x3 = y*y; + const auto x4 = radial_eval_alpha_squared*x3; + const auto x5 = radial_eval_alpha_squared*y; + const auto x6 = z*z; + const auto x7 = radial_eval_alpha_squared*x6; + const auto x8 = radial_eval_alpha_cubed*x; + const auto x9 = radial_eval_alpha_cubed*y; + const auto x10 = radial_eval_alpha_cubed*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = radial_eval_alpha + x1; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x2*y; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x2*z; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = radial_eval_alpha + x4; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x5*z; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = radial_eval_alpha + x7; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = 3.0*radial_eval_alpha + x1 + x4 + x7; + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = radial_eval_alpha_cubed*(x*x*x) + 5.0*x2 + x3*x8 + x6*x8; + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = radial_eval_alpha_cubed*(y*y*y) + x0*x9 + 5.0*x5 + x6*x9; + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = radial_eval_alpha_cubed*(z*z*z) + 5.0*radial_eval_alpha_squared*z + x0*x10 + x10*x3; + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_laplacian.hpp new file mode 100644 index 000000000..10a93a871 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l0_laplacian.hpp @@ -0,0 +1,173 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_0_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_0_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_cartesian_laplacian_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = radial_eval_alpha_squared*(x*x); + const auto x1 = radial_eval_alpha_squared*x; + const auto x2 = radial_eval_alpha_squared*(y*y); + const auto x3 = radial_eval_alpha_squared*(z*z); + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = 3.0*radial_eval_alpha + x0 + x2 + x3; + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1.hpp new file mode 100644 index 000000000..cd39a16f5 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1.hpp @@ -0,0 +1,148 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_1_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_1_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_cartesian_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x; + basis_eval[ipt + 1*npts] = radial_eval*y; + basis_eval[ipt + 2*npts] = radial_eval*z; + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*x; + ang_eval_1 = radial_eval*y; + ang_eval_2 = radial_eval*z; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_gradient.hpp new file mode 100644 index 000000000..1d88d3075 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_gradient.hpp @@ -0,0 +1,195 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_1_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_1_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_cartesian_gradient_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = radial_eval_alpha*x; + const auto x1 = x0*y; + const auto x2 = x0*z; + const auto x3 = radial_eval_alpha*y*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x; + basis_eval[ipt + 1*npts] = radial_eval*y; + basis_eval[ipt + 2*npts] = radial_eval*z; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*(x*x); + basis_x_eval[ipt + 1*npts] = x1; + basis_x_eval[ipt + 2*npts] = x2; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x1; + basis_y_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*(y*y); + basis_y_eval[ipt + 2*npts] = x3; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x2; + basis_z_eval[ipt + 1*npts] = x3; + basis_z_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*(z*z); + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*x; + ang_eval_1 = radial_eval*y; + ang_eval_2 = radial_eval*z; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = radial_eval + radial_eval_alpha*(x*x); + dang_eval_y_0 = x1; + dang_eval_z_0 = x2; + dang_eval_x_1 = x1; + dang_eval_y_1 = radial_eval + radial_eval_alpha*(y*y); + dang_eval_z_1 = x3; + dang_eval_x_2 = x2; + dang_eval_y_2 = x3; + dang_eval_z_2 = radial_eval + radial_eval_alpha*(z*z); + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_hessian.hpp new file mode 100644 index 000000000..8ad19ffbf --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_hessian.hpp @@ -0,0 +1,252 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_1_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_1_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_hessian_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = radial_eval_alpha*x; + const auto x2 = x1*y; + const auto x3 = x1*z; + const auto x4 = y*y; + const auto x5 = y*z; + const auto x6 = radial_eval_alpha*x5; + const auto x7 = z*z; + const auto x8 = 3.0*radial_eval_alpha; + const auto x9 = radial_eval_alpha_squared*x0; + const auto x10 = radial_eval_alpha + x9; + const auto x11 = x10*y; + const auto x12 = x10*z; + const auto x13 = radial_eval_alpha_squared*x4; + const auto x14 = radial_eval_alpha + x13; + const auto x15 = x*x14; + const auto x16 = radial_eval_alpha_squared*x*x5; + const auto x17 = radial_eval_alpha_squared*x7; + const auto x18 = radial_eval_alpha + x17; + const auto x19 = x*x18; + const auto x20 = x14*z; + const auto x21 = x18*y; + const auto x22 = 5.0*radial_eval_alpha + x13 + x17 + x9; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x; + basis_eval[ipt + 1*npts] = radial_eval*y; + basis_eval[ipt + 2*npts] = radial_eval*z; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*x0; + basis_x_eval[ipt + 1*npts] = x2; + basis_x_eval[ipt + 2*npts] = x3; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x2; + basis_y_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*x4; + basis_y_eval[ipt + 2*npts] = x6; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x3; + basis_z_eval[ipt + 1*npts] = x6; + basis_z_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*x7; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x*(x8 + x9); + basis_xx_eval[ipt + 1*npts] = x11; + basis_xx_eval[ipt + 2*npts] = x12; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x11; + basis_xy_eval[ipt + 1*npts] = x15; + basis_xy_eval[ipt + 2*npts] = x16; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x12; + basis_xz_eval[ipt + 1*npts] = x16; + basis_xz_eval[ipt + 2*npts] = x19; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x15; + basis_yy_eval[ipt + 1*npts] = y*(x13 + x8); + basis_yy_eval[ipt + 2*npts] = x20; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x16; + basis_yz_eval[ipt + 1*npts] = x20; + basis_yz_eval[ipt + 2*npts] = x21; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x19; + basis_zz_eval[ipt + 1*npts] = x21; + basis_zz_eval[ipt + 2*npts] = z*(x17 + x8); + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*x; + ang_eval_1 = radial_eval*y; + ang_eval_2 = radial_eval*z; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = radial_eval + radial_eval_alpha*x0; + dang_eval_y_0 = x2; + dang_eval_z_0 = x3; + dang_eval_x_1 = x2; + dang_eval_y_1 = radial_eval + radial_eval_alpha*x4; + dang_eval_z_1 = x6; + dang_eval_x_2 = x3; + dang_eval_y_2 = x6; + dang_eval_z_2 = radial_eval + radial_eval_alpha*x7; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_lapgrad.hpp new file mode 100644 index 000000000..82edba796 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_lapgrad.hpp @@ -0,0 +1,292 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_1_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_1_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_lapgrad_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = x0; + const auto x2 = radial_eval_alpha*x; + const auto x3 = x2*y; + const auto x4 = x2*z; + const auto x5 = y*y; + const auto x6 = x5; + const auto x7 = y*z; + const auto x8 = radial_eval_alpha*x7; + const auto x9 = z*z; + const auto x10 = x9; + const auto x11 = 3.0*radial_eval_alpha; + const auto x12 = radial_eval_alpha_squared*x1; + const auto x13 = radial_eval_alpha + x12; + const auto x14 = x13*y; + const auto x15 = x13*z; + const auto x16 = radial_eval_alpha_squared*x6; + const auto x17 = radial_eval_alpha + x16; + const auto x18 = x*x17; + const auto x19 = radial_eval_alpha_squared*x*x7; + const auto x20 = radial_eval_alpha_squared*x10; + const auto x21 = radial_eval_alpha + x20; + const auto x22 = x*x21; + const auto x23 = x17*z; + const auto x24 = x21*y; + const auto x25 = 5.0*radial_eval_alpha; + const auto x26 = x16 + x20 + x25; + const auto x27 = x12 + x26; + const auto x28 = 3.0*radial_eval_alpha_squared; + const auto x29 = radial_eval_alpha_cubed*(x*x*x); + const auto x30 = radial_eval_alpha_cubed*x6 + radial_eval_alpha_squared; + const auto x31 = radial_eval_alpha_cubed*x10 + radial_eval_alpha_squared; + const auto x32 = 5.0*radial_eval_alpha_squared; + const auto x33 = x*x30 + x*x31 + x*x32 + x29; + const auto x34 = radial_eval_alpha_cubed*(y*y*y); + const auto x35 = radial_eval_alpha_cubed*x1 + radial_eval_alpha_squared; + const auto x36 = x31*y + x32*y + x34 + x35*y; + const auto x37 = x12 + x25; + const auto x38 = radial_eval_alpha_cubed*(z*z*z); + const auto x39 = x30*z + x32*z + x35*z + x38; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x; + basis_eval[ipt + 1*npts] = radial_eval*y; + basis_eval[ipt + 2*npts] = radial_eval*z; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*x1; + basis_x_eval[ipt + 1*npts] = x3; + basis_x_eval[ipt + 2*npts] = x4; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x3; + basis_y_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*x6; + basis_y_eval[ipt + 2*npts] = x8; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x4; + basis_z_eval[ipt + 1*npts] = x8; + basis_z_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*x10; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x*(x11 + x12); + basis_xx_eval[ipt + 1*npts] = x14; + basis_xx_eval[ipt + 2*npts] = x15; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x14; + basis_xy_eval[ipt + 1*npts] = x18; + basis_xy_eval[ipt + 2*npts] = x19; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x15; + basis_xz_eval[ipt + 1*npts] = x19; + basis_xz_eval[ipt + 2*npts] = x22; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x18; + basis_yy_eval[ipt + 1*npts] = y*(x11 + x16); + basis_yy_eval[ipt + 2*npts] = x23; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x19; + basis_yz_eval[ipt + 1*npts] = x23; + basis_yz_eval[ipt + 2*npts] = x24; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x22; + basis_zz_eval[ipt + 1*npts] = x24; + basis_zz_eval[ipt + 2*npts] = z*(x11 + x20); + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x*x27; + basis_lapl_eval[ipt + 1*npts] = x27*y; + basis_lapl_eval[ipt + 2*npts] = x27*z; + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = x*(x*x28 + x29) + x0*x30 + x0*x31 + x1*x28 + x26; + basis_lapl_x_eval[ipt + 1*npts] = x33*y; + basis_lapl_x_eval[ipt + 2*npts] = x33*z; + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x*x36; + basis_lapl_y_eval[ipt + 1*npts] = x20 + x28*x6 + x31*x5 + x35*x5 + x37 + y*(x28*y + x34); + basis_lapl_y_eval[ipt + 2*npts] = x36*z; + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x*x39; + basis_lapl_z_eval[ipt + 1*npts] = x39*y; + basis_lapl_z_eval[ipt + 2*npts] = x10*x28 + x16 + x30*x9 + x35*x9 + x37 + z*(x28*z + x38); + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*x; + ang_eval_1 = radial_eval*y; + ang_eval_2 = radial_eval*z; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = radial_eval + radial_eval_alpha*x1; + dang_eval_y_0 = x3; + dang_eval_z_0 = x4; + dang_eval_x_1 = x3; + dang_eval_y_1 = radial_eval + radial_eval_alpha*x6; + dang_eval_z_1 = x8; + dang_eval_x_2 = x4; + dang_eval_y_2 = x8; + dang_eval_z_2 = radial_eval + radial_eval_alpha*x10; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_laplacian.hpp new file mode 100644 index 000000000..b3b650693 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l1_laplacian.hpp @@ -0,0 +1,222 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_1_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_1_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_laplacian_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = radial_eval_alpha*x; + const auto x2 = x1*y; + const auto x3 = x1*z; + const auto x4 = y*y; + const auto x5 = y*z; + const auto x6 = radial_eval_alpha*x5; + const auto x7 = z*z; + const auto x8 = 3.0*radial_eval_alpha; + const auto x9 = radial_eval_alpha_squared*x0; + const auto x10 = radial_eval_alpha + x9; + const auto x11 = x10*y; + const auto x12 = x10*z; + const auto x13 = radial_eval_alpha_squared*x4; + const auto x14 = radial_eval_alpha + x13; + const auto x15 = x*x14; + const auto x16 = radial_eval_alpha_squared*x*x5; + const auto x17 = radial_eval_alpha_squared*x7; + const auto x18 = radial_eval_alpha + x17; + const auto x19 = x*x18; + const auto x20 = x14*z; + const auto x21 = x18*y; + const auto x22 = 5.0*radial_eval_alpha + x13 + x17 + x9; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x; + basis_eval[ipt + 1*npts] = radial_eval*y; + basis_eval[ipt + 2*npts] = radial_eval*z; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*x0; + basis_x_eval[ipt + 1*npts] = x2; + basis_x_eval[ipt + 2*npts] = x3; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x2; + basis_y_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*x4; + basis_y_eval[ipt + 2*npts] = x6; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x3; + basis_z_eval[ipt + 1*npts] = x6; + basis_z_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*x7; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x*x22; + basis_lapl_eval[ipt + 1*npts] = x22*y; + basis_lapl_eval[ipt + 2*npts] = x22*z; + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*x; + ang_eval_1 = radial_eval*y; + ang_eval_2 = radial_eval*z; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = radial_eval + radial_eval_alpha*x0; + dang_eval_y_0 = x2; + dang_eval_z_0 = x3; + dang_eval_x_1 = x2; + dang_eval_y_1 = radial_eval + radial_eval_alpha*x4; + dang_eval_z_1 = x6; + dang_eval_x_2 = x3; + dang_eval_y_2 = x6; + dang_eval_z_2 = radial_eval + radial_eval_alpha*x7; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2.hpp new file mode 100644 index 000000000..69697708d --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2.hpp @@ -0,0 +1,160 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_2_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_2_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_cartesian_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + const auto x0 = radial_eval*x; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*(x*x); + basis_eval[ipt + 1*npts] = x0*y; + basis_eval[ipt + 2*npts] = x0*z; + basis_eval[ipt + 3*npts] = radial_eval*(y*y); + basis_eval[ipt + 4*npts] = radial_eval*y*z; + basis_eval[ipt + 5*npts] = radial_eval*(z*z); + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*(x*x); + ang_eval_1 = x0*y; + ang_eval_2 = x0*z; + ang_eval_3 = radial_eval*(y*y); + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*y*z; + ang_eval_1 = radial_eval*(z*z); + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_gradient.hpp new file mode 100644 index 000000000..3c8f94691 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_gradient.hpp @@ -0,0 +1,243 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_2_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_2_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_cartesian_gradient_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = radial_eval*x; + const auto x2 = y*y; + const auto x3 = radial_eval*y; + const auto x4 = z*z; + const auto x5 = radial_eval + radial_eval_alpha*x0; + const auto x6 = radial_eval_alpha*x; + const auto x7 = x6*y*z; + const auto x8 = radial_eval_alpha*y; + const auto x9 = radial_eval + radial_eval_alpha*x2; + const auto x10 = radial_eval_alpha*z; + const auto x11 = radial_eval + radial_eval_alpha*x4; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*y; + basis_eval[ipt + 2*npts] = x1*z; + basis_eval[ipt + 3*npts] = radial_eval*x2; + basis_eval[ipt + 4*npts] = x3*z; + basis_eval[ipt + 5*npts] = radial_eval*x4; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*(x*x*x) + 2.0*x1; + basis_x_eval[ipt + 1*npts] = x5*y; + basis_x_eval[ipt + 2*npts] = x5*z; + basis_x_eval[ipt + 3*npts] = x2*x6; + basis_x_eval[ipt + 4*npts] = x7; + basis_x_eval[ipt + 5*npts] = x4*x6; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x8; + basis_y_eval[ipt + 1*npts] = x*x9; + basis_y_eval[ipt + 2*npts] = x7; + basis_y_eval[ipt + 3*npts] = radial_eval_alpha*(y*y*y) + 2.0*x3; + basis_y_eval[ipt + 4*npts] = x9*z; + basis_y_eval[ipt + 5*npts] = x4*x8; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x10; + basis_z_eval[ipt + 1*npts] = x7; + basis_z_eval[ipt + 2*npts] = x*x11; + basis_z_eval[ipt + 3*npts] = x10*x2; + basis_z_eval[ipt + 4*npts] = x11*y; + basis_z_eval[ipt + 5*npts] = z*(2.0*radial_eval + radial_eval_alpha*(z*z)); + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*y; + ang_eval_2 = x1*z; + ang_eval_3 = radial_eval*x2; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x3*z; + ang_eval_1 = radial_eval*x4; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*(x*x*x) + 2.0*x1; + dang_eval_y_0 = x0*x8; + dang_eval_z_0 = x0*x10; + dang_eval_x_1 = x5*y; + dang_eval_y_1 = x*x9; + dang_eval_z_1 = x7; + dang_eval_x_2 = x5*z; + dang_eval_y_2 = x7; + dang_eval_z_2 = x*x11; + dang_eval_x_3 = x2*x6; + dang_eval_y_3 = radial_eval_alpha*(y*y*y) + 2.0*x3; + dang_eval_z_3 = x10*x2; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x7; + dang_eval_y_0 = x9*z; + dang_eval_z_0 = x11*y; + dang_eval_x_1 = x4*x6; + dang_eval_y_1 = x4*x8; + dang_eval_z_1 = z*(2.0*radial_eval + radial_eval_alpha*(z*z)); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_hessian.hpp new file mode 100644 index 000000000..123906d2a --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_hessian.hpp @@ -0,0 +1,333 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_2_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_2_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_hessian_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = x*y; + const auto x2 = x*z; + const auto x3 = y*y; + const auto x4 = y*z; + const auto x5 = z*z; + const auto x6 = 2.0*radial_eval; + const auto x7 = x*x*x; + const auto x8 = radial_eval + radial_eval_alpha*x0; + const auto x9 = radial_eval_alpha*x; + const auto x10 = x4*x9; + const auto x11 = radial_eval_alpha*y; + const auto x12 = radial_eval_alpha*x3; + const auto x13 = radial_eval + x12; + const auto x14 = y*y*y; + const auto x15 = radial_eval_alpha*z; + const auto x16 = radial_eval_alpha*x5; + const auto x17 = radial_eval + x16; + const auto x18 = z*z*z; + const auto x19 = 4.0*radial_eval_alpha; + const auto x20 = radial_eval_alpha_squared*x0; + const auto x21 = radial_eval_alpha + x20; + const auto x22 = x0*x19 + x0*x21 + x6; + const auto x23 = 3.0*radial_eval_alpha; + const auto x24 = x20 + x23; + const auto x25 = x21*x3; + const auto x26 = x21*x4; + const auto x27 = x21*x5; + const auto x28 = radial_eval_alpha_squared*x7 + 2.0*x9; + const auto x29 = radial_eval_alpha_squared*x14 + 2.0*x11; + const auto x30 = radial_eval_alpha_squared*x3; + const auto x31 = radial_eval_alpha + x30; + const auto x32 = x2*x31; + const auto x33 = radial_eval_alpha_squared*x5; + const auto x34 = radial_eval_alpha + x33; + const auto x35 = x1*x34; + const auto x36 = radial_eval_alpha_squared*x18 + 2.0*x15; + const auto x37 = x0*x31; + const auto x38 = x23 + x30; + const auto x39 = x19*x3 + x3*x31 + x6; + const auto x40 = x31*x5; + const auto x41 = x0*x34; + const auto x42 = x23 + x33; + const auto x43 = x3*x34; + const auto x44 = x19*x5 + x34*x5 + x6; + const auto x45 = 7.0*radial_eval_alpha + x20 + x30 + x33; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = radial_eval*x1; + basis_eval[ipt + 2*npts] = radial_eval*x2; + basis_eval[ipt + 3*npts] = radial_eval*x3; + basis_eval[ipt + 4*npts] = radial_eval*x4; + basis_eval[ipt + 5*npts] = radial_eval*x5; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x7 + x*x6; + basis_x_eval[ipt + 1*npts] = x8*y; + basis_x_eval[ipt + 2*npts] = x8*z; + basis_x_eval[ipt + 3*npts] = x3*x9; + basis_x_eval[ipt + 4*npts] = x10; + basis_x_eval[ipt + 5*npts] = x5*x9; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x11; + basis_y_eval[ipt + 1*npts] = x*x13; + basis_y_eval[ipt + 2*npts] = x10; + basis_y_eval[ipt + 3*npts] = radial_eval_alpha*x14 + x6*y; + basis_y_eval[ipt + 4*npts] = x13*z; + basis_y_eval[ipt + 5*npts] = x11*x5; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x15; + basis_z_eval[ipt + 1*npts] = x10; + basis_z_eval[ipt + 2*npts] = x*x17; + basis_z_eval[ipt + 3*npts] = x15*x3; + basis_z_eval[ipt + 4*npts] = x17*y; + basis_z_eval[ipt + 5*npts] = radial_eval_alpha*x18 + x6*z; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x22; + basis_xx_eval[ipt + 1*npts] = x1*x24; + basis_xx_eval[ipt + 2*npts] = x2*x24; + basis_xx_eval[ipt + 3*npts] = x25; + basis_xx_eval[ipt + 4*npts] = x26; + basis_xx_eval[ipt + 5*npts] = x27; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x28*y; + basis_xy_eval[ipt + 1*npts] = radial_eval_alpha_squared*x0*x3 + x12 + x8; + basis_xy_eval[ipt + 2*npts] = x26; + basis_xy_eval[ipt + 3*npts] = x*x29; + basis_xy_eval[ipt + 4*npts] = x32; + basis_xy_eval[ipt + 5*npts] = radial_eval_alpha_squared*x1*x5; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x28*z; + basis_xz_eval[ipt + 1*npts] = x26; + basis_xz_eval[ipt + 2*npts] = radial_eval_alpha_squared*x0*x5 + x16 + x8; + basis_xz_eval[ipt + 3*npts] = radial_eval_alpha_squared*x2*x3; + basis_xz_eval[ipt + 4*npts] = x35; + basis_xz_eval[ipt + 5*npts] = x*x36; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x37; + basis_yy_eval[ipt + 1*npts] = x1*x38; + basis_yy_eval[ipt + 2*npts] = x32; + basis_yy_eval[ipt + 3*npts] = x39; + basis_yy_eval[ipt + 4*npts] = x38*x4; + basis_yy_eval[ipt + 5*npts] = x40; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*x0*x4; + basis_yz_eval[ipt + 1*npts] = x32; + basis_yz_eval[ipt + 2*npts] = x35; + basis_yz_eval[ipt + 3*npts] = x29*z; + basis_yz_eval[ipt + 4*npts] = radial_eval_alpha_squared*x3*x5 + x13 + x16; + basis_yz_eval[ipt + 5*npts] = x36*y; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x41; + basis_zz_eval[ipt + 1*npts] = x35; + basis_zz_eval[ipt + 2*npts] = x2*x42; + basis_zz_eval[ipt + 3*npts] = x43; + basis_zz_eval[ipt + 4*npts] = x4*x42; + basis_zz_eval[ipt + 5*npts] = x44; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = radial_eval*x1; + ang_eval_2 = radial_eval*x2; + ang_eval_3 = radial_eval*x3; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x4; + ang_eval_1 = radial_eval*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x7 + x*x6; + dang_eval_y_0 = x0*x11; + dang_eval_z_0 = x0*x15; + dang_eval_x_1 = x8*y; + dang_eval_y_1 = x*x13; + dang_eval_z_1 = x10; + dang_eval_x_2 = x8*z; + dang_eval_y_2 = x10; + dang_eval_z_2 = x*x17; + dang_eval_x_3 = x3*x9; + dang_eval_y_3 = radial_eval_alpha*x14 + x6*y; + dang_eval_z_3 = x15*x3; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x10; + dang_eval_y_0 = x13*z; + dang_eval_z_0 = x17*y; + dang_eval_x_1 = x5*x9; + dang_eval_y_1 = x11*x5; + dang_eval_z_1 = radial_eval_alpha*x18 + x6*z; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_lapgrad.hpp new file mode 100644 index 000000000..05bb8eab8 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_lapgrad.hpp @@ -0,0 +1,407 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_2_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_2_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_lapgrad_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = x0; + const auto x2 = x*y; + const auto x3 = x*z; + const auto x4 = y*y; + const auto x5 = x4; + const auto x6 = y*z; + const auto x7 = z*z; + const auto x8 = x7; + const auto x9 = 2.0*radial_eval; + const auto x10 = x*x*x; + const auto x11 = radial_eval + radial_eval_alpha*x1; + const auto x12 = radial_eval_alpha*x; + const auto x13 = x12*x6; + const auto x14 = radial_eval_alpha*y; + const auto x15 = radial_eval_alpha*x5; + const auto x16 = radial_eval + x15; + const auto x17 = y*y*y; + const auto x18 = radial_eval_alpha*z; + const auto x19 = radial_eval_alpha*x8; + const auto x20 = radial_eval + x19; + const auto x21 = z*z*z; + const auto x22 = 4.0*radial_eval_alpha; + const auto x23 = radial_eval_alpha_squared*x1; + const auto x24 = radial_eval_alpha + x23; + const auto x25 = x1*x22 + x1*x24 + x9; + const auto x26 = 3.0*radial_eval_alpha; + const auto x27 = x23 + x26; + const auto x28 = x24*x5; + const auto x29 = x24*x6; + const auto x30 = x24*x8; + const auto x31 = 2.0*x12; + const auto x32 = radial_eval_alpha_squared*x10 + x31; + const auto x33 = 2.0*x14; + const auto x34 = radial_eval_alpha_squared*x17 + x33; + const auto x35 = radial_eval_alpha_squared*x5; + const auto x36 = radial_eval_alpha + x35; + const auto x37 = x3*x36; + const auto x38 = radial_eval_alpha_squared*x8; + const auto x39 = radial_eval_alpha + x38; + const auto x40 = x2*x39; + const auto x41 = 2.0*x18; + const auto x42 = radial_eval_alpha_squared*x21 + x41; + const auto x43 = x1*x36; + const auto x44 = x26 + x35; + const auto x45 = x22*x5 + x36*x5 + x9; + const auto x46 = x36*x8; + const auto x47 = x1*x39; + const auto x48 = x26 + x38; + const auto x49 = x39*x5; + const auto x50 = x22*x8 + x39*x8 + x9; + const auto x51 = x35 + x38; + const auto x52 = 7.0*radial_eval_alpha + x23 + x51; + const auto x53 = 2.0*x; + const auto x54 = radial_eval_alpha_cubed*x5 + radial_eval_alpha_squared; + const auto x55 = x1*x54; + const auto x56 = radial_eval_alpha_cubed*x8 + radial_eval_alpha_squared; + const auto x57 = x1*x56; + const auto x58 = radial_eval_alpha_squared*x; + const auto x59 = radial_eval_alpha_cubed*x10; + const auto x60 = 3.0*x58 + x59; + const auto x61 = 2.0*radial_eval_alpha_squared; + const auto x62 = x*x60 + x0*x54 + x0*x56 + x1*x61 + x22 + 3.0*x24 + x51; + const auto x63 = 4.0*x58; + const auto x64 = x5*x54; + const auto x65 = x5*x56; + const auto x66 = x54*x8; + const auto x67 = x56*x8; + const auto x68 = radial_eval_alpha_squared*y; + const auto x69 = 4.0*x68; + const auto x70 = radial_eval_alpha_cubed*x1 + radial_eval_alpha_squared; + const auto x71 = x1*x70; + const auto x72 = radial_eval_alpha_cubed*x17; + const auto x73 = 3.0*x68 + x72; + const auto x74 = x22 + x23; + const auto x75 = 3.0*x36 + x38 + x4*x56 + x4*x70 + x5*x61 + x73*y + x74; + const auto x76 = 2.0*y; + const auto x77 = x5*x70; + const auto x78 = x70*x8; + const auto x79 = radial_eval_alpha_squared*z; + const auto x80 = 4.0*x79; + const auto x81 = radial_eval_alpha_cubed*x21; + const auto x82 = 3.0*x79 + x81; + const auto x83 = x35 + 3.0*x39 + x54*x7 + x61*x8 + x7*x70 + x74 + x82*z; + const auto x84 = 2.0*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x1; + basis_eval[ipt + 1*npts] = radial_eval*x2; + basis_eval[ipt + 2*npts] = radial_eval*x3; + basis_eval[ipt + 3*npts] = radial_eval*x5; + basis_eval[ipt + 4*npts] = radial_eval*x6; + basis_eval[ipt + 5*npts] = radial_eval*x8; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x10 + x*x9; + basis_x_eval[ipt + 1*npts] = x11*y; + basis_x_eval[ipt + 2*npts] = x11*z; + basis_x_eval[ipt + 3*npts] = x12*x5; + basis_x_eval[ipt + 4*npts] = x13; + basis_x_eval[ipt + 5*npts] = x12*x8; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x1*x14; + basis_y_eval[ipt + 1*npts] = x*x16; + basis_y_eval[ipt + 2*npts] = x13; + basis_y_eval[ipt + 3*npts] = radial_eval_alpha*x17 + x9*y; + basis_y_eval[ipt + 4*npts] = x16*z; + basis_y_eval[ipt + 5*npts] = x14*x8; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x1*x18; + basis_z_eval[ipt + 1*npts] = x13; + basis_z_eval[ipt + 2*npts] = x*x20; + basis_z_eval[ipt + 3*npts] = x18*x5; + basis_z_eval[ipt + 4*npts] = x20*y; + basis_z_eval[ipt + 5*npts] = radial_eval_alpha*x21 + x9*z; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x25; + basis_xx_eval[ipt + 1*npts] = x2*x27; + basis_xx_eval[ipt + 2*npts] = x27*x3; + basis_xx_eval[ipt + 3*npts] = x28; + basis_xx_eval[ipt + 4*npts] = x29; + basis_xx_eval[ipt + 5*npts] = x30; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x32*y; + basis_xy_eval[ipt + 1*npts] = radial_eval_alpha_squared*x1*x5 + x11 + x15; + basis_xy_eval[ipt + 2*npts] = x29; + basis_xy_eval[ipt + 3*npts] = x*x34; + basis_xy_eval[ipt + 4*npts] = x37; + basis_xy_eval[ipt + 5*npts] = radial_eval_alpha_squared*x2*x8; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x32*z; + basis_xz_eval[ipt + 1*npts] = x29; + basis_xz_eval[ipt + 2*npts] = radial_eval_alpha_squared*x1*x8 + x11 + x19; + basis_xz_eval[ipt + 3*npts] = radial_eval_alpha_squared*x3*x5; + basis_xz_eval[ipt + 4*npts] = x40; + basis_xz_eval[ipt + 5*npts] = x*x42; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x43; + basis_yy_eval[ipt + 1*npts] = x2*x44; + basis_yy_eval[ipt + 2*npts] = x37; + basis_yy_eval[ipt + 3*npts] = x45; + basis_yy_eval[ipt + 4*npts] = x44*x6; + basis_yy_eval[ipt + 5*npts] = x46; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*x1*x6; + basis_yz_eval[ipt + 1*npts] = x37; + basis_yz_eval[ipt + 2*npts] = x40; + basis_yz_eval[ipt + 3*npts] = x34*z; + basis_yz_eval[ipt + 4*npts] = radial_eval_alpha_squared*x5*x8 + x16 + x19; + basis_yz_eval[ipt + 5*npts] = x42*y; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x47; + basis_zz_eval[ipt + 1*npts] = x40; + basis_zz_eval[ipt + 2*npts] = x3*x48; + basis_zz_eval[ipt + 3*npts] = x49; + basis_zz_eval[ipt + 4*npts] = x48*x6; + basis_zz_eval[ipt + 5*npts] = x50; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x25 + x43 + x47; + basis_lapl_eval[ipt + 1*npts] = x2*x52; + basis_lapl_eval[ipt + 2*npts] = x3*x52; + basis_lapl_eval[ipt + 3*npts] = x28 + x45 + x49; + basis_lapl_eval[ipt + 4*npts] = x52*x6; + basis_lapl_eval[ipt + 5*npts] = x30 + x46 + x50; + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = 6.0*x*x24 + x*x55 + x*x57 + x1*x60 + 6.0*x12 + x36*x53 + x39*x53; + basis_lapl_x_eval[ipt + 1*npts] = x62*y; + basis_lapl_x_eval[ipt + 2*npts] = x62*z; + basis_lapl_x_eval[ipt + 3*npts] = x*x64 + x*x65 + x31 + x5*x60 + x5*x63; + basis_lapl_x_eval[ipt + 4*npts] = x6*(x*x54 + x*x56 + 7.0*x58 + x59); + basis_lapl_x_eval[ipt + 5*npts] = x*x66 + x*x67 + x31 + x60*x8 + x63*x8; + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x1*x69 + x1*x73 + x33 + x57*y + x71*y; + basis_lapl_y_eval[ipt + 1*npts] = x*x75; + basis_lapl_y_eval[ipt + 2*npts] = x3*(x56*y + 7.0*x68 + x70*y + x72); + basis_lapl_y_eval[ipt + 3*npts] = 6.0*x14 + x24*x76 + 6.0*x36*y + x39*x76 + x5*x73 + x65*y + x77*y; + basis_lapl_y_eval[ipt + 4*npts] = x75*z; + basis_lapl_y_eval[ipt + 5*npts] = x33 + x67*y + x69*x8 + x73*x8 + x78*y; + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x1*x80 + x1*x82 + x41 + x55*z + x71*z; + basis_lapl_z_eval[ipt + 1*npts] = x2*(x54*z + x70*z + 7.0*x79 + x81); + basis_lapl_z_eval[ipt + 2*npts] = x*x83; + basis_lapl_z_eval[ipt + 3*npts] = x41 + x5*x80 + x5*x82 + x64*z + x77*z; + basis_lapl_z_eval[ipt + 4*npts] = x83*y; + basis_lapl_z_eval[ipt + 5*npts] = 6.0*x18 + x24*x84 + x36*x84 + 6.0*x39*z + x66*z + x78*z + x8*x82; + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x1; + ang_eval_1 = radial_eval*x2; + ang_eval_2 = radial_eval*x3; + ang_eval_3 = radial_eval*x5; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x6; + ang_eval_1 = radial_eval*x8; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x10 + x*x9; + dang_eval_y_0 = x1*x14; + dang_eval_z_0 = x1*x18; + dang_eval_x_1 = x11*y; + dang_eval_y_1 = x*x16; + dang_eval_z_1 = x13; + dang_eval_x_2 = x11*z; + dang_eval_y_2 = x13; + dang_eval_z_2 = x*x20; + dang_eval_x_3 = x12*x5; + dang_eval_y_3 = radial_eval_alpha*x17 + x9*y; + dang_eval_z_3 = x18*x5; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x13; + dang_eval_y_0 = x16*z; + dang_eval_z_0 = x20*y; + dang_eval_x_1 = x12*x8; + dang_eval_y_1 = x14*x8; + dang_eval_z_1 = radial_eval_alpha*x21 + x9*z; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_laplacian.hpp new file mode 100644 index 000000000..4190fdfe4 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l2_laplacian.hpp @@ -0,0 +1,288 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_2_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_2_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_laplacian_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = x*y; + const auto x2 = x*z; + const auto x3 = y*y; + const auto x4 = y*z; + const auto x5 = z*z; + const auto x6 = 2.0*radial_eval; + const auto x7 = x*x*x; + const auto x8 = radial_eval + radial_eval_alpha*x0; + const auto x9 = radial_eval_alpha*x; + const auto x10 = x4*x9; + const auto x11 = radial_eval_alpha*y; + const auto x12 = radial_eval_alpha*x3; + const auto x13 = radial_eval + x12; + const auto x14 = y*y*y; + const auto x15 = radial_eval_alpha*z; + const auto x16 = radial_eval_alpha*x5; + const auto x17 = radial_eval + x16; + const auto x18 = z*z*z; + const auto x19 = 4.0*radial_eval_alpha; + const auto x20 = radial_eval_alpha_squared*x0; + const auto x21 = radial_eval_alpha + x20; + const auto x22 = x0*x19 + x0*x21 + x6; + const auto x23 = 3.0*radial_eval_alpha; + const auto x24 = x20 + x23; + const auto x25 = x21*x3; + const auto x26 = x21*x4; + const auto x27 = x21*x5; + const auto x28 = radial_eval_alpha_squared*x7 + 2.0*x9; + const auto x29 = radial_eval_alpha_squared*x14 + 2.0*x11; + const auto x30 = radial_eval_alpha_squared*x3; + const auto x31 = radial_eval_alpha + x30; + const auto x32 = x2*x31; + const auto x33 = radial_eval_alpha_squared*x5; + const auto x34 = radial_eval_alpha + x33; + const auto x35 = x1*x34; + const auto x36 = radial_eval_alpha_squared*x18 + 2.0*x15; + const auto x37 = x0*x31; + const auto x38 = x23 + x30; + const auto x39 = x19*x3 + x3*x31 + x6; + const auto x40 = x31*x5; + const auto x41 = x0*x34; + const auto x42 = x23 + x33; + const auto x43 = x3*x34; + const auto x44 = x19*x5 + x34*x5 + x6; + const auto x45 = 7.0*radial_eval_alpha + x20 + x30 + x33; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = radial_eval*x1; + basis_eval[ipt + 2*npts] = radial_eval*x2; + basis_eval[ipt + 3*npts] = radial_eval*x3; + basis_eval[ipt + 4*npts] = radial_eval*x4; + basis_eval[ipt + 5*npts] = radial_eval*x5; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x7 + x*x6; + basis_x_eval[ipt + 1*npts] = x8*y; + basis_x_eval[ipt + 2*npts] = x8*z; + basis_x_eval[ipt + 3*npts] = x3*x9; + basis_x_eval[ipt + 4*npts] = x10; + basis_x_eval[ipt + 5*npts] = x5*x9; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x11; + basis_y_eval[ipt + 1*npts] = x*x13; + basis_y_eval[ipt + 2*npts] = x10; + basis_y_eval[ipt + 3*npts] = radial_eval_alpha*x14 + x6*y; + basis_y_eval[ipt + 4*npts] = x13*z; + basis_y_eval[ipt + 5*npts] = x11*x5; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x15; + basis_z_eval[ipt + 1*npts] = x10; + basis_z_eval[ipt + 2*npts] = x*x17; + basis_z_eval[ipt + 3*npts] = x15*x3; + basis_z_eval[ipt + 4*npts] = x17*y; + basis_z_eval[ipt + 5*npts] = radial_eval_alpha*x18 + x6*z; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x22 + x37 + x41; + basis_lapl_eval[ipt + 1*npts] = x1*x45; + basis_lapl_eval[ipt + 2*npts] = x2*x45; + basis_lapl_eval[ipt + 3*npts] = x25 + x39 + x43; + basis_lapl_eval[ipt + 4*npts] = x4*x45; + basis_lapl_eval[ipt + 5*npts] = x27 + x40 + x44; + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = radial_eval*x1; + ang_eval_2 = radial_eval*x2; + ang_eval_3 = radial_eval*x3; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x4; + ang_eval_1 = radial_eval*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x7 + x*x6; + dang_eval_y_0 = x0*x11; + dang_eval_z_0 = x0*x15; + dang_eval_x_1 = x8*y; + dang_eval_y_1 = x*x13; + dang_eval_z_1 = x10; + dang_eval_x_2 = x8*z; + dang_eval_y_2 = x10; + dang_eval_z_2 = x*x17; + dang_eval_x_3 = x3*x9; + dang_eval_y_3 = radial_eval_alpha*x14 + x6*y; + dang_eval_z_3 = x15*x3; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x10; + dang_eval_y_0 = x13*z; + dang_eval_z_0 = x17*y; + dang_eval_x_1 = x5*x9; + dang_eval_y_1 = x11*x5; + dang_eval_z_1 = radial_eval_alpha*x18 + x6*z; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3.hpp new file mode 100644 index 000000000..b16e2c473 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3.hpp @@ -0,0 +1,178 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_3_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_3_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_cartesian_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + const auto x0 = radial_eval*y; + const auto x1 = x*x; + const auto x2 = radial_eval*z; + const auto x3 = radial_eval*x; + const auto x4 = y*y; + const auto x5 = z*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*(x*x*x); + basis_eval[ipt + 1*npts] = x0*x1; + basis_eval[ipt + 2*npts] = x1*x2; + basis_eval[ipt + 3*npts] = x3*x4; + basis_eval[ipt + 4*npts] = x*x0*z; + basis_eval[ipt + 5*npts] = x3*x5; + basis_eval[ipt + 6*npts] = radial_eval*(y*y*y); + basis_eval[ipt + 7*npts] = x2*x4; + basis_eval[ipt + 8*npts] = x0*x5; + basis_eval[ipt + 9*npts] = radial_eval*(z*z*z); + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*(x*x*x); + ang_eval_1 = x0*x1; + ang_eval_2 = x1*x2; + ang_eval_3 = x3*x4; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x*x0*z; + ang_eval_1 = x3*x5; + ang_eval_2 = radial_eval*(y*y*y); + ang_eval_3 = x2*x4; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x0*x5; + ang_eval_1 = radial_eval*(z*z*z); + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_gradient.hpp new file mode 100644 index 000000000..e781f6df5 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_gradient.hpp @@ -0,0 +1,308 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_3_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_3_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_cartesian_gradient_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x; + const auto x3 = radial_eval*z; + const auto x4 = radial_eval*x; + const auto x5 = y*y; + const auto x6 = x*z; + const auto x7 = z*z; + const auto x8 = y*y*y; + const auto x9 = z*z*z; + const auto x10 = 3.0*radial_eval; + const auto x11 = radial_eval_alpha*x0 + 2.0*x4; + const auto x12 = radial_eval*x5; + const auto x13 = radial_eval_alpha*x2*x5; + const auto x14 = y*z; + const auto x15 = radial_eval*x7; + const auto x16 = radial_eval_alpha*x2*x7; + const auto x17 = radial_eval_alpha*x; + const auto x18 = x17*x5*z; + const auto x19 = x17*x7*y; + const auto x20 = radial_eval_alpha*y; + const auto x21 = radial_eval*x2; + const auto x22 = radial_eval_alpha*x14*x2; + const auto x23 = radial_eval_alpha*x8 + 2.0*x1; + const auto x24 = radial_eval_alpha*x5*x7; + const auto x25 = radial_eval_alpha*z; + const auto x26 = radial_eval_alpha*x9 + 2.0*x3; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x2; + basis_eval[ipt + 2*npts] = x2*x3; + basis_eval[ipt + 3*npts] = x4*x5; + basis_eval[ipt + 4*npts] = x1*x6; + basis_eval[ipt + 5*npts] = x4*x7; + basis_eval[ipt + 6*npts] = radial_eval*x8; + basis_eval[ipt + 7*npts] = x3*x5; + basis_eval[ipt + 8*npts] = x1*x7; + basis_eval[ipt + 9*npts] = radial_eval*x9; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*(x*x*x*x) + x10*x2; + basis_x_eval[ipt + 1*npts] = x11*y; + basis_x_eval[ipt + 2*npts] = x11*z; + basis_x_eval[ipt + 3*npts] = x12 + x13; + basis_x_eval[ipt + 4*npts] = x14*(radial_eval + radial_eval_alpha*x2); + basis_x_eval[ipt + 5*npts] = x15 + x16; + basis_x_eval[ipt + 6*npts] = x17*x8; + basis_x_eval[ipt + 7*npts] = x18; + basis_x_eval[ipt + 8*npts] = x19; + basis_x_eval[ipt + 9*npts] = x17*x9; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x20; + basis_y_eval[ipt + 1*npts] = x13 + x21; + basis_y_eval[ipt + 2*npts] = x22; + basis_y_eval[ipt + 3*npts] = x*x23; + basis_y_eval[ipt + 4*npts] = x6*(radial_eval + radial_eval_alpha*x5); + basis_y_eval[ipt + 5*npts] = x19; + basis_y_eval[ipt + 6*npts] = radial_eval_alpha*(y*y*y*y) + x10*x5; + basis_y_eval[ipt + 7*npts] = x23*z; + basis_y_eval[ipt + 8*npts] = x15 + x24; + basis_y_eval[ipt + 9*npts] = x20*x9; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x25; + basis_z_eval[ipt + 1*npts] = x22; + basis_z_eval[ipt + 2*npts] = x16 + x21; + basis_z_eval[ipt + 3*npts] = x18; + basis_z_eval[ipt + 4*npts] = x*y*(radial_eval + radial_eval_alpha*x7); + basis_z_eval[ipt + 5*npts] = x*x26; + basis_z_eval[ipt + 6*npts] = x25*x8; + basis_z_eval[ipt + 7*npts] = x12 + x24; + basis_z_eval[ipt + 8*npts] = x26*y; + basis_z_eval[ipt + 9*npts] = radial_eval_alpha*(z*z*z*z) + x10*x7; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x2; + ang_eval_2 = x2*x3; + ang_eval_3 = x4*x5; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x1*x6; + ang_eval_1 = x4*x7; + ang_eval_2 = radial_eval*x8; + ang_eval_3 = x3*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x1*x7; + ang_eval_1 = radial_eval*x9; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*(x*x*x*x) + x10*x2; + dang_eval_y_0 = x0*x20; + dang_eval_z_0 = x0*x25; + dang_eval_x_1 = x11*y; + dang_eval_y_1 = x13 + x21; + dang_eval_z_1 = x22; + dang_eval_x_2 = x11*z; + dang_eval_y_2 = x22; + dang_eval_z_2 = x16 + x21; + dang_eval_x_3 = x12 + x13; + dang_eval_y_3 = x*x23; + dang_eval_z_3 = x18; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x14*(radial_eval + radial_eval_alpha*x2); + dang_eval_y_0 = x6*(radial_eval + radial_eval_alpha*x5); + dang_eval_z_0 = x*y*(radial_eval + radial_eval_alpha*x7); + dang_eval_x_1 = x15 + x16; + dang_eval_y_1 = x19; + dang_eval_z_1 = x*x26; + dang_eval_x_2 = x17*x8; + dang_eval_y_2 = radial_eval_alpha*(y*y*y*y) + x10*x5; + dang_eval_z_2 = x25*x8; + dang_eval_x_3 = x18; + dang_eval_y_3 = x23*z; + dang_eval_z_3 = x12 + x24; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x19; + dang_eval_y_0 = x15 + x24; + dang_eval_z_0 = x26*y; + dang_eval_x_1 = x17*x9; + dang_eval_y_1 = x20*x9; + dang_eval_z_1 = radial_eval_alpha*(z*z*z*z) + x10*x7; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_hessian.hpp new file mode 100644 index 000000000..4bd0f6338 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_hessian.hpp @@ -0,0 +1,454 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_3_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_3_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_hessian_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x; + const auto x3 = radial_eval*z; + const auto x4 = radial_eval*x; + const auto x5 = y*y; + const auto x6 = x*z; + const auto x7 = z*z; + const auto x8 = y*y*y; + const auto x9 = z*z*z; + const auto x10 = x*x*x*x; + const auto x11 = 3.0*radial_eval; + const auto x12 = radial_eval_alpha*x0 + 2.0*x4; + const auto x13 = radial_eval*x5; + const auto x14 = x2*x5; + const auto x15 = radial_eval_alpha*x14; + const auto x16 = y*z; + const auto x17 = radial_eval_alpha*x2; + const auto x18 = radial_eval + x17; + const auto x19 = radial_eval*x7; + const auto x20 = x2*x7; + const auto x21 = radial_eval_alpha*x20; + const auto x22 = radial_eval_alpha*x; + const auto x23 = x22*x5*z; + const auto x24 = x22*x7*y; + const auto x25 = radial_eval_alpha*y; + const auto x26 = radial_eval*x2; + const auto x27 = radial_eval_alpha*x16*x2; + const auto x28 = radial_eval_alpha*x8 + 2.0*x1; + const auto x29 = radial_eval_alpha*x5; + const auto x30 = radial_eval + x29; + const auto x31 = y*y*y*y; + const auto x32 = x5*x7; + const auto x33 = radial_eval_alpha*x32; + const auto x34 = radial_eval_alpha*z; + const auto x35 = x*y; + const auto x36 = radial_eval_alpha*x7; + const auto x37 = radial_eval_alpha*x9 + 2.0*x3; + const auto x38 = z*z*z*z; + const auto x39 = 6.0*radial_eval_alpha; + const auto x40 = radial_eval_alpha_squared*x2; + const auto x41 = radial_eval_alpha + x40; + const auto x42 = x0*x39 + x0*x41 + 6.0*x4; + const auto x43 = 4.0*radial_eval_alpha; + const auto x44 = 2.0*radial_eval; + const auto x45 = x2*x41 + x44; + const auto x46 = x2*x43 + x45; + const auto x47 = 2.0*radial_eval_alpha; + const auto x48 = x47*x5; + const auto x49 = x41*x5; + const auto x50 = x*x16; + const auto x51 = 3.0*radial_eval_alpha; + const auto x52 = x47*x7; + const auto x53 = x41*x7; + const auto x54 = x41*x8; + const auto x55 = x41*x9; + const auto x56 = radial_eval_alpha_squared*x10 + x2*x51; + const auto x57 = 2.0*x22; + const auto x58 = x16*(radial_eval_alpha_squared*x0 + x57); + const auto x59 = 2.0*x25; + const auto x60 = radial_eval_alpha_squared*x14; + const auto x61 = x29 + x60; + const auto x62 = radial_eval_alpha_squared*x20; + const auto x63 = x36 + x62; + const auto x64 = radial_eval_alpha_squared*x31 + x5*x51; + const auto x65 = x6*(radial_eval_alpha_squared*x8 + x59); + const auto x66 = radial_eval_alpha_squared*x32; + const auto x67 = x36 + x66; + const auto x68 = 2.0*x34; + const auto x69 = x35*(radial_eval_alpha_squared*x9 + x68); + const auto x70 = radial_eval_alpha_squared*x38 + x51*x7; + const auto x71 = radial_eval_alpha_squared*x5; + const auto x72 = radial_eval_alpha + x71; + const auto x73 = x0*x72; + const auto x74 = x2*x47; + const auto x75 = x2*x72; + const auto x76 = x44 + x5*x72; + const auto x77 = x43*x5 + x76; + const auto x78 = x7*x72; + const auto x79 = 6.0*x1 + x39*x8 + x72*x8; + const auto x80 = x72*x9; + const auto x81 = radial_eval_alpha_squared*x7; + const auto x82 = radial_eval_alpha + x81; + const auto x83 = x0*x82; + const auto x84 = x2*x82; + const auto x85 = x5*x82; + const auto x86 = x44 + x7*x82; + const auto x87 = x43*x7 + x86; + const auto x88 = x8*x82; + const auto x89 = 6.0*x3 + x39*x9 + x82*x9; + const auto x90 = x2*x39 + x45 + x75 + x84; + const auto x91 = x39*x5 + x49 + x76 + x85; + const auto x92 = x39*x7 + x53 + x78 + x86; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x2; + basis_eval[ipt + 2*npts] = x2*x3; + basis_eval[ipt + 3*npts] = x4*x5; + basis_eval[ipt + 4*npts] = x1*x6; + basis_eval[ipt + 5*npts] = x4*x7; + basis_eval[ipt + 6*npts] = radial_eval*x8; + basis_eval[ipt + 7*npts] = x3*x5; + basis_eval[ipt + 8*npts] = x1*x7; + basis_eval[ipt + 9*npts] = radial_eval*x9; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x10 + x11*x2; + basis_x_eval[ipt + 1*npts] = x12*y; + basis_x_eval[ipt + 2*npts] = x12*z; + basis_x_eval[ipt + 3*npts] = x13 + x15; + basis_x_eval[ipt + 4*npts] = x16*x18; + basis_x_eval[ipt + 5*npts] = x19 + x21; + basis_x_eval[ipt + 6*npts] = x22*x8; + basis_x_eval[ipt + 7*npts] = x23; + basis_x_eval[ipt + 8*npts] = x24; + basis_x_eval[ipt + 9*npts] = x22*x9; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x25; + basis_y_eval[ipt + 1*npts] = x15 + x26; + basis_y_eval[ipt + 2*npts] = x27; + basis_y_eval[ipt + 3*npts] = x*x28; + basis_y_eval[ipt + 4*npts] = x30*x6; + basis_y_eval[ipt + 5*npts] = x24; + basis_y_eval[ipt + 6*npts] = radial_eval_alpha*x31 + x11*x5; + basis_y_eval[ipt + 7*npts] = x28*z; + basis_y_eval[ipt + 8*npts] = x19 + x33; + basis_y_eval[ipt + 9*npts] = x25*x9; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x34; + basis_z_eval[ipt + 1*npts] = x27; + basis_z_eval[ipt + 2*npts] = x21 + x26; + basis_z_eval[ipt + 3*npts] = x23; + basis_z_eval[ipt + 4*npts] = x35*(radial_eval + x36); + basis_z_eval[ipt + 5*npts] = x*x37; + basis_z_eval[ipt + 6*npts] = x34*x8; + basis_z_eval[ipt + 7*npts] = x13 + x33; + basis_z_eval[ipt + 8*npts] = x37*y; + basis_z_eval[ipt + 9*npts] = radial_eval_alpha*x38 + x11*x7; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x42; + basis_xx_eval[ipt + 1*npts] = x46*y; + basis_xx_eval[ipt + 2*npts] = x46*z; + basis_xx_eval[ipt + 3*npts] = x*(x48 + x49); + basis_xx_eval[ipt + 4*npts] = x50*(x40 + x51); + basis_xx_eval[ipt + 5*npts] = x*(x52 + x53); + basis_xx_eval[ipt + 6*npts] = x54; + basis_xx_eval[ipt + 7*npts] = x49*z; + basis_xx_eval[ipt + 8*npts] = x53*y; + basis_xx_eval[ipt + 9*npts] = x55; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x56*y; + basis_xy_eval[ipt + 1*npts] = radial_eval_alpha_squared*x0*x5 + x12 + x5*x57; + basis_xy_eval[ipt + 2*npts] = x58; + basis_xy_eval[ipt + 3*npts] = radial_eval_alpha_squared*x2*x8 + x2*x59 + x28; + basis_xy_eval[ipt + 4*npts] = z*(x18 + x61); + basis_xy_eval[ipt + 5*npts] = x63*y; + basis_xy_eval[ipt + 6*npts] = x*x64; + basis_xy_eval[ipt + 7*npts] = x65; + basis_xy_eval[ipt + 8*npts] = x*x67; + basis_xy_eval[ipt + 9*npts] = radial_eval_alpha_squared*x35*x9; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x56*z; + basis_xz_eval[ipt + 1*npts] = x58; + basis_xz_eval[ipt + 2*npts] = radial_eval_alpha_squared*x0*x7 + x12 + x57*x7; + basis_xz_eval[ipt + 3*npts] = x61*z; + basis_xz_eval[ipt + 4*npts] = y*(x18 + x63); + basis_xz_eval[ipt + 5*npts] = radial_eval_alpha_squared*x2*x9 + x2*x68 + x37; + basis_xz_eval[ipt + 6*npts] = radial_eval_alpha_squared*x6*x8; + basis_xz_eval[ipt + 7*npts] = x*(x29 + x66); + basis_xz_eval[ipt + 8*npts] = x69; + basis_xz_eval[ipt + 9*npts] = x*x70; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x73; + basis_yy_eval[ipt + 1*npts] = y*(x74 + x75); + basis_yy_eval[ipt + 2*npts] = x75*z; + basis_yy_eval[ipt + 3*npts] = x*x77; + basis_yy_eval[ipt + 4*npts] = x50*(x51 + x71); + basis_yy_eval[ipt + 5*npts] = x*x78; + basis_yy_eval[ipt + 6*npts] = x79; + basis_yy_eval[ipt + 7*npts] = x77*z; + basis_yy_eval[ipt + 8*npts] = y*(x52 + x78); + basis_yy_eval[ipt + 9*npts] = x80; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*x0*x16; + basis_yz_eval[ipt + 1*npts] = z*(x17 + x60); + basis_yz_eval[ipt + 2*npts] = y*(x17 + x62); + basis_yz_eval[ipt + 3*npts] = x65; + basis_yz_eval[ipt + 4*npts] = x*(x30 + x67); + basis_yz_eval[ipt + 5*npts] = x69; + basis_yz_eval[ipt + 6*npts] = x64*z; + basis_yz_eval[ipt + 7*npts] = radial_eval_alpha_squared*x7*x8 + x28 + x59*x7; + basis_yz_eval[ipt + 8*npts] = radial_eval_alpha_squared*x5*x9 + x37 + x5*x68; + basis_yz_eval[ipt + 9*npts] = x70*y; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x83; + basis_zz_eval[ipt + 1*npts] = x84*y; + basis_zz_eval[ipt + 2*npts] = z*(x74 + x84); + basis_zz_eval[ipt + 3*npts] = x*x85; + basis_zz_eval[ipt + 4*npts] = x50*(x51 + x81); + basis_zz_eval[ipt + 5*npts] = x*x87; + basis_zz_eval[ipt + 6*npts] = x88; + basis_zz_eval[ipt + 7*npts] = z*(x48 + x85); + basis_zz_eval[ipt + 8*npts] = x87*y; + basis_zz_eval[ipt + 9*npts] = x89; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x2; + ang_eval_2 = x2*x3; + ang_eval_3 = x4*x5; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x1*x6; + ang_eval_1 = x4*x7; + ang_eval_2 = radial_eval*x8; + ang_eval_3 = x3*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x1*x7; + ang_eval_1 = radial_eval*x9; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x10 + x11*x2; + dang_eval_y_0 = x0*x25; + dang_eval_z_0 = x0*x34; + dang_eval_x_1 = x12*y; + dang_eval_y_1 = x15 + x26; + dang_eval_z_1 = x27; + dang_eval_x_2 = x12*z; + dang_eval_y_2 = x27; + dang_eval_z_2 = x21 + x26; + dang_eval_x_3 = x13 + x15; + dang_eval_y_3 = x*x28; + dang_eval_z_3 = x23; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x16*x18; + dang_eval_y_0 = x30*x6; + dang_eval_z_0 = x35*(radial_eval + x36); + dang_eval_x_1 = x19 + x21; + dang_eval_y_1 = x24; + dang_eval_z_1 = x*x37; + dang_eval_x_2 = x22*x8; + dang_eval_y_2 = radial_eval_alpha*x31 + x11*x5; + dang_eval_z_2 = x34*x8; + dang_eval_x_3 = x23; + dang_eval_y_3 = x28*z; + dang_eval_z_3 = x13 + x33; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x24; + dang_eval_y_0 = x19 + x33; + dang_eval_z_0 = x37*y; + dang_eval_x_1 = x22*x9; + dang_eval_y_1 = x25*x9; + dang_eval_z_1 = radial_eval_alpha*x38 + x11*x7; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_lapgrad.hpp new file mode 100644 index 000000000..fbeefba7d --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_lapgrad.hpp @@ -0,0 +1,572 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_3_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_3_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_lapgrad_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x; + const auto x3 = x2; + const auto x4 = radial_eval*z; + const auto x5 = radial_eval*x; + const auto x6 = y*y; + const auto x7 = x6; + const auto x8 = x*z; + const auto x9 = z*z; + const auto x10 = x9; + const auto x11 = y*y*y; + const auto x12 = z*z*z; + const auto x13 = x*x*x*x; + const auto x14 = 3.0*radial_eval; + const auto x15 = radial_eval_alpha*x0 + 2.0*x5; + const auto x16 = radial_eval*x7; + const auto x17 = x3*x7; + const auto x18 = radial_eval_alpha*x17; + const auto x19 = y*z; + const auto x20 = radial_eval_alpha*x3; + const auto x21 = radial_eval + x20; + const auto x22 = radial_eval*x10; + const auto x23 = x10*x3; + const auto x24 = radial_eval_alpha*x23; + const auto x25 = radial_eval_alpha*x; + const auto x26 = x25*x7*z; + const auto x27 = x10*x25*y; + const auto x28 = radial_eval_alpha*y; + const auto x29 = radial_eval*x3; + const auto x30 = radial_eval_alpha*x19*x3; + const auto x31 = radial_eval_alpha*x11 + 2.0*x1; + const auto x32 = radial_eval_alpha*x7; + const auto x33 = radial_eval + x32; + const auto x34 = y*y*y*y; + const auto x35 = x10*x7; + const auto x36 = radial_eval_alpha*x35; + const auto x37 = radial_eval_alpha*z; + const auto x38 = x*y; + const auto x39 = radial_eval_alpha*x10; + const auto x40 = radial_eval_alpha*x12 + 2.0*x4; + const auto x41 = z*z*z*z; + const auto x42 = 6.0*radial_eval_alpha; + const auto x43 = radial_eval_alpha_squared*x3; + const auto x44 = radial_eval_alpha + x43; + const auto x45 = x0*x42 + x0*x44 + 6.0*x5; + const auto x46 = 4.0*radial_eval_alpha; + const auto x47 = 2.0*radial_eval; + const auto x48 = x3*x44; + const auto x49 = x47 + x48; + const auto x50 = x3*x46 + x49; + const auto x51 = 2.0*radial_eval_alpha; + const auto x52 = x51*x7; + const auto x53 = x44*x7; + const auto x54 = x*x19; + const auto x55 = 3.0*radial_eval_alpha; + const auto x56 = x10*x51; + const auto x57 = x10*x44; + const auto x58 = x11*x44; + const auto x59 = x12*x44; + const auto x60 = radial_eval_alpha_squared*x13 + x3*x55; + const auto x61 = 2.0*x25; + const auto x62 = x19*(radial_eval_alpha_squared*x0 + x61); + const auto x63 = 2.0*x28; + const auto x64 = radial_eval_alpha_squared*x17; + const auto x65 = x32 + x64; + const auto x66 = radial_eval_alpha_squared*x23; + const auto x67 = x39 + x66; + const auto x68 = radial_eval_alpha_squared*x34 + x55*x7; + const auto x69 = x8*(radial_eval_alpha_squared*x11 + x63); + const auto x70 = radial_eval_alpha_squared*x35; + const auto x71 = x39 + x70; + const auto x72 = 2.0*x37; + const auto x73 = x38*(radial_eval_alpha_squared*x12 + x72); + const auto x74 = radial_eval_alpha_squared*x41 + x10*x55; + const auto x75 = radial_eval_alpha_squared*x7; + const auto x76 = radial_eval_alpha + x75; + const auto x77 = x0*x76; + const auto x78 = x3*x51; + const auto x79 = x3*x76; + const auto x80 = x7*x76; + const auto x81 = x47 + x80; + const auto x82 = x46*x7 + x81; + const auto x83 = x10*x76; + const auto x84 = 6.0*x1 + x11*x42 + x11*x76; + const auto x85 = x12*x76; + const auto x86 = radial_eval_alpha_squared*x10; + const auto x87 = radial_eval_alpha + x86; + const auto x88 = x0*x87; + const auto x89 = x3*x87; + const auto x90 = x7*x87; + const auto x91 = x10*x87; + const auto x92 = x47 + x91; + const auto x93 = x10*x46 + x92; + const auto x94 = x11*x87; + const auto x95 = x12*x42 + x12*x87 + 6.0*x4; + const auto x96 = x3*x42 + x49 + x79 + x89; + const auto x97 = x42*x7 + x53 + x81 + x90; + const auto x98 = x75 + x86; + const auto x99 = x10*x42 + x57 + x83 + x92; + const auto x100 = 6.0*radial_eval; + const auto x101 = 18.0*radial_eval_alpha; + const auto x102 = 3.0*x79; + const auto x103 = 3.0*x89; + const auto x104 = radial_eval_alpha_cubed*x7 + radial_eval_alpha_squared; + const auto x105 = x0*x104; + const auto x106 = radial_eval_alpha_cubed*x10 + radial_eval_alpha_squared; + const auto x107 = x0*x106; + const auto x108 = 3.0*radial_eval_alpha_squared; + const auto x109 = radial_eval_alpha_cubed*x0 + x*x108; + const auto x110 = 2.0*radial_eval_alpha_squared; + const auto x111 = 6.0*x; + const auto x112 = 2.0*x; + const auto x113 = x104*x3; + const auto x114 = x106*x3; + const auto x115 = x*x113 + x*x114 + x0*x110 + x109*x3 + x111*x44 + x112*x76 + x112*x87 + 10.0*x25; + const auto x116 = 3.0*x53; + const auto x117 = x109*x7; + const auto x118 = x104*x7; + const auto x119 = x106*x7; + const auto x120 = 4.0*radial_eval_alpha_squared; + const auto x121 = x120*x17; + const auto x122 = 3.0*x57; + const auto x123 = x10*x109; + const auto x124 = x10*x104; + const auto x125 = x10*x106; + const auto x126 = x120*x23; + const auto x127 = 6.0*y; + const auto x128 = x127*x25; + const auto x129 = radial_eval_alpha_squared*x111; + const auto x130 = x104*x11; + const auto x131 = x106*x11; + const auto x132 = 6.0*z; + const auto x133 = x132*x25; + const auto x134 = x104*x12; + const auto x135 = x106*x12; + const auto x136 = radial_eval_alpha_squared*x127; + const auto x137 = radial_eval_alpha_cubed*x3 + radial_eval_alpha_squared; + const auto x138 = x0*x137; + const auto x139 = radial_eval_alpha_cubed*x11 + x108*y; + const auto x140 = x139*x3; + const auto x141 = x137*x3; + const auto x142 = 2.0*y; + const auto x143 = x137*x7; + const auto x144 = x11*x110 + x119*y + x127*x76 + x139*x7 + x142*x44 + x142*x87 + x143*y + 10.0*x28; + const auto x145 = x42 + x43; + const auto x146 = x10*x137; + const auto x147 = x10*x139; + const auto x148 = 3.0*x90; + const auto x149 = x11*x137; + const auto x150 = 3.0*x83; + const auto x151 = x120*x35; + const auto x152 = x19*x42; + const auto x153 = x12*x137; + const auto x154 = radial_eval_alpha_squared*x132; + const auto x155 = radial_eval_alpha_cubed*x12 + x108*z; + const auto x156 = x155*x3; + const auto x157 = x155*x7; + const auto x158 = 2.0*z; + const auto x159 = x10*x155 + x110*x12 + x124*z + x132*x87 + x146*z + x158*x44 + x158*x76 + 10.0*x37; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x3; + basis_eval[ipt + 2*npts] = x3*x4; + basis_eval[ipt + 3*npts] = x5*x7; + basis_eval[ipt + 4*npts] = x1*x8; + basis_eval[ipt + 5*npts] = x10*x5; + basis_eval[ipt + 6*npts] = radial_eval*x11; + basis_eval[ipt + 7*npts] = x4*x7; + basis_eval[ipt + 8*npts] = x1*x10; + basis_eval[ipt + 9*npts] = radial_eval*x12; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x13 + x14*x3; + basis_x_eval[ipt + 1*npts] = x15*y; + basis_x_eval[ipt + 2*npts] = x15*z; + basis_x_eval[ipt + 3*npts] = x16 + x18; + basis_x_eval[ipt + 4*npts] = x19*x21; + basis_x_eval[ipt + 5*npts] = x22 + x24; + basis_x_eval[ipt + 6*npts] = x11*x25; + basis_x_eval[ipt + 7*npts] = x26; + basis_x_eval[ipt + 8*npts] = x27; + basis_x_eval[ipt + 9*npts] = x12*x25; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x28; + basis_y_eval[ipt + 1*npts] = x18 + x29; + basis_y_eval[ipt + 2*npts] = x30; + basis_y_eval[ipt + 3*npts] = x*x31; + basis_y_eval[ipt + 4*npts] = x33*x8; + basis_y_eval[ipt + 5*npts] = x27; + basis_y_eval[ipt + 6*npts] = radial_eval_alpha*x34 + x14*x7; + basis_y_eval[ipt + 7*npts] = x31*z; + basis_y_eval[ipt + 8*npts] = x22 + x36; + basis_y_eval[ipt + 9*npts] = x12*x28; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x37; + basis_z_eval[ipt + 1*npts] = x30; + basis_z_eval[ipt + 2*npts] = x24 + x29; + basis_z_eval[ipt + 3*npts] = x26; + basis_z_eval[ipt + 4*npts] = x38*(radial_eval + x39); + basis_z_eval[ipt + 5*npts] = x*x40; + basis_z_eval[ipt + 6*npts] = x11*x37; + basis_z_eval[ipt + 7*npts] = x16 + x36; + basis_z_eval[ipt + 8*npts] = x40*y; + basis_z_eval[ipt + 9*npts] = radial_eval_alpha*x41 + x10*x14; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x45; + basis_xx_eval[ipt + 1*npts] = x50*y; + basis_xx_eval[ipt + 2*npts] = x50*z; + basis_xx_eval[ipt + 3*npts] = x*(x52 + x53); + basis_xx_eval[ipt + 4*npts] = x54*(x43 + x55); + basis_xx_eval[ipt + 5*npts] = x*(x56 + x57); + basis_xx_eval[ipt + 6*npts] = x58; + basis_xx_eval[ipt + 7*npts] = x53*z; + basis_xx_eval[ipt + 8*npts] = x57*y; + basis_xx_eval[ipt + 9*npts] = x59; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x60*y; + basis_xy_eval[ipt + 1*npts] = radial_eval_alpha_squared*x0*x7 + x15 + x61*x7; + basis_xy_eval[ipt + 2*npts] = x62; + basis_xy_eval[ipt + 3*npts] = radial_eval_alpha_squared*x11*x3 + x3*x63 + x31; + basis_xy_eval[ipt + 4*npts] = z*(x21 + x65); + basis_xy_eval[ipt + 5*npts] = x67*y; + basis_xy_eval[ipt + 6*npts] = x*x68; + basis_xy_eval[ipt + 7*npts] = x69; + basis_xy_eval[ipt + 8*npts] = x*x71; + basis_xy_eval[ipt + 9*npts] = radial_eval_alpha_squared*x12*x38; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x60*z; + basis_xz_eval[ipt + 1*npts] = x62; + basis_xz_eval[ipt + 2*npts] = radial_eval_alpha_squared*x0*x10 + x10*x61 + x15; + basis_xz_eval[ipt + 3*npts] = x65*z; + basis_xz_eval[ipt + 4*npts] = y*(x21 + x67); + basis_xz_eval[ipt + 5*npts] = radial_eval_alpha_squared*x12*x3 + x3*x72 + x40; + basis_xz_eval[ipt + 6*npts] = radial_eval_alpha_squared*x11*x8; + basis_xz_eval[ipt + 7*npts] = x*(x32 + x70); + basis_xz_eval[ipt + 8*npts] = x73; + basis_xz_eval[ipt + 9*npts] = x*x74; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x77; + basis_yy_eval[ipt + 1*npts] = y*(x78 + x79); + basis_yy_eval[ipt + 2*npts] = x79*z; + basis_yy_eval[ipt + 3*npts] = x*x82; + basis_yy_eval[ipt + 4*npts] = x54*(x55 + x75); + basis_yy_eval[ipt + 5*npts] = x*x83; + basis_yy_eval[ipt + 6*npts] = x84; + basis_yy_eval[ipt + 7*npts] = x82*z; + basis_yy_eval[ipt + 8*npts] = y*(x56 + x83); + basis_yy_eval[ipt + 9*npts] = x85; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*x0*x19; + basis_yz_eval[ipt + 1*npts] = z*(x20 + x64); + basis_yz_eval[ipt + 2*npts] = y*(x20 + x66); + basis_yz_eval[ipt + 3*npts] = x69; + basis_yz_eval[ipt + 4*npts] = x*(x33 + x71); + basis_yz_eval[ipt + 5*npts] = x73; + basis_yz_eval[ipt + 6*npts] = x68*z; + basis_yz_eval[ipt + 7*npts] = radial_eval_alpha_squared*x10*x11 + x10*x63 + x31; + basis_yz_eval[ipt + 8*npts] = radial_eval_alpha_squared*x12*x7 + x40 + x7*x72; + basis_yz_eval[ipt + 9*npts] = x74*y; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x88; + basis_zz_eval[ipt + 1*npts] = x89*y; + basis_zz_eval[ipt + 2*npts] = z*(x78 + x89); + basis_zz_eval[ipt + 3*npts] = x*x90; + basis_zz_eval[ipt + 4*npts] = x54*(x55 + x86); + basis_zz_eval[ipt + 5*npts] = x*x93; + basis_zz_eval[ipt + 6*npts] = x94; + basis_zz_eval[ipt + 7*npts] = z*(x52 + x90); + basis_zz_eval[ipt + 8*npts] = x93*y; + basis_zz_eval[ipt + 9*npts] = x95; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x45 + x77 + x88; + basis_lapl_eval[ipt + 1*npts] = x96*y; + basis_lapl_eval[ipt + 2*npts] = x96*z; + basis_lapl_eval[ipt + 3*npts] = x*x97; + basis_lapl_eval[ipt + 4*npts] = x54*(9.0*radial_eval_alpha + x43 + x98); + basis_lapl_eval[ipt + 5*npts] = x*x99; + basis_lapl_eval[ipt + 6*npts] = x58 + x84 + x94; + basis_lapl_eval[ipt + 7*npts] = x97*z; + basis_lapl_eval[ipt + 8*npts] = x99*y; + basis_lapl_eval[ipt + 9*npts] = x59 + x85 + x95; + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = x*x105 + x*x107 + x0*x109 + x100 + x101*x3 + x102 + x103 + 9.0*x48; + basis_lapl_x_eval[ipt + 1*npts] = x115*y; + basis_lapl_x_eval[ipt + 2*npts] = x115*z; + basis_lapl_x_eval[ipt + 3*npts] = x*x117 + x116 + x118*x2 + x119*x2 + x121 + x78 + x82 + x90; + basis_lapl_x_eval[ipt + 4*npts] = x19*(x*x109 + x104*x2 + x106*x2 + x120*x3 + x42 + 3.0*x44 + x98); + basis_lapl_x_eval[ipt + 5*npts] = x*x123 + x122 + x124*x2 + x125*x2 + x126 + x78 + x83 + x93; + basis_lapl_x_eval[ipt + 6*npts] = x*x130 + x*x131 + x109*x11 + x11*x129 + x128; + basis_lapl_x_eval[ipt + 7*npts] = z*(x*x118 + x*x119 + x117 + x129*x7 + x61); + basis_lapl_x_eval[ipt + 8*npts] = y*(x*x124 + x*x125 + x10*x129 + x123 + x61); + basis_lapl_x_eval[ipt + 9*npts] = x*x134 + x*x135 + x109*x12 + x12*x129 + x133; + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x0*x136 + x0*x139 + x107*y + x128 + x138*y; + basis_lapl_y_eval[ipt + 1*npts] = x102 + x114*x6 + x121 + x140*y + x141*x6 + x50 + x52 + x89; + basis_lapl_y_eval[ipt + 2*npts] = z*(x114*y + x136*x3 + x140 + x141*y + x63); + basis_lapl_y_eval[ipt + 3*npts] = x*x144; + basis_lapl_y_eval[ipt + 4*npts] = x8*(x106*x6 + x120*x7 + x137*x6 + x139*y + x145 + 3.0*x76 + x86); + basis_lapl_y_eval[ipt + 5*npts] = x*(x10*x136 + x125*y + x146*y + x147 + x63); + basis_lapl_y_eval[ipt + 6*npts] = x100 + x101*x7 + x11*x139 + x116 + x131*y + x148 + x149*y + 9.0*x80; + basis_lapl_y_eval[ipt + 7*npts] = x144*z; + basis_lapl_y_eval[ipt + 8*npts] = x125*x6 + x146*x6 + x147*y + x150 + x151 + x52 + x57 + x93; + basis_lapl_y_eval[ipt + 9*npts] = x12*x136 + x12*x139 + x135*y + x152 + x153*y; + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x0*x154 + x0*x155 + x105*z + x133 + x138*z; + basis_lapl_z_eval[ipt + 1*npts] = y*(x113*z + x141*z + x154*x3 + x156 + x72); + basis_lapl_z_eval[ipt + 2*npts] = x103 + x113*x9 + x126 + x141*x9 + x156*z + x50 + x56 + x79; + basis_lapl_z_eval[ipt + 3*npts] = x*(x118*z + x143*z + x154*x7 + x157 + x72); + basis_lapl_z_eval[ipt + 4*npts] = x38*(x10*x120 + x104*x9 + x137*x9 + x145 + x155*z + x75 + 3.0*x87); + basis_lapl_z_eval[ipt + 5*npts] = x*x159; + basis_lapl_z_eval[ipt + 6*npts] = x11*x154 + x11*x155 + x130*z + x149*z + x152; + basis_lapl_z_eval[ipt + 7*npts] = x118*x9 + x143*x9 + x148 + x151 + x157*z + x53 + x56 + x82; + basis_lapl_z_eval[ipt + 8*npts] = x159*y; + basis_lapl_z_eval[ipt + 9*npts] = x10*x101 + x100 + x12*x155 + x122 + x134*z + x150 + x153*z + 9.0*x91; + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x3; + ang_eval_2 = x3*x4; + ang_eval_3 = x5*x7; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x1*x8; + ang_eval_1 = x10*x5; + ang_eval_2 = radial_eval*x11; + ang_eval_3 = x4*x7; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x1*x10; + ang_eval_1 = radial_eval*x12; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x13 + x14*x3; + dang_eval_y_0 = x0*x28; + dang_eval_z_0 = x0*x37; + dang_eval_x_1 = x15*y; + dang_eval_y_1 = x18 + x29; + dang_eval_z_1 = x30; + dang_eval_x_2 = x15*z; + dang_eval_y_2 = x30; + dang_eval_z_2 = x24 + x29; + dang_eval_x_3 = x16 + x18; + dang_eval_y_3 = x*x31; + dang_eval_z_3 = x26; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x19*x21; + dang_eval_y_0 = x33*x8; + dang_eval_z_0 = x38*(radial_eval + x39); + dang_eval_x_1 = x22 + x24; + dang_eval_y_1 = x27; + dang_eval_z_1 = x*x40; + dang_eval_x_2 = x11*x25; + dang_eval_y_2 = radial_eval_alpha*x34 + x14*x7; + dang_eval_z_2 = x11*x37; + dang_eval_x_3 = x26; + dang_eval_y_3 = x31*z; + dang_eval_z_3 = x16 + x36; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x27; + dang_eval_y_0 = x22 + x36; + dang_eval_z_0 = x40*y; + dang_eval_x_1 = x12*x25; + dang_eval_y_1 = x12*x28; + dang_eval_z_1 = radial_eval_alpha*x41 + x10*x14; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_laplacian.hpp new file mode 100644 index 000000000..bfb30ca91 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l3_laplacian.hpp @@ -0,0 +1,389 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_3_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_3_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_laplacian_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x; + const auto x3 = radial_eval*z; + const auto x4 = radial_eval*x; + const auto x5 = y*y; + const auto x6 = x*z; + const auto x7 = z*z; + const auto x8 = y*y*y; + const auto x9 = z*z*z; + const auto x10 = x*x*x*x; + const auto x11 = 3.0*radial_eval; + const auto x12 = radial_eval_alpha*x0 + 2.0*x4; + const auto x13 = radial_eval*x5; + const auto x14 = x2*x5; + const auto x15 = radial_eval_alpha*x14; + const auto x16 = y*z; + const auto x17 = radial_eval_alpha*x2; + const auto x18 = radial_eval + x17; + const auto x19 = radial_eval*x7; + const auto x20 = x2*x7; + const auto x21 = radial_eval_alpha*x20; + const auto x22 = radial_eval_alpha*x; + const auto x23 = x22*x5*z; + const auto x24 = x22*x7*y; + const auto x25 = radial_eval_alpha*y; + const auto x26 = radial_eval*x2; + const auto x27 = radial_eval_alpha*x16*x2; + const auto x28 = radial_eval_alpha*x8 + 2.0*x1; + const auto x29 = radial_eval_alpha*x5; + const auto x30 = radial_eval + x29; + const auto x31 = y*y*y*y; + const auto x32 = x5*x7; + const auto x33 = radial_eval_alpha*x32; + const auto x34 = radial_eval_alpha*z; + const auto x35 = x*y; + const auto x36 = radial_eval_alpha*x7; + const auto x37 = radial_eval_alpha*x9 + 2.0*x3; + const auto x38 = z*z*z*z; + const auto x39 = 6.0*radial_eval_alpha; + const auto x40 = radial_eval_alpha_squared*x2; + const auto x41 = radial_eval_alpha + x40; + const auto x42 = x0*x39 + x0*x41 + 6.0*x4; + const auto x43 = 4.0*radial_eval_alpha; + const auto x44 = 2.0*radial_eval; + const auto x45 = x2*x41 + x44; + const auto x46 = x2*x43 + x45; + const auto x47 = 2.0*radial_eval_alpha; + const auto x48 = x47*x5; + const auto x49 = x41*x5; + const auto x50 = x*x16; + const auto x51 = 3.0*radial_eval_alpha; + const auto x52 = x47*x7; + const auto x53 = x41*x7; + const auto x54 = x41*x8; + const auto x55 = x41*x9; + const auto x56 = radial_eval_alpha_squared*x10 + x2*x51; + const auto x57 = 2.0*x22; + const auto x58 = x16*(radial_eval_alpha_squared*x0 + x57); + const auto x59 = 2.0*x25; + const auto x60 = radial_eval_alpha_squared*x14; + const auto x61 = x29 + x60; + const auto x62 = radial_eval_alpha_squared*x20; + const auto x63 = x36 + x62; + const auto x64 = radial_eval_alpha_squared*x31 + x5*x51; + const auto x65 = x6*(radial_eval_alpha_squared*x8 + x59); + const auto x66 = radial_eval_alpha_squared*x32; + const auto x67 = x36 + x66; + const auto x68 = 2.0*x34; + const auto x69 = x35*(radial_eval_alpha_squared*x9 + x68); + const auto x70 = radial_eval_alpha_squared*x38 + x51*x7; + const auto x71 = radial_eval_alpha_squared*x5; + const auto x72 = radial_eval_alpha + x71; + const auto x73 = x0*x72; + const auto x74 = x2*x47; + const auto x75 = x2*x72; + const auto x76 = x44 + x5*x72; + const auto x77 = x43*x5 + x76; + const auto x78 = x7*x72; + const auto x79 = 6.0*x1 + x39*x8 + x72*x8; + const auto x80 = x72*x9; + const auto x81 = radial_eval_alpha_squared*x7; + const auto x82 = radial_eval_alpha + x81; + const auto x83 = x0*x82; + const auto x84 = x2*x82; + const auto x85 = x5*x82; + const auto x86 = x44 + x7*x82; + const auto x87 = x43*x7 + x86; + const auto x88 = x8*x82; + const auto x89 = 6.0*x3 + x39*x9 + x82*x9; + const auto x90 = x2*x39 + x45 + x75 + x84; + const auto x91 = x39*x5 + x49 + x76 + x85; + const auto x92 = x39*x7 + x53 + x78 + x86; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x2; + basis_eval[ipt + 2*npts] = x2*x3; + basis_eval[ipt + 3*npts] = x4*x5; + basis_eval[ipt + 4*npts] = x1*x6; + basis_eval[ipt + 5*npts] = x4*x7; + basis_eval[ipt + 6*npts] = radial_eval*x8; + basis_eval[ipt + 7*npts] = x3*x5; + basis_eval[ipt + 8*npts] = x1*x7; + basis_eval[ipt + 9*npts] = radial_eval*x9; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x10 + x11*x2; + basis_x_eval[ipt + 1*npts] = x12*y; + basis_x_eval[ipt + 2*npts] = x12*z; + basis_x_eval[ipt + 3*npts] = x13 + x15; + basis_x_eval[ipt + 4*npts] = x16*x18; + basis_x_eval[ipt + 5*npts] = x19 + x21; + basis_x_eval[ipt + 6*npts] = x22*x8; + basis_x_eval[ipt + 7*npts] = x23; + basis_x_eval[ipt + 8*npts] = x24; + basis_x_eval[ipt + 9*npts] = x22*x9; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x25; + basis_y_eval[ipt + 1*npts] = x15 + x26; + basis_y_eval[ipt + 2*npts] = x27; + basis_y_eval[ipt + 3*npts] = x*x28; + basis_y_eval[ipt + 4*npts] = x30*x6; + basis_y_eval[ipt + 5*npts] = x24; + basis_y_eval[ipt + 6*npts] = radial_eval_alpha*x31 + x11*x5; + basis_y_eval[ipt + 7*npts] = x28*z; + basis_y_eval[ipt + 8*npts] = x19 + x33; + basis_y_eval[ipt + 9*npts] = x25*x9; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x34; + basis_z_eval[ipt + 1*npts] = x27; + basis_z_eval[ipt + 2*npts] = x21 + x26; + basis_z_eval[ipt + 3*npts] = x23; + basis_z_eval[ipt + 4*npts] = x35*(radial_eval + x36); + basis_z_eval[ipt + 5*npts] = x*x37; + basis_z_eval[ipt + 6*npts] = x34*x8; + basis_z_eval[ipt + 7*npts] = x13 + x33; + basis_z_eval[ipt + 8*npts] = x37*y; + basis_z_eval[ipt + 9*npts] = radial_eval_alpha*x38 + x11*x7; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x42 + x73 + x83; + basis_lapl_eval[ipt + 1*npts] = x90*y; + basis_lapl_eval[ipt + 2*npts] = x90*z; + basis_lapl_eval[ipt + 3*npts] = x*x91; + basis_lapl_eval[ipt + 4*npts] = x50*(9.0*radial_eval_alpha + x40 + x71 + x81); + basis_lapl_eval[ipt + 5*npts] = x*x92; + basis_lapl_eval[ipt + 6*npts] = x54 + x79 + x88; + basis_lapl_eval[ipt + 7*npts] = x91*z; + basis_lapl_eval[ipt + 8*npts] = x92*y; + basis_lapl_eval[ipt + 9*npts] = x55 + x80 + x89; + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x2; + ang_eval_2 = x2*x3; + ang_eval_3 = x4*x5; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x1*x6; + ang_eval_1 = x4*x7; + ang_eval_2 = radial_eval*x8; + ang_eval_3 = x3*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x1*x7; + ang_eval_1 = radial_eval*x9; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x10 + x11*x2; + dang_eval_y_0 = x0*x25; + dang_eval_z_0 = x0*x34; + dang_eval_x_1 = x12*y; + dang_eval_y_1 = x15 + x26; + dang_eval_z_1 = x27; + dang_eval_x_2 = x12*z; + dang_eval_y_2 = x27; + dang_eval_z_2 = x21 + x26; + dang_eval_x_3 = x13 + x15; + dang_eval_y_3 = x*x28; + dang_eval_z_3 = x23; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x16*x18; + dang_eval_y_0 = x30*x6; + dang_eval_z_0 = x35*(radial_eval + x36); + dang_eval_x_1 = x19 + x21; + dang_eval_y_1 = x24; + dang_eval_z_1 = x*x37; + dang_eval_x_2 = x22*x8; + dang_eval_y_2 = radial_eval_alpha*x31 + x11*x5; + dang_eval_z_2 = x34*x8; + dang_eval_x_3 = x23; + dang_eval_y_3 = x28*z; + dang_eval_z_3 = x13 + x33; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x24; + dang_eval_y_0 = x19 + x33; + dang_eval_z_0 = x37*y; + dang_eval_x_1 = x22*x9; + dang_eval_y_1 = x25*x9; + dang_eval_z_1 = radial_eval_alpha*x38 + x11*x7; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4.hpp new file mode 100644 index 000000000..eda599dad --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4.hpp @@ -0,0 +1,197 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_4_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_4_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_cartesian_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + const auto x0 = radial_eval*y; + const auto x1 = x*x*x; + const auto x2 = radial_eval*z; + const auto x3 = x*x; + const auto x4 = y*y; + const auto x5 = z*z; + const auto x6 = radial_eval*x; + const auto x7 = y*y*y; + const auto x8 = z*z*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*(x*x*x*x); + basis_eval[ipt + 1*npts] = x0*x1; + basis_eval[ipt + 2*npts] = x1*x2; + basis_eval[ipt + 3*npts] = radial_eval*x3*x4; + basis_eval[ipt + 4*npts] = x0*x3*z; + basis_eval[ipt + 5*npts] = radial_eval*x3*x5; + basis_eval[ipt + 6*npts] = x6*x7; + basis_eval[ipt + 7*npts] = x*x2*x4; + basis_eval[ipt + 8*npts] = x*x0*x5; + basis_eval[ipt + 9*npts] = x6*x8; + basis_eval[ipt + 10*npts] = radial_eval*(y*y*y*y); + basis_eval[ipt + 11*npts] = x2*x7; + basis_eval[ipt + 12*npts] = radial_eval*x4*x5; + basis_eval[ipt + 13*npts] = x0*x8; + basis_eval[ipt + 14*npts] = radial_eval*(z*z*z*z); + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*(x*x*x*x); + ang_eval_1 = x0*x1; + ang_eval_2 = x1*x2; + ang_eval_3 = radial_eval*x3*x4; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x0*x3*z; + ang_eval_1 = radial_eval*x3*x5; + ang_eval_2 = x6*x7; + ang_eval_3 = x*x2*x4; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x*x0*x5; + ang_eval_1 = x6*x8; + ang_eval_2 = radial_eval*(y*y*y*y); + ang_eval_3 = x2*x7; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + basis_eval[ipt + 10*npts] = ang_eval_2; + basis_eval[ipt + 11*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x4*x5; + ang_eval_1 = x0*x8; + ang_eval_2 = radial_eval*(z*z*z*z); + basis_eval[ipt + 12*npts] = ang_eval_0; + basis_eval[ipt + 13*npts] = ang_eval_1; + basis_eval[ipt + 14*npts] = ang_eval_2; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_gradient.hpp new file mode 100644 index 000000000..3b5722e79 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_gradient.hpp @@ -0,0 +1,388 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_gradient_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_gradient_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = x*x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x*x; + const auto x3 = radial_eval*z; + const auto x4 = x*x; + const auto x5 = y*y; + const auto x6 = x4*x5; + const auto x7 = z*z; + const auto x8 = x4*x7; + const auto x9 = radial_eval*x; + const auto x10 = y*y*y; + const auto x11 = z*z*z; + const auto x12 = y*y*y*y; + const auto x13 = x5*x7; + const auto x14 = z*z*z*z; + const auto x15 = 4.0*radial_eval; + const auto x16 = 3.0*radial_eval; + const auto x17 = radial_eval_alpha*x0 + x16*x4; + const auto x18 = 2.0*x9; + const auto x19 = radial_eval_alpha*x2*x5; + const auto x20 = y*z; + const auto x21 = radial_eval_alpha*x2*x7; + const auto x22 = radial_eval*x10; + const auto x23 = radial_eval_alpha*x10*x4; + const auto x24 = radial_eval*x5; + const auto x25 = radial_eval_alpha*x6; + const auto x26 = radial_eval*x7; + const auto x27 = radial_eval_alpha*x8; + const auto x28 = radial_eval*x11; + const auto x29 = radial_eval_alpha*x11*x4; + const auto x30 = radial_eval_alpha*x; + const auto x31 = x10*x30*z; + const auto x32 = x11*x30*y; + const auto x33 = radial_eval_alpha*y; + const auto x34 = radial_eval*x2; + const auto x35 = radial_eval_alpha*x2*x20; + const auto x36 = 2.0*x1; + const auto x37 = radial_eval*x4; + const auto x38 = radial_eval_alpha*x12 + x16*x5; + const auto x39 = radial_eval_alpha*x13; + const auto x40 = radial_eval_alpha*x10*x7; + const auto x41 = radial_eval_alpha*x11*x5; + const auto x42 = radial_eval_alpha*z; + const auto x43 = 2.0*x3; + const auto x44 = radial_eval_alpha*x14 + x16*x7; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x2; + basis_eval[ipt + 2*npts] = x2*x3; + basis_eval[ipt + 3*npts] = radial_eval*x6; + basis_eval[ipt + 4*npts] = x1*x4*z; + basis_eval[ipt + 5*npts] = radial_eval*x8; + basis_eval[ipt + 6*npts] = x10*x9; + basis_eval[ipt + 7*npts] = x*x3*x5; + basis_eval[ipt + 8*npts] = x*x1*x7; + basis_eval[ipt + 9*npts] = x11*x9; + basis_eval[ipt + 10*npts] = radial_eval*x12; + basis_eval[ipt + 11*npts] = x10*x3; + basis_eval[ipt + 12*npts] = radial_eval*x13; + basis_eval[ipt + 13*npts] = x1*x11; + basis_eval[ipt + 14*npts] = radial_eval*x14; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*(x*x*x*x*x) + x15*x2; + basis_x_eval[ipt + 1*npts] = x17*y; + basis_x_eval[ipt + 2*npts] = x17*z; + basis_x_eval[ipt + 3*npts] = x18*x5 + x19; + basis_x_eval[ipt + 4*npts] = x20*(radial_eval_alpha*x2 + x18); + basis_x_eval[ipt + 5*npts] = x18*x7 + x21; + basis_x_eval[ipt + 6*npts] = x22 + x23; + basis_x_eval[ipt + 7*npts] = z*(x24 + x25); + basis_x_eval[ipt + 8*npts] = y*(x26 + x27); + basis_x_eval[ipt + 9*npts] = x28 + x29; + basis_x_eval[ipt + 10*npts] = x12*x30; + basis_x_eval[ipt + 11*npts] = x31; + basis_x_eval[ipt + 12*npts] = x13*x30; + basis_x_eval[ipt + 13*npts] = x32; + basis_x_eval[ipt + 14*npts] = x14*x30; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x33; + basis_y_eval[ipt + 1*npts] = x19 + x34; + basis_y_eval[ipt + 2*npts] = x35; + basis_y_eval[ipt + 3*npts] = x23 + x36*x4; + basis_y_eval[ipt + 4*npts] = z*(x25 + x37); + basis_y_eval[ipt + 5*npts] = x33*x8; + basis_y_eval[ipt + 6*npts] = x*x38; + basis_y_eval[ipt + 7*npts] = x*z*(radial_eval_alpha*x10 + x36); + basis_y_eval[ipt + 8*npts] = x*(x26 + x39); + basis_y_eval[ipt + 9*npts] = x32; + basis_y_eval[ipt + 10*npts] = radial_eval_alpha*(y*y*y*y*y) + x10*x15; + basis_y_eval[ipt + 11*npts] = x38*z; + basis_y_eval[ipt + 12*npts] = x36*x7 + x40; + basis_y_eval[ipt + 13*npts] = x28 + x41; + basis_y_eval[ipt + 14*npts] = x14*x33; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x42; + basis_z_eval[ipt + 1*npts] = x35; + basis_z_eval[ipt + 2*npts] = x21 + x34; + basis_z_eval[ipt + 3*npts] = x42*x6; + basis_z_eval[ipt + 4*npts] = y*(x27 + x37); + basis_z_eval[ipt + 5*npts] = x29 + x4*x43; + basis_z_eval[ipt + 6*npts] = x31; + basis_z_eval[ipt + 7*npts] = x*(x24 + x39); + basis_z_eval[ipt + 8*npts] = x*y*(radial_eval_alpha*x11 + x43); + basis_z_eval[ipt + 9*npts] = x*x44; + basis_z_eval[ipt + 10*npts] = x12*x42; + basis_z_eval[ipt + 11*npts] = x22 + x40; + basis_z_eval[ipt + 12*npts] = x41 + x43*x5; + basis_z_eval[ipt + 13*npts] = x44*y; + basis_z_eval[ipt + 14*npts] = radial_eval_alpha*(z*z*z*z*z) + x11*x15; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x2; + ang_eval_2 = x2*x3; + ang_eval_3 = radial_eval*x6; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x1*x4*z; + ang_eval_1 = radial_eval*x8; + ang_eval_2 = x10*x9; + ang_eval_3 = x*x3*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x*x1*x7; + ang_eval_1 = x11*x9; + ang_eval_2 = radial_eval*x12; + ang_eval_3 = x10*x3; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + basis_eval[ipt + 10*npts] = ang_eval_2; + basis_eval[ipt + 11*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x13; + ang_eval_1 = x1*x11; + ang_eval_2 = radial_eval*x14; + basis_eval[ipt + 12*npts] = ang_eval_0; + basis_eval[ipt + 13*npts] = ang_eval_1; + basis_eval[ipt + 14*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*(x*x*x*x*x) + x15*x2; + dang_eval_y_0 = x0*x33; + dang_eval_z_0 = x0*x42; + dang_eval_x_1 = x17*y; + dang_eval_y_1 = x19 + x34; + dang_eval_z_1 = x35; + dang_eval_x_2 = x17*z; + dang_eval_y_2 = x35; + dang_eval_z_2 = x21 + x34; + dang_eval_x_3 = x18*x5 + x19; + dang_eval_y_3 = x23 + x36*x4; + dang_eval_z_3 = x42*x6; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x20*(radial_eval_alpha*x2 + x18); + dang_eval_y_0 = z*(x25 + x37); + dang_eval_z_0 = y*(x27 + x37); + dang_eval_x_1 = x18*x7 + x21; + dang_eval_y_1 = x33*x8; + dang_eval_z_1 = x29 + x4*x43; + dang_eval_x_2 = x22 + x23; + dang_eval_y_2 = x*x38; + dang_eval_z_2 = x31; + dang_eval_x_3 = z*(x24 + x25); + dang_eval_y_3 = x*z*(radial_eval_alpha*x10 + x36); + dang_eval_z_3 = x*(x24 + x39); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = y*(x26 + x27); + dang_eval_y_0 = x*(x26 + x39); + dang_eval_z_0 = x*y*(radial_eval_alpha*x11 + x43); + dang_eval_x_1 = x28 + x29; + dang_eval_y_1 = x32; + dang_eval_z_1 = x*x44; + dang_eval_x_2 = x12*x30; + dang_eval_y_2 = radial_eval_alpha*(y*y*y*y*y) + x10*x15; + dang_eval_z_2 = x12*x42; + dang_eval_x_3 = x31; + dang_eval_y_3 = x38*z; + dang_eval_z_3 = x22 + x40; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + basis_x_eval[ipt + 10*npts] = dang_eval_x_2; + basis_y_eval[ipt + 10*npts] = dang_eval_y_2; + basis_z_eval[ipt + 10*npts] = dang_eval_z_2; + basis_x_eval[ipt + 11*npts] = dang_eval_x_3; + basis_y_eval[ipt + 11*npts] = dang_eval_y_3; + basis_z_eval[ipt + 11*npts] = dang_eval_z_3; + + dang_eval_x_0 = x13*x30; + dang_eval_y_0 = x36*x7 + x40; + dang_eval_z_0 = x41 + x43*x5; + dang_eval_x_1 = x32; + dang_eval_y_1 = x28 + x41; + dang_eval_z_1 = x44*y; + dang_eval_x_2 = x14*x30; + dang_eval_y_2 = x14*x33; + dang_eval_z_2 = radial_eval_alpha*(z*z*z*z*z) + x11*x15; + basis_x_eval[ipt + 12*npts] = dang_eval_x_0; + basis_y_eval[ipt + 12*npts] = dang_eval_y_0; + basis_z_eval[ipt + 12*npts] = dang_eval_z_0; + basis_x_eval[ipt + 13*npts] = dang_eval_x_1; + basis_y_eval[ipt + 13*npts] = dang_eval_y_1; + basis_z_eval[ipt + 13*npts] = dang_eval_z_1; + basis_x_eval[ipt + 14*npts] = dang_eval_x_2; + basis_y_eval[ipt + 14*npts] = dang_eval_y_2; + basis_z_eval[ipt + 14*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_hessian.hpp new file mode 100644 index 000000000..aca72628e --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_hessian.hpp @@ -0,0 +1,611 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_hessian_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_hessian_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x*x; + const auto x3 = radial_eval*z; + const auto x4 = x*x; + const auto x5 = y*y; + const auto x6 = x4*x5; + const auto x7 = x1*z; + const auto x8 = z*z; + const auto x9 = x4*x8; + const auto x10 = radial_eval*x; + const auto x11 = y*y*y; + const auto x12 = x*x3; + const auto x13 = x*x1; + const auto x14 = z*z*z; + const auto x15 = y*y*y*y; + const auto x16 = x5*x8; + const auto x17 = z*z*z*z; + const auto x18 = x*x*x*x*x; + const auto x19 = 4.0*radial_eval; + const auto x20 = 3.0*radial_eval; + const auto x21 = radial_eval_alpha*x0 + x20*x4; + const auto x22 = 2.0*x10; + const auto x23 = x2*x5; + const auto x24 = radial_eval_alpha*x23; + const auto x25 = y*z; + const auto x26 = radial_eval_alpha*x2; + const auto x27 = x22 + x26; + const auto x28 = x2*x8; + const auto x29 = radial_eval_alpha*x28; + const auto x30 = radial_eval*x11; + const auto x31 = x11*x4; + const auto x32 = radial_eval_alpha*x31; + const auto x33 = radial_eval*x5; + const auto x34 = radial_eval_alpha*x6; + const auto x35 = x33 + x34; + const auto x36 = radial_eval*x8; + const auto x37 = radial_eval_alpha*x9; + const auto x38 = x36 + x37; + const auto x39 = radial_eval*x14; + const auto x40 = x14*x4; + const auto x41 = radial_eval_alpha*x40; + const auto x42 = radial_eval_alpha*x; + const auto x43 = x11*x42*z; + const auto x44 = x14*x42*y; + const auto x45 = radial_eval_alpha*y; + const auto x46 = radial_eval*x2; + const auto x47 = radial_eval_alpha*x2*x25; + const auto x48 = 2.0*x1; + const auto x49 = radial_eval*x4; + const auto x50 = x34 + x49; + const auto x51 = radial_eval_alpha*x15 + x20*x5; + const auto x52 = x*z; + const auto x53 = radial_eval_alpha*x11; + const auto x54 = x48 + x53; + const auto x55 = radial_eval_alpha*x16; + const auto x56 = y*y*y*y*y; + const auto x57 = x11*x8; + const auto x58 = radial_eval_alpha*x57; + const auto x59 = x14*x5; + const auto x60 = radial_eval_alpha*x59; + const auto x61 = radial_eval_alpha*z; + const auto x62 = 2.0*x3; + const auto x63 = x*y; + const auto x64 = radial_eval_alpha*x14; + const auto x65 = x62 + x64; + const auto x66 = radial_eval_alpha*x17 + x20*x8; + const auto x67 = z*z*z*z*z; + const auto x68 = 12.0*radial_eval; + const auto x69 = 8.0*radial_eval_alpha; + const auto x70 = radial_eval_alpha + radial_eval_alpha_squared*x4; + const auto x71 = x0*x69 + x0*x70 + x4*x68; + const auto x72 = 6.0*radial_eval_alpha; + const auto x73 = 6.0*x10 + x2*x70; + const auto x74 = x2*x72 + x73; + const auto x75 = 4.0*radial_eval_alpha; + const auto x76 = x6*x75; + const auto x77 = 2.0*radial_eval; + const auto x78 = x5*x77; + const auto x79 = x4*x5*x70 + x78; + const auto x80 = x4*x70 + x77; + const auto x81 = x75*x9; + const auto x82 = x77*x8; + const auto x83 = x4*x70*x8 + x82; + const auto x84 = 2.0*radial_eval_alpha; + const auto x85 = x11*x84; + const auto x86 = x11*x70; + const auto x87 = x5*x84; + const auto x88 = x5*x70; + const auto x89 = x8*x84; + const auto x90 = x70*x8; + const auto x91 = x14*x84; + const auto x92 = x14*x70; + const auto x93 = x15*x70; + const auto x94 = x5*x70*x8; + const auto x95 = x17*x70; + const auto x96 = radial_eval_alpha_squared*x18 + x2*x75; + const auto x97 = 3.0*radial_eval_alpha; + const auto x98 = x6*x97; + const auto x99 = x25*(radial_eval_alpha_squared*x0 + x4*x97); + const auto x100 = 2.0*x42; + const auto x101 = 2.0*x45; + const auto x102 = radial_eval_alpha_squared*x23; + const auto x103 = x100*x5 + x102; + const auto x104 = radial_eval_alpha_squared*x28; + const auto x105 = x100*x8 + x104; + const auto x106 = radial_eval_alpha_squared*x31; + const auto x107 = x101*x4 + x106; + const auto x108 = radial_eval_alpha_squared*x4*x5*x8; + const auto x109 = x108 + x55; + const auto x110 = radial_eval_alpha_squared*x40; + const auto x111 = radial_eval_alpha_squared*x56 + x11*x75; + const auto x112 = x52*(radial_eval_alpha_squared*x15 + x5*x97); + const auto x113 = radial_eval_alpha_squared*x57; + const auto x114 = x101*x8 + x113; + const auto x115 = radial_eval_alpha_squared*x59; + const auto x116 = x9*x97; + const auto x117 = 2.0*x61; + const auto x118 = x110 + x117*x4; + const auto x119 = x115 + x117*x5; + const auto x120 = x63*(radial_eval_alpha_squared*x17 + x8*x97); + const auto x121 = radial_eval_alpha_squared*x67 + x14*x75; + const auto x122 = radial_eval_alpha + radial_eval_alpha_squared*x5; + const auto x123 = x0*x122; + const auto x124 = x2*x84; + const auto x125 = x122*x2; + const auto x126 = x4*x77; + const auto x127 = x122*x4*x5 + x126; + const auto x128 = x4*x84; + const auto x129 = x122*x4; + const auto x130 = x122*x4*x8; + const auto x131 = 6.0*x1 + x11*x122; + const auto x132 = x11*x72 + x131; + const auto x133 = x122*x5 + x77; + const auto x134 = x122*x8; + const auto x135 = x122*x14; + const auto x136 = x122*x15 + x15*x69 + x5*x68; + const auto x137 = x16*x75; + const auto x138 = x122*x5*x8 + x82; + const auto x139 = x122*x17; + const auto x140 = x16*x97; + const auto x141 = radial_eval_alpha + radial_eval_alpha_squared*x8; + const auto x142 = x0*x141; + const auto x143 = x141*x2; + const auto x144 = x141*x4*x5; + const auto x145 = x141*x4; + const auto x146 = x126 + x141*x4*x8; + const auto x147 = x11*x141; + const auto x148 = x141*x5; + const auto x149 = x141*x8 + x77; + const auto x150 = x14*x141 + 6.0*x3; + const auto x151 = x14*x72 + x150; + const auto x152 = x141*x15; + const auto x153 = x141*x5*x8 + x78; + const auto x154 = x141*x17 + x17*x69 + x68*x8; + const auto x155 = x125 + x143 + x2*x69 + x73; + const auto x156 = x11*x69 + x131 + x147 + x86; + const auto x157 = x135 + x14*x69 + x150 + x92; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x2; + basis_eval[ipt + 2*npts] = x2*x3; + basis_eval[ipt + 3*npts] = radial_eval*x6; + basis_eval[ipt + 4*npts] = x4*x7; + basis_eval[ipt + 5*npts] = radial_eval*x9; + basis_eval[ipt + 6*npts] = x10*x11; + basis_eval[ipt + 7*npts] = x12*x5; + basis_eval[ipt + 8*npts] = x13*x8; + basis_eval[ipt + 9*npts] = x10*x14; + basis_eval[ipt + 10*npts] = radial_eval*x15; + basis_eval[ipt + 11*npts] = x11*x3; + basis_eval[ipt + 12*npts] = radial_eval*x16; + basis_eval[ipt + 13*npts] = x1*x14; + basis_eval[ipt + 14*npts] = radial_eval*x17; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x18 + x19*x2; + basis_x_eval[ipt + 1*npts] = x21*y; + basis_x_eval[ipt + 2*npts] = x21*z; + basis_x_eval[ipt + 3*npts] = x22*x5 + x24; + basis_x_eval[ipt + 4*npts] = x25*x27; + basis_x_eval[ipt + 5*npts] = x22*x8 + x29; + basis_x_eval[ipt + 6*npts] = x30 + x32; + basis_x_eval[ipt + 7*npts] = x35*z; + basis_x_eval[ipt + 8*npts] = x38*y; + basis_x_eval[ipt + 9*npts] = x39 + x41; + basis_x_eval[ipt + 10*npts] = x15*x42; + basis_x_eval[ipt + 11*npts] = x43; + basis_x_eval[ipt + 12*npts] = x16*x42; + basis_x_eval[ipt + 13*npts] = x44; + basis_x_eval[ipt + 14*npts] = x17*x42; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x45; + basis_y_eval[ipt + 1*npts] = x24 + x46; + basis_y_eval[ipt + 2*npts] = x47; + basis_y_eval[ipt + 3*npts] = x32 + x4*x48; + basis_y_eval[ipt + 4*npts] = x50*z; + basis_y_eval[ipt + 5*npts] = x45*x9; + basis_y_eval[ipt + 6*npts] = x*x51; + basis_y_eval[ipt + 7*npts] = x52*x54; + basis_y_eval[ipt + 8*npts] = x*(x36 + x55); + basis_y_eval[ipt + 9*npts] = x44; + basis_y_eval[ipt + 10*npts] = radial_eval_alpha*x56 + x11*x19; + basis_y_eval[ipt + 11*npts] = x51*z; + basis_y_eval[ipt + 12*npts] = x48*x8 + x58; + basis_y_eval[ipt + 13*npts] = x39 + x60; + basis_y_eval[ipt + 14*npts] = x17*x45; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x61; + basis_z_eval[ipt + 1*npts] = x47; + basis_z_eval[ipt + 2*npts] = x29 + x46; + basis_z_eval[ipt + 3*npts] = x6*x61; + basis_z_eval[ipt + 4*npts] = y*(x37 + x49); + basis_z_eval[ipt + 5*npts] = x4*x62 + x41; + basis_z_eval[ipt + 6*npts] = x43; + basis_z_eval[ipt + 7*npts] = x*(x33 + x55); + basis_z_eval[ipt + 8*npts] = x63*x65; + basis_z_eval[ipt + 9*npts] = x*x66; + basis_z_eval[ipt + 10*npts] = x15*x61; + basis_z_eval[ipt + 11*npts] = x30 + x58; + basis_z_eval[ipt + 12*npts] = x5*x62 + x60; + basis_z_eval[ipt + 13*npts] = x66*y; + basis_z_eval[ipt + 14*npts] = radial_eval_alpha*x67 + x14*x19; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x71; + basis_xx_eval[ipt + 1*npts] = x74*y; + basis_xx_eval[ipt + 2*npts] = x74*z; + basis_xx_eval[ipt + 3*npts] = x76 + x79; + basis_xx_eval[ipt + 4*npts] = x25*(x4*x75 + x80); + basis_xx_eval[ipt + 5*npts] = x81 + x83; + basis_xx_eval[ipt + 6*npts] = x*(x85 + x86); + basis_xx_eval[ipt + 7*npts] = x52*(x87 + x88); + basis_xx_eval[ipt + 8*npts] = x63*(x89 + x90); + basis_xx_eval[ipt + 9*npts] = x*(x91 + x92); + basis_xx_eval[ipt + 10*npts] = x93; + basis_xx_eval[ipt + 11*npts] = x86*z; + basis_xx_eval[ipt + 12*npts] = x94; + basis_xx_eval[ipt + 13*npts] = x92*y; + basis_xx_eval[ipt + 14*npts] = x95; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x96*y; + basis_xy_eval[ipt + 1*npts] = radial_eval_alpha_squared*x0*x5 + x21 + x98; + basis_xy_eval[ipt + 2*npts] = x99; + basis_xy_eval[ipt + 3*npts] = radial_eval_alpha_squared*x11*x2 + x100*x11 + x101*x2 + 4.0*x13; + basis_xy_eval[ipt + 4*npts] = z*(x103 + x27); + basis_xy_eval[ipt + 5*npts] = x105*y; + basis_xy_eval[ipt + 6*npts] = radial_eval_alpha_squared*x15*x4 + x51 + x98; + basis_xy_eval[ipt + 7*npts] = z*(x107 + x54); + basis_xy_eval[ipt + 8*npts] = x109 + x38; + basis_xy_eval[ipt + 9*npts] = y*(x110 + x64); + basis_xy_eval[ipt + 10*npts] = x*x111; + basis_xy_eval[ipt + 11*npts] = x112; + basis_xy_eval[ipt + 12*npts] = x*x114; + basis_xy_eval[ipt + 13*npts] = x*(x115 + x64); + basis_xy_eval[ipt + 14*npts] = radial_eval_alpha_squared*x17*x63; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x96*z; + basis_xz_eval[ipt + 1*npts] = x99; + basis_xz_eval[ipt + 2*npts] = radial_eval_alpha_squared*x0*x8 + x116 + x21; + basis_xz_eval[ipt + 3*npts] = x103*z; + basis_xz_eval[ipt + 4*npts] = y*(x105 + x27); + basis_xz_eval[ipt + 5*npts] = radial_eval_alpha_squared*x14*x2 + x100*x14 + x117*x2 + 4.0*x12; + basis_xz_eval[ipt + 6*npts] = z*(x106 + x53); + basis_xz_eval[ipt + 7*npts] = x109 + x35; + basis_xz_eval[ipt + 8*npts] = y*(x118 + x65); + basis_xz_eval[ipt + 9*npts] = radial_eval_alpha_squared*x17*x4 + x116 + x66; + basis_xz_eval[ipt + 10*npts] = radial_eval_alpha_squared*x15*x52; + basis_xz_eval[ipt + 11*npts] = x*(x113 + x53); + basis_xz_eval[ipt + 12*npts] = x*x119; + basis_xz_eval[ipt + 13*npts] = x120; + basis_xz_eval[ipt + 14*npts] = x*x121; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x123; + basis_yy_eval[ipt + 1*npts] = y*(x124 + x125); + basis_yy_eval[ipt + 2*npts] = x125*z; + basis_yy_eval[ipt + 3*npts] = x127 + x76; + basis_yy_eval[ipt + 4*npts] = x25*(x128 + x129); + basis_yy_eval[ipt + 5*npts] = x130; + basis_yy_eval[ipt + 6*npts] = x*x132; + basis_yy_eval[ipt + 7*npts] = x52*(x133 + x5*x75); + basis_yy_eval[ipt + 8*npts] = x63*(x134 + x89); + basis_yy_eval[ipt + 9*npts] = x*x135; + basis_yy_eval[ipt + 10*npts] = x136; + basis_yy_eval[ipt + 11*npts] = x132*z; + basis_yy_eval[ipt + 12*npts] = x137 + x138; + basis_yy_eval[ipt + 13*npts] = y*(x135 + x91); + basis_yy_eval[ipt + 14*npts] = x139; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*x0*x25; + basis_yz_eval[ipt + 1*npts] = z*(x102 + x26); + basis_yz_eval[ipt + 2*npts] = y*(x104 + x26); + basis_yz_eval[ipt + 3*npts] = x107*z; + basis_yz_eval[ipt + 4*npts] = x108 + x37 + x50; + basis_yz_eval[ipt + 5*npts] = x118*y; + basis_yz_eval[ipt + 6*npts] = x112; + basis_yz_eval[ipt + 7*npts] = x*(x114 + x54); + basis_yz_eval[ipt + 8*npts] = x*(x119 + x65); + basis_yz_eval[ipt + 9*npts] = x120; + basis_yz_eval[ipt + 10*npts] = x111*z; + basis_yz_eval[ipt + 11*npts] = radial_eval_alpha_squared*x15*x8 + x140 + x51; + basis_yz_eval[ipt + 12*npts] = radial_eval_alpha_squared*x11*x14 + x101*x14 + x11*x117 + 4.0*x7; + basis_yz_eval[ipt + 13*npts] = radial_eval_alpha_squared*x17*x5 + x140 + x66; + basis_yz_eval[ipt + 14*npts] = x121*y; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x142; + basis_zz_eval[ipt + 1*npts] = x143*y; + basis_zz_eval[ipt + 2*npts] = z*(x124 + x143); + basis_zz_eval[ipt + 3*npts] = x144; + basis_zz_eval[ipt + 4*npts] = x25*(x128 + x145); + basis_zz_eval[ipt + 5*npts] = x146 + x81; + basis_zz_eval[ipt + 6*npts] = x*x147; + basis_zz_eval[ipt + 7*npts] = x52*(x148 + x87); + basis_zz_eval[ipt + 8*npts] = x63*(x149 + x75*x8); + basis_zz_eval[ipt + 9*npts] = x*x151; + basis_zz_eval[ipt + 10*npts] = x152; + basis_zz_eval[ipt + 11*npts] = z*(x147 + x85); + basis_zz_eval[ipt + 12*npts] = x137 + x153; + basis_zz_eval[ipt + 13*npts] = x151*y; + basis_zz_eval[ipt + 14*npts] = x154; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x2; + ang_eval_2 = x2*x3; + ang_eval_3 = radial_eval*x6; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x4*x7; + ang_eval_1 = radial_eval*x9; + ang_eval_2 = x10*x11; + ang_eval_3 = x12*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x13*x8; + ang_eval_1 = x10*x14; + ang_eval_2 = radial_eval*x15; + ang_eval_3 = x11*x3; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + basis_eval[ipt + 10*npts] = ang_eval_2; + basis_eval[ipt + 11*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x16; + ang_eval_1 = x1*x14; + ang_eval_2 = radial_eval*x17; + basis_eval[ipt + 12*npts] = ang_eval_0; + basis_eval[ipt + 13*npts] = ang_eval_1; + basis_eval[ipt + 14*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x18 + x19*x2; + dang_eval_y_0 = x0*x45; + dang_eval_z_0 = x0*x61; + dang_eval_x_1 = x21*y; + dang_eval_y_1 = x24 + x46; + dang_eval_z_1 = x47; + dang_eval_x_2 = x21*z; + dang_eval_y_2 = x47; + dang_eval_z_2 = x29 + x46; + dang_eval_x_3 = x22*x5 + x24; + dang_eval_y_3 = x32 + x4*x48; + dang_eval_z_3 = x6*x61; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x25*x27; + dang_eval_y_0 = x50*z; + dang_eval_z_0 = y*(x37 + x49); + dang_eval_x_1 = x22*x8 + x29; + dang_eval_y_1 = x45*x9; + dang_eval_z_1 = x4*x62 + x41; + dang_eval_x_2 = x30 + x32; + dang_eval_y_2 = x*x51; + dang_eval_z_2 = x43; + dang_eval_x_3 = x35*z; + dang_eval_y_3 = x52*x54; + dang_eval_z_3 = x*(x33 + x55); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x38*y; + dang_eval_y_0 = x*(x36 + x55); + dang_eval_z_0 = x63*x65; + dang_eval_x_1 = x39 + x41; + dang_eval_y_1 = x44; + dang_eval_z_1 = x*x66; + dang_eval_x_2 = x15*x42; + dang_eval_y_2 = radial_eval_alpha*x56 + x11*x19; + dang_eval_z_2 = x15*x61; + dang_eval_x_3 = x43; + dang_eval_y_3 = x51*z; + dang_eval_z_3 = x30 + x58; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + basis_x_eval[ipt + 10*npts] = dang_eval_x_2; + basis_y_eval[ipt + 10*npts] = dang_eval_y_2; + basis_z_eval[ipt + 10*npts] = dang_eval_z_2; + basis_x_eval[ipt + 11*npts] = dang_eval_x_3; + basis_y_eval[ipt + 11*npts] = dang_eval_y_3; + basis_z_eval[ipt + 11*npts] = dang_eval_z_3; + + dang_eval_x_0 = x16*x42; + dang_eval_y_0 = x48*x8 + x58; + dang_eval_z_0 = x5*x62 + x60; + dang_eval_x_1 = x44; + dang_eval_y_1 = x39 + x60; + dang_eval_z_1 = x66*y; + dang_eval_x_2 = x17*x42; + dang_eval_y_2 = x17*x45; + dang_eval_z_2 = radial_eval_alpha*x67 + x14*x19; + basis_x_eval[ipt + 12*npts] = dang_eval_x_0; + basis_y_eval[ipt + 12*npts] = dang_eval_y_0; + basis_z_eval[ipt + 12*npts] = dang_eval_z_0; + basis_x_eval[ipt + 13*npts] = dang_eval_x_1; + basis_y_eval[ipt + 13*npts] = dang_eval_y_1; + basis_z_eval[ipt + 13*npts] = dang_eval_z_1; + basis_x_eval[ipt + 14*npts] = dang_eval_x_2; + basis_y_eval[ipt + 14*npts] = dang_eval_y_2; + basis_z_eval[ipt + 14*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_lapgrad.hpp new file mode 100644 index 000000000..c7f314162 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_lapgrad.hpp @@ -0,0 +1,796 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_lapgrad_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_lapgrad_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = x*x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x*x; + const auto x3 = radial_eval*z; + const auto x4 = x*x; + const auto x5 = x4; + const auto x6 = y*y; + const auto x7 = x6; + const auto x8 = x5*x7; + const auto x9 = x1*z; + const auto x10 = z*z; + const auto x11 = x10; + const auto x12 = x11*x5; + const auto x13 = radial_eval*x; + const auto x14 = y*y*y; + const auto x15 = x*x3; + const auto x16 = x*x1; + const auto x17 = z*z*z; + const auto x18 = y*y*y*y; + const auto x19 = x11*x7; + const auto x20 = z*z*z*z; + const auto x21 = x*x*x*x*x; + const auto x22 = 4.0*radial_eval; + const auto x23 = 3.0*radial_eval; + const auto x24 = radial_eval_alpha*x0 + x23*x5; + const auto x25 = 2.0*x13; + const auto x26 = x2*x7; + const auto x27 = radial_eval_alpha*x26; + const auto x28 = y*z; + const auto x29 = radial_eval_alpha*x2; + const auto x30 = x25 + x29; + const auto x31 = x11*x2; + const auto x32 = radial_eval_alpha*x31; + const auto x33 = radial_eval*x14; + const auto x34 = x14*x5; + const auto x35 = radial_eval_alpha*x34; + const auto x36 = radial_eval*x7; + const auto x37 = radial_eval_alpha*x8; + const auto x38 = x36 + x37; + const auto x39 = radial_eval*x11; + const auto x40 = radial_eval_alpha*x12; + const auto x41 = x39 + x40; + const auto x42 = radial_eval*x17; + const auto x43 = x17*x5; + const auto x44 = radial_eval_alpha*x43; + const auto x45 = radial_eval_alpha*x; + const auto x46 = x14*x45*z; + const auto x47 = x17*x45*y; + const auto x48 = radial_eval_alpha*y; + const auto x49 = radial_eval*x2; + const auto x50 = radial_eval_alpha*x2*x28; + const auto x51 = 2.0*x1; + const auto x52 = radial_eval*x5; + const auto x53 = x37 + x52; + const auto x54 = radial_eval_alpha*x18 + x23*x7; + const auto x55 = x*z; + const auto x56 = radial_eval_alpha*x14; + const auto x57 = x51 + x56; + const auto x58 = radial_eval_alpha*x19; + const auto x59 = y*y*y*y*y; + const auto x60 = x11*x14; + const auto x61 = radial_eval_alpha*x60; + const auto x62 = x17*x7; + const auto x63 = radial_eval_alpha*x62; + const auto x64 = radial_eval_alpha*z; + const auto x65 = 2.0*x3; + const auto x66 = x*y; + const auto x67 = radial_eval_alpha*x17; + const auto x68 = x65 + x67; + const auto x69 = radial_eval_alpha*x20 + x11*x23; + const auto x70 = z*z*z*z*z; + const auto x71 = 12.0*radial_eval; + const auto x72 = 8.0*radial_eval_alpha; + const auto x73 = radial_eval_alpha + radial_eval_alpha_squared*x5; + const auto x74 = x0*x72 + x0*x73 + x5*x71; + const auto x75 = 6.0*radial_eval_alpha; + const auto x76 = x2*x73; + const auto x77 = 6.0*x13 + x76; + const auto x78 = x2*x75 + x77; + const auto x79 = 4.0*radial_eval_alpha; + const auto x80 = x79*x8; + const auto x81 = 2.0*radial_eval; + const auto x82 = x7*x81; + const auto x83 = x5*x7*x73 + x82; + const auto x84 = x5*x73; + const auto x85 = x81 + x84; + const auto x86 = x12*x79; + const auto x87 = x11*x81; + const auto x88 = x11*x5*x73 + x87; + const auto x89 = 2.0*radial_eval_alpha; + const auto x90 = x14*x89; + const auto x91 = x14*x73; + const auto x92 = x7*x89; + const auto x93 = x7*x73; + const auto x94 = x11*x89; + const auto x95 = x11*x73; + const auto x96 = x17*x89; + const auto x97 = x17*x73; + const auto x98 = x18*x73; + const auto x99 = x11*x7*x73; + const auto x100 = x20*x73; + const auto x101 = radial_eval_alpha_squared*x21 + x2*x79; + const auto x102 = 3.0*radial_eval_alpha; + const auto x103 = x102*x8; + const auto x104 = x28*(radial_eval_alpha_squared*x0 + x102*x5); + const auto x105 = 2.0*x45; + const auto x106 = 2.0*x48; + const auto x107 = x105*x7; + const auto x108 = radial_eval_alpha_squared*x26; + const auto x109 = x107 + x108; + const auto x110 = x105*x11; + const auto x111 = radial_eval_alpha_squared*x31; + const auto x112 = x110 + x111; + const auto x113 = x106*x5; + const auto x114 = radial_eval_alpha_squared*x34; + const auto x115 = x113 + x114; + const auto x116 = radial_eval_alpha_squared*x11*x5*x7; + const auto x117 = x116 + x58; + const auto x118 = radial_eval_alpha_squared*x43; + const auto x119 = radial_eval_alpha_squared*x59 + x14*x79; + const auto x120 = x55*(radial_eval_alpha_squared*x18 + x102*x7); + const auto x121 = x106*x11; + const auto x122 = radial_eval_alpha_squared*x60; + const auto x123 = x121 + x122; + const auto x124 = radial_eval_alpha_squared*x62; + const auto x125 = x102*x12; + const auto x126 = 2.0*x64; + const auto x127 = x126*x5; + const auto x128 = x118 + x127; + const auto x129 = x126*x7; + const auto x130 = x124 + x129; + const auto x131 = x66*(radial_eval_alpha_squared*x20 + x102*x11); + const auto x132 = radial_eval_alpha_squared*x70 + x17*x79; + const auto x133 = radial_eval_alpha + radial_eval_alpha_squared*x7; + const auto x134 = x0*x133; + const auto x135 = x2*x89; + const auto x136 = x133*x2; + const auto x137 = x5*x81; + const auto x138 = x133*x5*x7 + x137; + const auto x139 = x5*x89; + const auto x140 = x133*x5; + const auto x141 = x11*x133*x5; + const auto x142 = x133*x14; + const auto x143 = 6.0*x1 + x142; + const auto x144 = x14*x75 + x143; + const auto x145 = x133*x7; + const auto x146 = x145 + x81; + const auto x147 = x11*x133; + const auto x148 = x133*x17; + const auto x149 = x133*x18 + x18*x72 + x7*x71; + const auto x150 = x19*x79; + const auto x151 = x11*x133*x7 + x87; + const auto x152 = x133*x20; + const auto x153 = x102*x19; + const auto x154 = radial_eval_alpha + radial_eval_alpha_squared*x11; + const auto x155 = x0*x154; + const auto x156 = x154*x2; + const auto x157 = x154*x5*x7; + const auto x158 = x154*x5; + const auto x159 = x11*x154*x5 + x137; + const auto x160 = x14*x154; + const auto x161 = x154*x7; + const auto x162 = x11*x154; + const auto x163 = x162 + x81; + const auto x164 = x154*x17; + const auto x165 = x164 + 6.0*x3; + const auto x166 = x165 + x17*x75; + const auto x167 = x154*x18; + const auto x168 = x11*x154*x7 + x82; + const auto x169 = x11*x71 + x154*x20 + x20*x72; + const auto x170 = x136 + x156 + x2*x72 + x77; + const auto x171 = x158 + x85; + const auto x172 = x14*x72 + x143 + x160 + x91; + const auto x173 = x146 + x161; + const auto x174 = x147 + x163; + const auto x175 = x148 + x165 + x17*x72 + x97; + const auto x176 = 36.0*radial_eval_alpha; + const auto x177 = radial_eval_alpha_cubed*x7 + radial_eval_alpha_squared; + const auto x178 = x0*x177; + const auto x179 = radial_eval_alpha_cubed*x11 + radial_eval_alpha_squared; + const auto x180 = x0*x179; + const auto x181 = radial_eval_alpha_squared*x; + const auto x182 = radial_eval_alpha_cubed*x2 + 3.0*x181; + const auto x183 = 6.0*radial_eval; + const auto x184 = 24.0*radial_eval_alpha; + const auto x185 = 2.0*radial_eval_alpha_squared; + const auto x186 = 3.0*x140; + const auto x187 = 3.0*x158; + const auto x188 = x177*x2; + const auto x189 = x179*x2; + const auto x190 = x*x188 + x*x189 + x0*x185 + x182*x2 + x183 + x184*x5 + x186 + x187 + 9.0*x84; + const auto x191 = 2.0*x; + const auto x192 = 4.0*radial_eval_alpha_squared; + const auto x193 = 6.0*x; + const auto x194 = 14.0*x45; + const auto x195 = x177*x5*x7; + const auto x196 = x179*x5*x7; + const auto x197 = 4.0*x13 + x135; + const auto x198 = x177*x5; + const auto x199 = x179*x5; + const auto x200 = x11*x177*x5; + const auto x201 = x11*x179*x5; + const auto x202 = x14*x182; + const auto x203 = x14*x177; + const auto x204 = x14*x179; + const auto x205 = 6.0*x48; + const auto x206 = 6.0*radial_eval_alpha_squared; + const auto x207 = 3.0*x93; + const auto x208 = x7*x75; + const auto x209 = x177*x7; + const auto x210 = x179*x7; + const auto x211 = x206*x8; + const auto x212 = 3.0*x95; + const auto x213 = x11*x75; + const auto x214 = x11*x177; + const auto x215 = x11*x179; + const auto x216 = x12*x206; + const auto x217 = x17*x182; + const auto x218 = x17*x177; + const auto x219 = x17*x179; + const auto x220 = 6.0*x64; + const auto x221 = 12.0*x45; + const auto x222 = 8.0*x181; + const auto x223 = x177*x18; + const auto x224 = x179*x18; + const auto x225 = 6.0*y; + const auto x226 = x225*x45; + const auto x227 = x11*x177*x7; + const auto x228 = x11*x179*x7; + const auto x229 = 6.0*z; + const auto x230 = x229*x45; + const auto x231 = x177*x20; + const auto x232 = x179*x20; + const auto x233 = 12.0*x48; + const auto x234 = radial_eval_alpha_squared*y; + const auto x235 = 8.0*x234; + const auto x236 = radial_eval_alpha_cubed*x5 + radial_eval_alpha_squared; + const auto x237 = x0*x236; + const auto x238 = radial_eval_alpha_cubed*x14 + 3.0*x234; + const auto x239 = x2*x238; + const auto x240 = x2*x236; + const auto x241 = 6.0*x45; + const auto x242 = 2.0*y; + const auto x243 = 14.0*x48; + const auto x244 = x236*x5*x7; + const auto x245 = 4.0*x1 + x90; + const auto x246 = x5*x75; + const auto x247 = x236*x5; + const auto x248 = x11*x236*x5; + const auto x249 = 3.0*x161; + const auto x250 = x14*x236; + const auto x251 = x14*x238 + 9.0*x145 + x18*x185 + x183 + x184*x7 + x204*y + x207 + x249 + x250*y; + const auto x252 = x236*x7; + const auto x253 = 3.0*x147; + const auto x254 = x11*x236; + const auto x255 = x19*x206; + const auto x256 = x28*x75; + const auto x257 = x17*x236; + const auto x258 = x17*x238; + const auto x259 = x18*x236; + const auto x260 = x11*x236*x7; + const auto x261 = x20*x236; + const auto x262 = 12.0*x64; + const auto x263 = radial_eval_alpha_squared*z; + const auto x264 = 8.0*x263; + const auto x265 = radial_eval_alpha_cubed*x17 + 3.0*x263; + const auto x266 = x2*x265; + const auto x267 = 2.0*z; + const auto x268 = 14.0*x64; + const auto x269 = 4.0*x3 + x96; + const auto x270 = x14*x265; + const auto x271 = x11*x184 + 9.0*x162 + x17*x265 + x183 + x185*x20 + x212 + x218*z + x253 + x257*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x2; + basis_eval[ipt + 2*npts] = x2*x3; + basis_eval[ipt + 3*npts] = radial_eval*x8; + basis_eval[ipt + 4*npts] = x5*x9; + basis_eval[ipt + 5*npts] = radial_eval*x12; + basis_eval[ipt + 6*npts] = x13*x14; + basis_eval[ipt + 7*npts] = x15*x7; + basis_eval[ipt + 8*npts] = x11*x16; + basis_eval[ipt + 9*npts] = x13*x17; + basis_eval[ipt + 10*npts] = radial_eval*x18; + basis_eval[ipt + 11*npts] = x14*x3; + basis_eval[ipt + 12*npts] = radial_eval*x19; + basis_eval[ipt + 13*npts] = x1*x17; + basis_eval[ipt + 14*npts] = radial_eval*x20; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x21 + x2*x22; + basis_x_eval[ipt + 1*npts] = x24*y; + basis_x_eval[ipt + 2*npts] = x24*z; + basis_x_eval[ipt + 3*npts] = x25*x7 + x27; + basis_x_eval[ipt + 4*npts] = x28*x30; + basis_x_eval[ipt + 5*npts] = x11*x25 + x32; + basis_x_eval[ipt + 6*npts] = x33 + x35; + basis_x_eval[ipt + 7*npts] = x38*z; + basis_x_eval[ipt + 8*npts] = x41*y; + basis_x_eval[ipt + 9*npts] = x42 + x44; + basis_x_eval[ipt + 10*npts] = x18*x45; + basis_x_eval[ipt + 11*npts] = x46; + basis_x_eval[ipt + 12*npts] = x19*x45; + basis_x_eval[ipt + 13*npts] = x47; + basis_x_eval[ipt + 14*npts] = x20*x45; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x48; + basis_y_eval[ipt + 1*npts] = x27 + x49; + basis_y_eval[ipt + 2*npts] = x50; + basis_y_eval[ipt + 3*npts] = x35 + x5*x51; + basis_y_eval[ipt + 4*npts] = x53*z; + basis_y_eval[ipt + 5*npts] = x12*x48; + basis_y_eval[ipt + 6*npts] = x*x54; + basis_y_eval[ipt + 7*npts] = x55*x57; + basis_y_eval[ipt + 8*npts] = x*(x39 + x58); + basis_y_eval[ipt + 9*npts] = x47; + basis_y_eval[ipt + 10*npts] = radial_eval_alpha*x59 + x14*x22; + basis_y_eval[ipt + 11*npts] = x54*z; + basis_y_eval[ipt + 12*npts] = x11*x51 + x61; + basis_y_eval[ipt + 13*npts] = x42 + x63; + basis_y_eval[ipt + 14*npts] = x20*x48; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x64; + basis_z_eval[ipt + 1*npts] = x50; + basis_z_eval[ipt + 2*npts] = x32 + x49; + basis_z_eval[ipt + 3*npts] = x64*x8; + basis_z_eval[ipt + 4*npts] = y*(x40 + x52); + basis_z_eval[ipt + 5*npts] = x44 + x5*x65; + basis_z_eval[ipt + 6*npts] = x46; + basis_z_eval[ipt + 7*npts] = x*(x36 + x58); + basis_z_eval[ipt + 8*npts] = x66*x68; + basis_z_eval[ipt + 9*npts] = x*x69; + basis_z_eval[ipt + 10*npts] = x18*x64; + basis_z_eval[ipt + 11*npts] = x33 + x61; + basis_z_eval[ipt + 12*npts] = x63 + x65*x7; + basis_z_eval[ipt + 13*npts] = x69*y; + basis_z_eval[ipt + 14*npts] = radial_eval_alpha*x70 + x17*x22; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x74; + basis_xx_eval[ipt + 1*npts] = x78*y; + basis_xx_eval[ipt + 2*npts] = x78*z; + basis_xx_eval[ipt + 3*npts] = x80 + x83; + basis_xx_eval[ipt + 4*npts] = x28*(x5*x79 + x85); + basis_xx_eval[ipt + 5*npts] = x86 + x88; + basis_xx_eval[ipt + 6*npts] = x*(x90 + x91); + basis_xx_eval[ipt + 7*npts] = x55*(x92 + x93); + basis_xx_eval[ipt + 8*npts] = x66*(x94 + x95); + basis_xx_eval[ipt + 9*npts] = x*(x96 + x97); + basis_xx_eval[ipt + 10*npts] = x98; + basis_xx_eval[ipt + 11*npts] = x91*z; + basis_xx_eval[ipt + 12*npts] = x99; + basis_xx_eval[ipt + 13*npts] = x97*y; + basis_xx_eval[ipt + 14*npts] = x100; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x101*y; + basis_xy_eval[ipt + 1*npts] = radial_eval_alpha_squared*x0*x7 + x103 + x24; + basis_xy_eval[ipt + 2*npts] = x104; + basis_xy_eval[ipt + 3*npts] = radial_eval_alpha_squared*x14*x2 + x105*x14 + x106*x2 + 4.0*x16; + basis_xy_eval[ipt + 4*npts] = z*(x109 + x30); + basis_xy_eval[ipt + 5*npts] = x112*y; + basis_xy_eval[ipt + 6*npts] = radial_eval_alpha_squared*x18*x5 + x103 + x54; + basis_xy_eval[ipt + 7*npts] = z*(x115 + x57); + basis_xy_eval[ipt + 8*npts] = x117 + x41; + basis_xy_eval[ipt + 9*npts] = y*(x118 + x67); + basis_xy_eval[ipt + 10*npts] = x*x119; + basis_xy_eval[ipt + 11*npts] = x120; + basis_xy_eval[ipt + 12*npts] = x*x123; + basis_xy_eval[ipt + 13*npts] = x*(x124 + x67); + basis_xy_eval[ipt + 14*npts] = radial_eval_alpha_squared*x20*x66; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x101*z; + basis_xz_eval[ipt + 1*npts] = x104; + basis_xz_eval[ipt + 2*npts] = radial_eval_alpha_squared*x0*x11 + x125 + x24; + basis_xz_eval[ipt + 3*npts] = x109*z; + basis_xz_eval[ipt + 4*npts] = y*(x112 + x30); + basis_xz_eval[ipt + 5*npts] = radial_eval_alpha_squared*x17*x2 + x105*x17 + x126*x2 + 4.0*x15; + basis_xz_eval[ipt + 6*npts] = z*(x114 + x56); + basis_xz_eval[ipt + 7*npts] = x117 + x38; + basis_xz_eval[ipt + 8*npts] = y*(x128 + x68); + basis_xz_eval[ipt + 9*npts] = radial_eval_alpha_squared*x20*x5 + x125 + x69; + basis_xz_eval[ipt + 10*npts] = radial_eval_alpha_squared*x18*x55; + basis_xz_eval[ipt + 11*npts] = x*(x122 + x56); + basis_xz_eval[ipt + 12*npts] = x*x130; + basis_xz_eval[ipt + 13*npts] = x131; + basis_xz_eval[ipt + 14*npts] = x*x132; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x134; + basis_yy_eval[ipt + 1*npts] = y*(x135 + x136); + basis_yy_eval[ipt + 2*npts] = x136*z; + basis_yy_eval[ipt + 3*npts] = x138 + x80; + basis_yy_eval[ipt + 4*npts] = x28*(x139 + x140); + basis_yy_eval[ipt + 5*npts] = x141; + basis_yy_eval[ipt + 6*npts] = x*x144; + basis_yy_eval[ipt + 7*npts] = x55*(x146 + x7*x79); + basis_yy_eval[ipt + 8*npts] = x66*(x147 + x94); + basis_yy_eval[ipt + 9*npts] = x*x148; + basis_yy_eval[ipt + 10*npts] = x149; + basis_yy_eval[ipt + 11*npts] = x144*z; + basis_yy_eval[ipt + 12*npts] = x150 + x151; + basis_yy_eval[ipt + 13*npts] = y*(x148 + x96); + basis_yy_eval[ipt + 14*npts] = x152; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*x0*x28; + basis_yz_eval[ipt + 1*npts] = z*(x108 + x29); + basis_yz_eval[ipt + 2*npts] = y*(x111 + x29); + basis_yz_eval[ipt + 3*npts] = x115*z; + basis_yz_eval[ipt + 4*npts] = x116 + x40 + x53; + basis_yz_eval[ipt + 5*npts] = x128*y; + basis_yz_eval[ipt + 6*npts] = x120; + basis_yz_eval[ipt + 7*npts] = x*(x123 + x57); + basis_yz_eval[ipt + 8*npts] = x*(x130 + x68); + basis_yz_eval[ipt + 9*npts] = x131; + basis_yz_eval[ipt + 10*npts] = x119*z; + basis_yz_eval[ipt + 11*npts] = radial_eval_alpha_squared*x11*x18 + x153 + x54; + basis_yz_eval[ipt + 12*npts] = radial_eval_alpha_squared*x14*x17 + x106*x17 + x126*x14 + 4.0*x9; + basis_yz_eval[ipt + 13*npts] = radial_eval_alpha_squared*x20*x7 + x153 + x69; + basis_yz_eval[ipt + 14*npts] = x132*y; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x155; + basis_zz_eval[ipt + 1*npts] = x156*y; + basis_zz_eval[ipt + 2*npts] = z*(x135 + x156); + basis_zz_eval[ipt + 3*npts] = x157; + basis_zz_eval[ipt + 4*npts] = x28*(x139 + x158); + basis_zz_eval[ipt + 5*npts] = x159 + x86; + basis_zz_eval[ipt + 6*npts] = x*x160; + basis_zz_eval[ipt + 7*npts] = x55*(x161 + x92); + basis_zz_eval[ipt + 8*npts] = x66*(x11*x79 + x163); + basis_zz_eval[ipt + 9*npts] = x*x166; + basis_zz_eval[ipt + 10*npts] = x167; + basis_zz_eval[ipt + 11*npts] = z*(x160 + x90); + basis_zz_eval[ipt + 12*npts] = x150 + x168; + basis_zz_eval[ipt + 13*npts] = x166*y; + basis_zz_eval[ipt + 14*npts] = x169; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x134 + x155 + x74; + basis_lapl_eval[ipt + 1*npts] = x170*y; + basis_lapl_eval[ipt + 2*npts] = x170*z; + basis_lapl_eval[ipt + 3*npts] = x138 + x157 + x72*x8 + x83; + basis_lapl_eval[ipt + 4*npts] = x28*(x140 + x171 + x5*x72); + basis_lapl_eval[ipt + 5*npts] = x12*x72 + x141 + x159 + x88; + basis_lapl_eval[ipt + 6*npts] = x*x172; + basis_lapl_eval[ipt + 7*npts] = x55*(x173 + x7*x72 + x93); + basis_lapl_eval[ipt + 8*npts] = x66*(x11*x72 + x174 + x95); + basis_lapl_eval[ipt + 9*npts] = x*x175; + basis_lapl_eval[ipt + 10*npts] = x149 + x167 + x98; + basis_lapl_eval[ipt + 11*npts] = x172*z; + basis_lapl_eval[ipt + 12*npts] = x151 + x168 + x19*x72 + x99; + basis_lapl_eval[ipt + 13*npts] = x175*y; + basis_lapl_eval[ipt + 14*npts] = x100 + x152 + x169; + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = x*x178 + x*x180 + x0*x182 + 24.0*x13 + 4.0*x136 + 4.0*x156 + x176*x2 + 12.0*x76; + basis_lapl_x_eval[ipt + 1*npts] = x190*y; + basis_lapl_x_eval[ipt + 2*npts] = x190*z; + basis_lapl_x_eval[ipt + 3*npts] = x*x195 + x*x196 + x145*x191 + x161*x191 + x182*x5*x7 + x192*x26 + x193*x93 + x194*x7 + x197; + basis_lapl_x_eval[ipt + 4*npts] = x28*(x*x198 + x*x199 + x133*x191 + x154*x191 + x182*x5 + x192*x2 + x193*x73 + x194); + basis_lapl_x_eval[ipt + 5*npts] = x*x200 + x*x201 + x11*x182*x5 + x11*x194 + x147*x191 + x162*x191 + x192*x31 + x193*x95 + x197; + basis_lapl_x_eval[ipt + 6*npts] = x*x202 + x144 + x160 + x203*x4 + x204*x4 + x205*x5 + x206*x34 + 3.0*x91; + basis_lapl_x_eval[ipt + 7*npts] = z*(x*x182*x7 + x139 + x173 + x207 + x208 + x209*x4 + x210*x4 + x211); + basis_lapl_x_eval[ipt + 8*npts] = y*(x*x11*x182 + x139 + x174 + x212 + x213 + x214*x4 + x215*x4 + x216); + basis_lapl_x_eval[ipt + 9*npts] = x*x217 + x148 + x166 + x206*x43 + x218*x4 + x219*x4 + x220*x5 + 3.0*x97; + basis_lapl_x_eval[ipt + 10*npts] = x*x223 + x*x224 + x18*x182 + x18*x222 + x221*x7; + basis_lapl_x_eval[ipt + 11*npts] = z*(x*x203 + x*x204 + x14*x222 + x202 + x226); + basis_lapl_x_eval[ipt + 12*npts] = x*x227 + x*x228 + x107 + x11*x182*x7 + x110 + x19*x222; + basis_lapl_x_eval[ipt + 13*npts] = y*(x*x218 + x*x219 + x17*x222 + x217 + x230); + basis_lapl_x_eval[ipt + 14*npts] = x*x231 + x*x232 + x11*x221 + x182*x20 + x20*x222; + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x0*x235 + x0*x238 + x180*y + x233*x5 + x237*y; + basis_lapl_y_eval[ipt + 1*npts] = 3.0*x136 + x156 + x189*x6 + x206*x26 + x239*y + x240*x6 + x241*x7 + x78; + basis_lapl_y_eval[ipt + 2*npts] = z*(x189*y + x2*x235 + x226 + x239 + x240*y); + basis_lapl_y_eval[ipt + 3*npts] = x140*x225 + x158*x242 + x192*x34 + x196*y + x238*x5*x7 + x242*x84 + x243*x5 + x244*y + x245; + basis_lapl_y_eval[ipt + 4*npts] = z*(x171 + x186 + x199*x6 + x211 + x238*x5*y + x246 + x247*x6 + x92); + basis_lapl_y_eval[ipt + 5*npts] = x11*x238*x5 + x113 + x12*x235 + x121 + x201*y + x248*y; + basis_lapl_y_eval[ipt + 6*npts] = x*x251; + basis_lapl_y_eval[ipt + 7*npts] = x55*(x133*x225 + x14*x192 + x154*x242 + x210*y + x238*x7 + x242*x73 + x243 + x252*y); + basis_lapl_y_eval[ipt + 8*npts] = x*(x11*x238*y + x163 + x213 + x215*x6 + x253 + x254*x6 + x255 + x92 + x95); + basis_lapl_y_eval[ipt + 9*npts] = x*(x17*x235 + x219*y + x256 + x257*y + x258); + basis_lapl_y_eval[ipt + 10*npts] = 24.0*x1 + x14*x176 + 12.0*x142 + 4.0*x160 + x18*x238 + x224*y + x259*y + 4.0*x91; + basis_lapl_y_eval[ipt + 11*npts] = x251*z; + basis_lapl_y_eval[ipt + 12*npts] = x11*x238*x7 + x11*x243 + x147*x225 + x162*x242 + x192*x60 + x228*y + x242*x95 + x245 + x260*y; + basis_lapl_y_eval[ipt + 13*npts] = 3.0*x148 + x166 + x206*x62 + x219*x6 + x220*x7 + x257*x6 + x258*y + x97; + basis_lapl_y_eval[ipt + 14*npts] = x11*x233 + x20*x235 + x20*x238 + x232*y + x261*y; + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x0*x264 + x0*x265 + x178*z + x237*z + x262*x5; + basis_lapl_z_eval[ipt + 1*npts] = y*(x188*z + x2*x264 + x230 + x240*z + x266); + basis_lapl_z_eval[ipt + 2*npts] = x10*x188 + x10*x240 + x11*x241 + x136 + 3.0*x156 + x206*x31 + x266*z + x78; + basis_lapl_z_eval[ipt + 3*npts] = x127 + x129 + x195*z + x244*z + x264*x8 + x265*x5*x7; + basis_lapl_z_eval[ipt + 4*npts] = y*(x10*x198 + x10*x247 + x140 + x187 + x216 + x246 + x265*x5*z + x85 + x94); + basis_lapl_z_eval[ipt + 5*npts] = x11*x265*x5 + x140*x267 + x158*x229 + x192*x43 + x200*z + x248*z + x267*x84 + x268*x5 + x269; + basis_lapl_z_eval[ipt + 6*npts] = x*(x14*x264 + x203*z + x250*z + x256 + x270); + basis_lapl_z_eval[ipt + 7*npts] = x*(x10*x209 + x10*x252 + x146 + x208 + x249 + x255 + x265*x7*z + x93 + x94); + basis_lapl_z_eval[ipt + 8*npts] = x66*(x11*x265 + x133*x267 + x154*x229 + x17*x192 + x214*z + x254*z + x267*x73 + x268); + basis_lapl_z_eval[ipt + 9*npts] = x*x271; + basis_lapl_z_eval[ipt + 10*npts] = x18*x264 + x18*x265 + x223*z + x259*z + x262*x7; + basis_lapl_z_eval[ipt + 11*npts] = x10*x203 + x10*x250 + x11*x205 + x144 + 3.0*x160 + x206*x60 + x270*z + x91; + basis_lapl_z_eval[ipt + 12*npts] = x11*x265*x7 + x145*x267 + x161*x229 + x192*x62 + x227*z + x260*z + x267*x93 + x268*x7 + x269; + basis_lapl_z_eval[ipt + 13*npts] = x271*y; + basis_lapl_z_eval[ipt + 14*npts] = 4.0*x148 + 12.0*x164 + x17*x176 + x20*x265 + x231*z + x261*z + 24.0*x3 + 4.0*x97; + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x2; + ang_eval_2 = x2*x3; + ang_eval_3 = radial_eval*x8; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x5*x9; + ang_eval_1 = radial_eval*x12; + ang_eval_2 = x13*x14; + ang_eval_3 = x15*x7; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x11*x16; + ang_eval_1 = x13*x17; + ang_eval_2 = radial_eval*x18; + ang_eval_3 = x14*x3; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + basis_eval[ipt + 10*npts] = ang_eval_2; + basis_eval[ipt + 11*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x19; + ang_eval_1 = x1*x17; + ang_eval_2 = radial_eval*x20; + basis_eval[ipt + 12*npts] = ang_eval_0; + basis_eval[ipt + 13*npts] = ang_eval_1; + basis_eval[ipt + 14*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x21 + x2*x22; + dang_eval_y_0 = x0*x48; + dang_eval_z_0 = x0*x64; + dang_eval_x_1 = x24*y; + dang_eval_y_1 = x27 + x49; + dang_eval_z_1 = x50; + dang_eval_x_2 = x24*z; + dang_eval_y_2 = x50; + dang_eval_z_2 = x32 + x49; + dang_eval_x_3 = x25*x7 + x27; + dang_eval_y_3 = x35 + x5*x51; + dang_eval_z_3 = x64*x8; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x28*x30; + dang_eval_y_0 = x53*z; + dang_eval_z_0 = y*(x40 + x52); + dang_eval_x_1 = x11*x25 + x32; + dang_eval_y_1 = x12*x48; + dang_eval_z_1 = x44 + x5*x65; + dang_eval_x_2 = x33 + x35; + dang_eval_y_2 = x*x54; + dang_eval_z_2 = x46; + dang_eval_x_3 = x38*z; + dang_eval_y_3 = x55*x57; + dang_eval_z_3 = x*(x36 + x58); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x41*y; + dang_eval_y_0 = x*(x39 + x58); + dang_eval_z_0 = x66*x68; + dang_eval_x_1 = x42 + x44; + dang_eval_y_1 = x47; + dang_eval_z_1 = x*x69; + dang_eval_x_2 = x18*x45; + dang_eval_y_2 = radial_eval_alpha*x59 + x14*x22; + dang_eval_z_2 = x18*x64; + dang_eval_x_3 = x46; + dang_eval_y_3 = x54*z; + dang_eval_z_3 = x33 + x61; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + basis_x_eval[ipt + 10*npts] = dang_eval_x_2; + basis_y_eval[ipt + 10*npts] = dang_eval_y_2; + basis_z_eval[ipt + 10*npts] = dang_eval_z_2; + basis_x_eval[ipt + 11*npts] = dang_eval_x_3; + basis_y_eval[ipt + 11*npts] = dang_eval_y_3; + basis_z_eval[ipt + 11*npts] = dang_eval_z_3; + + dang_eval_x_0 = x19*x45; + dang_eval_y_0 = x11*x51 + x61; + dang_eval_z_0 = x63 + x65*x7; + dang_eval_x_1 = x47; + dang_eval_y_1 = x42 + x63; + dang_eval_z_1 = x69*y; + dang_eval_x_2 = x20*x45; + dang_eval_y_2 = x20*x48; + dang_eval_z_2 = radial_eval_alpha*x70 + x17*x22; + basis_x_eval[ipt + 12*npts] = dang_eval_x_0; + basis_y_eval[ipt + 12*npts] = dang_eval_y_0; + basis_z_eval[ipt + 12*npts] = dang_eval_z_0; + basis_x_eval[ipt + 13*npts] = dang_eval_x_1; + basis_y_eval[ipt + 13*npts] = dang_eval_y_1; + basis_z_eval[ipt + 13*npts] = dang_eval_z_1; + basis_x_eval[ipt + 14*npts] = dang_eval_x_2; + basis_y_eval[ipt + 14*npts] = dang_eval_y_2; + basis_z_eval[ipt + 14*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_laplacian.hpp new file mode 100644 index 000000000..617b9a10c --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_cartesian_l4_laplacian.hpp @@ -0,0 +1,521 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_cartesian_laplacian_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_cartesian_laplacian_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = x*x*x*x; + const auto x1 = radial_eval*y; + const auto x2 = x*x*x; + const auto x3 = radial_eval*z; + const auto x4 = x*x; + const auto x5 = y*y; + const auto x6 = x4*x5; + const auto x7 = x1*z; + const auto x8 = z*z; + const auto x9 = x4*x8; + const auto x10 = radial_eval*x; + const auto x11 = y*y*y; + const auto x12 = x*x3; + const auto x13 = x*x1; + const auto x14 = z*z*z; + const auto x15 = y*y*y*y; + const auto x16 = x5*x8; + const auto x17 = z*z*z*z; + const auto x18 = x*x*x*x*x; + const auto x19 = 4.0*radial_eval; + const auto x20 = 3.0*radial_eval; + const auto x21 = radial_eval_alpha*x0 + x20*x4; + const auto x22 = 2.0*x10; + const auto x23 = x2*x5; + const auto x24 = radial_eval_alpha*x23; + const auto x25 = y*z; + const auto x26 = radial_eval_alpha*x2; + const auto x27 = x22 + x26; + const auto x28 = x2*x8; + const auto x29 = radial_eval_alpha*x28; + const auto x30 = radial_eval*x11; + const auto x31 = x11*x4; + const auto x32 = radial_eval_alpha*x31; + const auto x33 = radial_eval*x5; + const auto x34 = radial_eval_alpha*x6; + const auto x35 = x33 + x34; + const auto x36 = radial_eval*x8; + const auto x37 = radial_eval_alpha*x9; + const auto x38 = x36 + x37; + const auto x39 = radial_eval*x14; + const auto x40 = x14*x4; + const auto x41 = radial_eval_alpha*x40; + const auto x42 = radial_eval_alpha*x; + const auto x43 = x11*x42*z; + const auto x44 = x14*x42*y; + const auto x45 = radial_eval_alpha*y; + const auto x46 = radial_eval*x2; + const auto x47 = radial_eval_alpha*x2*x25; + const auto x48 = 2.0*x1; + const auto x49 = radial_eval*x4; + const auto x50 = x34 + x49; + const auto x51 = radial_eval_alpha*x15 + x20*x5; + const auto x52 = x*z; + const auto x53 = radial_eval_alpha*x11; + const auto x54 = x48 + x53; + const auto x55 = radial_eval_alpha*x16; + const auto x56 = y*y*y*y*y; + const auto x57 = x11*x8; + const auto x58 = radial_eval_alpha*x57; + const auto x59 = x14*x5; + const auto x60 = radial_eval_alpha*x59; + const auto x61 = radial_eval_alpha*z; + const auto x62 = 2.0*x3; + const auto x63 = x*y; + const auto x64 = radial_eval_alpha*x14; + const auto x65 = x62 + x64; + const auto x66 = radial_eval_alpha*x17 + x20*x8; + const auto x67 = z*z*z*z*z; + const auto x68 = 12.0*radial_eval; + const auto x69 = 8.0*radial_eval_alpha; + const auto x70 = radial_eval_alpha + radial_eval_alpha_squared*x4; + const auto x71 = x0*x69 + x0*x70 + x4*x68; + const auto x72 = 6.0*radial_eval_alpha; + const auto x73 = 6.0*x10 + x2*x70; + const auto x74 = x2*x72 + x73; + const auto x75 = 4.0*radial_eval_alpha; + const auto x76 = x6*x75; + const auto x77 = 2.0*radial_eval; + const auto x78 = x5*x77; + const auto x79 = x4*x5*x70 + x78; + const auto x80 = x4*x70 + x77; + const auto x81 = x75*x9; + const auto x82 = x77*x8; + const auto x83 = x4*x70*x8 + x82; + const auto x84 = 2.0*radial_eval_alpha; + const auto x85 = x11*x84; + const auto x86 = x11*x70; + const auto x87 = x5*x84; + const auto x88 = x5*x70; + const auto x89 = x8*x84; + const auto x90 = x70*x8; + const auto x91 = x14*x84; + const auto x92 = x14*x70; + const auto x93 = x15*x70; + const auto x94 = x5*x70*x8; + const auto x95 = x17*x70; + const auto x96 = radial_eval_alpha_squared*x18 + x2*x75; + const auto x97 = 3.0*radial_eval_alpha; + const auto x98 = x6*x97; + const auto x99 = x25*(radial_eval_alpha_squared*x0 + x4*x97); + const auto x100 = 2.0*x42; + const auto x101 = 2.0*x45; + const auto x102 = radial_eval_alpha_squared*x23; + const auto x103 = x100*x5 + x102; + const auto x104 = radial_eval_alpha_squared*x28; + const auto x105 = x100*x8 + x104; + const auto x106 = radial_eval_alpha_squared*x31; + const auto x107 = x101*x4 + x106; + const auto x108 = radial_eval_alpha_squared*x4*x5*x8; + const auto x109 = x108 + x55; + const auto x110 = radial_eval_alpha_squared*x40; + const auto x111 = radial_eval_alpha_squared*x56 + x11*x75; + const auto x112 = x52*(radial_eval_alpha_squared*x15 + x5*x97); + const auto x113 = radial_eval_alpha_squared*x57; + const auto x114 = x101*x8 + x113; + const auto x115 = radial_eval_alpha_squared*x59; + const auto x116 = x9*x97; + const auto x117 = 2.0*x61; + const auto x118 = x110 + x117*x4; + const auto x119 = x115 + x117*x5; + const auto x120 = x63*(radial_eval_alpha_squared*x17 + x8*x97); + const auto x121 = radial_eval_alpha_squared*x67 + x14*x75; + const auto x122 = radial_eval_alpha + radial_eval_alpha_squared*x5; + const auto x123 = x0*x122; + const auto x124 = x2*x84; + const auto x125 = x122*x2; + const auto x126 = x4*x77; + const auto x127 = x122*x4*x5 + x126; + const auto x128 = x4*x84; + const auto x129 = x122*x4; + const auto x130 = x122*x4*x8; + const auto x131 = 6.0*x1 + x11*x122; + const auto x132 = x11*x72 + x131; + const auto x133 = x122*x5 + x77; + const auto x134 = x122*x8; + const auto x135 = x122*x14; + const auto x136 = x122*x15 + x15*x69 + x5*x68; + const auto x137 = x16*x75; + const auto x138 = x122*x5*x8 + x82; + const auto x139 = x122*x17; + const auto x140 = x16*x97; + const auto x141 = radial_eval_alpha + radial_eval_alpha_squared*x8; + const auto x142 = x0*x141; + const auto x143 = x141*x2; + const auto x144 = x141*x4*x5; + const auto x145 = x141*x4; + const auto x146 = x126 + x141*x4*x8; + const auto x147 = x11*x141; + const auto x148 = x141*x5; + const auto x149 = x141*x8 + x77; + const auto x150 = x14*x141 + 6.0*x3; + const auto x151 = x14*x72 + x150; + const auto x152 = x141*x15; + const auto x153 = x141*x5*x8 + x78; + const auto x154 = x141*x17 + x17*x69 + x68*x8; + const auto x155 = x125 + x143 + x2*x69 + x73; + const auto x156 = x11*x69 + x131 + x147 + x86; + const auto x157 = x135 + x14*x69 + x150 + x92; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x0; + basis_eval[ipt + 1*npts] = x1*x2; + basis_eval[ipt + 2*npts] = x2*x3; + basis_eval[ipt + 3*npts] = radial_eval*x6; + basis_eval[ipt + 4*npts] = x4*x7; + basis_eval[ipt + 5*npts] = radial_eval*x9; + basis_eval[ipt + 6*npts] = x10*x11; + basis_eval[ipt + 7*npts] = x12*x5; + basis_eval[ipt + 8*npts] = x13*x8; + basis_eval[ipt + 9*npts] = x10*x14; + basis_eval[ipt + 10*npts] = radial_eval*x15; + basis_eval[ipt + 11*npts] = x11*x3; + basis_eval[ipt + 12*npts] = radial_eval*x16; + basis_eval[ipt + 13*npts] = x1*x14; + basis_eval[ipt + 14*npts] = radial_eval*x17; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x18 + x19*x2; + basis_x_eval[ipt + 1*npts] = x21*y; + basis_x_eval[ipt + 2*npts] = x21*z; + basis_x_eval[ipt + 3*npts] = x22*x5 + x24; + basis_x_eval[ipt + 4*npts] = x25*x27; + basis_x_eval[ipt + 5*npts] = x22*x8 + x29; + basis_x_eval[ipt + 6*npts] = x30 + x32; + basis_x_eval[ipt + 7*npts] = x35*z; + basis_x_eval[ipt + 8*npts] = x38*y; + basis_x_eval[ipt + 9*npts] = x39 + x41; + basis_x_eval[ipt + 10*npts] = x15*x42; + basis_x_eval[ipt + 11*npts] = x43; + basis_x_eval[ipt + 12*npts] = x16*x42; + basis_x_eval[ipt + 13*npts] = x44; + basis_x_eval[ipt + 14*npts] = x17*x42; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*x45; + basis_y_eval[ipt + 1*npts] = x24 + x46; + basis_y_eval[ipt + 2*npts] = x47; + basis_y_eval[ipt + 3*npts] = x32 + x4*x48; + basis_y_eval[ipt + 4*npts] = x50*z; + basis_y_eval[ipt + 5*npts] = x45*x9; + basis_y_eval[ipt + 6*npts] = x*x51; + basis_y_eval[ipt + 7*npts] = x52*x54; + basis_y_eval[ipt + 8*npts] = x*(x36 + x55); + basis_y_eval[ipt + 9*npts] = x44; + basis_y_eval[ipt + 10*npts] = radial_eval_alpha*x56 + x11*x19; + basis_y_eval[ipt + 11*npts] = x51*z; + basis_y_eval[ipt + 12*npts] = x48*x8 + x58; + basis_y_eval[ipt + 13*npts] = x39 + x60; + basis_y_eval[ipt + 14*npts] = x17*x45; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x0*x61; + basis_z_eval[ipt + 1*npts] = x47; + basis_z_eval[ipt + 2*npts] = x29 + x46; + basis_z_eval[ipt + 3*npts] = x6*x61; + basis_z_eval[ipt + 4*npts] = y*(x37 + x49); + basis_z_eval[ipt + 5*npts] = x4*x62 + x41; + basis_z_eval[ipt + 6*npts] = x43; + basis_z_eval[ipt + 7*npts] = x*(x33 + x55); + basis_z_eval[ipt + 8*npts] = x63*x65; + basis_z_eval[ipt + 9*npts] = x*x66; + basis_z_eval[ipt + 10*npts] = x15*x61; + basis_z_eval[ipt + 11*npts] = x30 + x58; + basis_z_eval[ipt + 12*npts] = x5*x62 + x60; + basis_z_eval[ipt + 13*npts] = x66*y; + basis_z_eval[ipt + 14*npts] = radial_eval_alpha*x67 + x14*x19; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x123 + x142 + x71; + basis_lapl_eval[ipt + 1*npts] = x155*y; + basis_lapl_eval[ipt + 2*npts] = x155*z; + basis_lapl_eval[ipt + 3*npts] = x127 + x144 + x6*x69 + x79; + basis_lapl_eval[ipt + 4*npts] = x25*(x129 + x145 + x4*x69 + x80); + basis_lapl_eval[ipt + 5*npts] = x130 + x146 + x69*x9 + x83; + basis_lapl_eval[ipt + 6*npts] = x*x156; + basis_lapl_eval[ipt + 7*npts] = x52*(x133 + x148 + x5*x69 + x88); + basis_lapl_eval[ipt + 8*npts] = x63*(x134 + x149 + x69*x8 + x90); + basis_lapl_eval[ipt + 9*npts] = x*x157; + basis_lapl_eval[ipt + 10*npts] = x136 + x152 + x93; + basis_lapl_eval[ipt + 11*npts] = x156*z; + basis_lapl_eval[ipt + 12*npts] = x138 + x153 + x16*x69 + x94; + basis_lapl_eval[ipt + 13*npts] = x157*y; + basis_lapl_eval[ipt + 14*npts] = x139 + x154 + x95; + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x0; + ang_eval_1 = x1*x2; + ang_eval_2 = x2*x3; + ang_eval_3 = radial_eval*x6; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x4*x7; + ang_eval_1 = radial_eval*x9; + ang_eval_2 = x10*x11; + ang_eval_3 = x12*x5; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = x13*x8; + ang_eval_1 = x10*x14; + ang_eval_2 = radial_eval*x15; + ang_eval_3 = x11*x3; + basis_eval[ipt + 8*npts] = ang_eval_0; + basis_eval[ipt + 9*npts] = ang_eval_1; + basis_eval[ipt + 10*npts] = ang_eval_2; + basis_eval[ipt + 11*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x16; + ang_eval_1 = x1*x14; + ang_eval_2 = radial_eval*x17; + basis_eval[ipt + 12*npts] = ang_eval_0; + basis_eval[ipt + 13*npts] = ang_eval_1; + basis_eval[ipt + 14*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = radial_eval_alpha*x18 + x19*x2; + dang_eval_y_0 = x0*x45; + dang_eval_z_0 = x0*x61; + dang_eval_x_1 = x21*y; + dang_eval_y_1 = x24 + x46; + dang_eval_z_1 = x47; + dang_eval_x_2 = x21*z; + dang_eval_y_2 = x47; + dang_eval_z_2 = x29 + x46; + dang_eval_x_3 = x22*x5 + x24; + dang_eval_y_3 = x32 + x4*x48; + dang_eval_z_3 = x6*x61; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x25*x27; + dang_eval_y_0 = x50*z; + dang_eval_z_0 = y*(x37 + x49); + dang_eval_x_1 = x22*x8 + x29; + dang_eval_y_1 = x45*x9; + dang_eval_z_1 = x4*x62 + x41; + dang_eval_x_2 = x30 + x32; + dang_eval_y_2 = x*x51; + dang_eval_z_2 = x43; + dang_eval_x_3 = x35*z; + dang_eval_y_3 = x52*x54; + dang_eval_z_3 = x*(x33 + x55); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x38*y; + dang_eval_y_0 = x*(x36 + x55); + dang_eval_z_0 = x63*x65; + dang_eval_x_1 = x39 + x41; + dang_eval_y_1 = x44; + dang_eval_z_1 = x*x66; + dang_eval_x_2 = x15*x42; + dang_eval_y_2 = radial_eval_alpha*x56 + x11*x19; + dang_eval_z_2 = x15*x61; + dang_eval_x_3 = x43; + dang_eval_y_3 = x51*z; + dang_eval_z_3 = x30 + x58; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + basis_x_eval[ipt + 9*npts] = dang_eval_x_1; + basis_y_eval[ipt + 9*npts] = dang_eval_y_1; + basis_z_eval[ipt + 9*npts] = dang_eval_z_1; + basis_x_eval[ipt + 10*npts] = dang_eval_x_2; + basis_y_eval[ipt + 10*npts] = dang_eval_y_2; + basis_z_eval[ipt + 10*npts] = dang_eval_z_2; + basis_x_eval[ipt + 11*npts] = dang_eval_x_3; + basis_y_eval[ipt + 11*npts] = dang_eval_y_3; + basis_z_eval[ipt + 11*npts] = dang_eval_z_3; + + dang_eval_x_0 = x16*x42; + dang_eval_y_0 = x48*x8 + x58; + dang_eval_z_0 = x5*x62 + x60; + dang_eval_x_1 = x44; + dang_eval_y_1 = x39 + x60; + dang_eval_z_1 = x66*y; + dang_eval_x_2 = x17*x42; + dang_eval_y_2 = x17*x45; + dang_eval_z_2 = radial_eval_alpha*x67 + x14*x19; + basis_x_eval[ipt + 12*npts] = dang_eval_x_0; + basis_y_eval[ipt + 12*npts] = dang_eval_y_0; + basis_z_eval[ipt + 12*npts] = dang_eval_z_0; + basis_x_eval[ipt + 13*npts] = dang_eval_x_1; + basis_y_eval[ipt + 13*npts] = dang_eval_y_1; + basis_z_eval[ipt + 13*npts] = dang_eval_z_1; + basis_x_eval[ipt + 14*npts] = dang_eval_x_2; + basis_y_eval[ipt + 14*npts] = dang_eval_y_2; + basis_z_eval[ipt + 14*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0.hpp new file mode 100644 index 000000000..90bd73c8a --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0.hpp @@ -0,0 +1,140 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_0_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_0_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_spherical_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_gradient.hpp new file mode 100644 index 000000000..398174edb --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_gradient.hpp @@ -0,0 +1,163 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_0_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_0_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_spherical_gradient_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_hessian.hpp new file mode 100644 index 000000000..8dbf7eca6 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_hessian.hpp @@ -0,0 +1,193 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_0_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_0_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_spherical_hessian_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = radial_eval_alpha_squared*(x*x); + const auto x1 = radial_eval_alpha_squared*x; + const auto x2 = radial_eval_alpha_squared*(y*y); + const auto x3 = radial_eval_alpha_squared*(z*z); + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = radial_eval_alpha + x0; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x1*y; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x1*z; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = radial_eval_alpha + x2; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = radial_eval_alpha_squared*y*z; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = radial_eval_alpha + x3; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_lapgrad.hpp new file mode 100644 index 000000000..b1c79105d --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_lapgrad.hpp @@ -0,0 +1,215 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_0_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_0_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_spherical_lapgrad_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = x*x; + const auto x1 = radial_eval_alpha_squared*x0; + const auto x2 = radial_eval_alpha_squared*x; + const auto x3 = y*y; + const auto x4 = radial_eval_alpha_squared*x3; + const auto x5 = radial_eval_alpha_squared*y; + const auto x6 = z*z; + const auto x7 = radial_eval_alpha_squared*x6; + const auto x8 = radial_eval_alpha_cubed*x; + const auto x9 = radial_eval_alpha_cubed*y; + const auto x10 = radial_eval_alpha_cubed*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = radial_eval_alpha + x1; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x2*y; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x2*z; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = radial_eval_alpha + x4; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x5*z; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = radial_eval_alpha + x7; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = 3.0*radial_eval_alpha + x1 + x4 + x7; + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = radial_eval_alpha_cubed*(x*x*x) + 5.0*x2 + x3*x8 + x6*x8; + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = radial_eval_alpha_cubed*(y*y*y) + x0*x9 + 5.0*x5 + x6*x9; + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = radial_eval_alpha_cubed*(z*z*z) + 5.0*radial_eval_alpha_squared*z + x0*x10 + x10*x3; + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_laplacian.hpp new file mode 100644 index 000000000..94969ab88 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l0_laplacian.hpp @@ -0,0 +1,173 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_0_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_0_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_spherical_laplacian_0( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = radial_eval_alpha_squared*(x*x); + const auto x1 = radial_eval_alpha_squared*x; + const auto x2 = radial_eval_alpha_squared*(y*y); + const auto x3 = radial_eval_alpha_squared*(z*z); + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = radial_eval_alpha*x; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval_alpha*y; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = radial_eval_alpha*z; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = 3.0*radial_eval_alpha + x0 + x2 + x3; + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + + + ang_eval_0 = radial_eval; + basis_eval[ipt + 0*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + + dang_eval_x_0 = radial_eval_alpha*x; + dang_eval_y_0 = radial_eval_alpha*y; + dang_eval_z_0 = radial_eval_alpha*z; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1.hpp new file mode 100644 index 000000000..b3c937f6e --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1.hpp @@ -0,0 +1,148 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_1_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_1_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_spherical_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*y; + basis_eval[ipt + 1*npts] = radial_eval*z; + basis_eval[ipt + 2*npts] = radial_eval*x; + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*y; + ang_eval_1 = radial_eval*z; + ang_eval_2 = radial_eval*x; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_gradient.hpp new file mode 100644 index 000000000..6550449be --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_gradient.hpp @@ -0,0 +1,195 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_1_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_1_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_spherical_gradient_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = radial_eval_alpha*x; + const auto x1 = x0*y; + const auto x2 = x0*z; + const auto x3 = radial_eval_alpha*y*z; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*y; + basis_eval[ipt + 1*npts] = radial_eval*z; + basis_eval[ipt + 2*npts] = radial_eval*x; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1; + basis_x_eval[ipt + 1*npts] = x2; + basis_x_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*(x*x); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*(y*y); + basis_y_eval[ipt + 1*npts] = x3; + basis_y_eval[ipt + 2*npts] = x1; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x3; + basis_z_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*(z*z); + basis_z_eval[ipt + 2*npts] = x2; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*y; + ang_eval_1 = radial_eval*z; + ang_eval_2 = radial_eval*x; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = x1; + dang_eval_y_0 = radial_eval + radial_eval_alpha*(y*y); + dang_eval_z_0 = x3; + dang_eval_x_1 = x2; + dang_eval_y_1 = x3; + dang_eval_z_1 = radial_eval + radial_eval_alpha*(z*z); + dang_eval_x_2 = radial_eval + radial_eval_alpha*(x*x); + dang_eval_y_2 = x1; + dang_eval_z_2 = x2; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_hessian.hpp new file mode 100644 index 000000000..d4a583805 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_hessian.hpp @@ -0,0 +1,252 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_1_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_1_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_hessian_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = radial_eval_alpha*x; + const auto x1 = x0*y; + const auto x2 = x0*z; + const auto x3 = x*x; + const auto x4 = y*y; + const auto x5 = y*z; + const auto x6 = radial_eval_alpha*x5; + const auto x7 = z*z; + const auto x8 = radial_eval_alpha_squared*x3; + const auto x9 = radial_eval_alpha + x8; + const auto x10 = x9*y; + const auto x11 = x9*z; + const auto x12 = 3.0*radial_eval_alpha; + const auto x13 = radial_eval_alpha_squared*x4; + const auto x14 = radial_eval_alpha + x13; + const auto x15 = x*x14; + const auto x16 = radial_eval_alpha_squared*x*x5; + const auto x17 = radial_eval_alpha_squared*x7; + const auto x18 = radial_eval_alpha + x17; + const auto x19 = x*x18; + const auto x20 = x14*z; + const auto x21 = x18*y; + const auto x22 = 5.0*radial_eval_alpha + x13 + x17 + x8; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*y; + basis_eval[ipt + 1*npts] = radial_eval*z; + basis_eval[ipt + 2*npts] = radial_eval*x; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1; + basis_x_eval[ipt + 1*npts] = x2; + basis_x_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*x3; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*x4; + basis_y_eval[ipt + 1*npts] = x6; + basis_y_eval[ipt + 2*npts] = x1; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x6; + basis_z_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*x7; + basis_z_eval[ipt + 2*npts] = x2; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x10; + basis_xx_eval[ipt + 1*npts] = x11; + basis_xx_eval[ipt + 2*npts] = x*(x12 + x8); + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x15; + basis_xy_eval[ipt + 1*npts] = x16; + basis_xy_eval[ipt + 2*npts] = x10; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x16; + basis_xz_eval[ipt + 1*npts] = x19; + basis_xz_eval[ipt + 2*npts] = x11; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = y*(x12 + x13); + basis_yy_eval[ipt + 1*npts] = x20; + basis_yy_eval[ipt + 2*npts] = x15; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x20; + basis_yz_eval[ipt + 1*npts] = x21; + basis_yz_eval[ipt + 2*npts] = x16; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x21; + basis_zz_eval[ipt + 1*npts] = z*(x12 + x17); + basis_zz_eval[ipt + 2*npts] = x19; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*y; + ang_eval_1 = radial_eval*z; + ang_eval_2 = radial_eval*x; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = x1; + dang_eval_y_0 = radial_eval + radial_eval_alpha*x4; + dang_eval_z_0 = x6; + dang_eval_x_1 = x2; + dang_eval_y_1 = x6; + dang_eval_z_1 = radial_eval + radial_eval_alpha*x7; + dang_eval_x_2 = radial_eval + radial_eval_alpha*x3; + dang_eval_y_2 = x1; + dang_eval_z_2 = x2; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_lapgrad.hpp new file mode 100644 index 000000000..a80cd7e50 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_lapgrad.hpp @@ -0,0 +1,292 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_1_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_1_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_lapgrad_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = radial_eval_alpha*x; + const auto x1 = x0*y; + const auto x2 = x0*z; + const auto x3 = x*x; + const auto x4 = x3; + const auto x5 = y*y; + const auto x6 = x5; + const auto x7 = y*z; + const auto x8 = radial_eval_alpha*x7; + const auto x9 = z*z; + const auto x10 = x9; + const auto x11 = radial_eval_alpha_squared*x4; + const auto x12 = radial_eval_alpha + x11; + const auto x13 = x12*y; + const auto x14 = x12*z; + const auto x15 = 3.0*radial_eval_alpha; + const auto x16 = radial_eval_alpha_squared*x6; + const auto x17 = radial_eval_alpha + x16; + const auto x18 = x*x17; + const auto x19 = radial_eval_alpha_squared*x*x7; + const auto x20 = radial_eval_alpha_squared*x10; + const auto x21 = radial_eval_alpha + x20; + const auto x22 = x*x21; + const auto x23 = x17*z; + const auto x24 = x21*y; + const auto x25 = 5.0*radial_eval_alpha; + const auto x26 = x16 + x20 + x25; + const auto x27 = x11 + x26; + const auto x28 = 5.0*radial_eval_alpha_squared; + const auto x29 = radial_eval_alpha_cubed*(x*x*x); + const auto x30 = radial_eval_alpha_cubed*x6 + radial_eval_alpha_squared; + const auto x31 = radial_eval_alpha_cubed*x10 + radial_eval_alpha_squared; + const auto x32 = x*x28 + x*x30 + x*x31 + x29; + const auto x33 = 3.0*radial_eval_alpha_squared; + const auto x34 = radial_eval_alpha_cubed*(y*y*y); + const auto x35 = radial_eval_alpha_cubed*x4 + radial_eval_alpha_squared; + const auto x36 = x11 + x25; + const auto x37 = x28*y + x31*y + x34 + x35*y; + const auto x38 = radial_eval_alpha_cubed*(z*z*z); + const auto x39 = x28*z + x30*z + x35*z + x38; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*y; + basis_eval[ipt + 1*npts] = radial_eval*z; + basis_eval[ipt + 2*npts] = radial_eval*x; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1; + basis_x_eval[ipt + 1*npts] = x2; + basis_x_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*x4; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*x6; + basis_y_eval[ipt + 1*npts] = x8; + basis_y_eval[ipt + 2*npts] = x1; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x8; + basis_z_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*x10; + basis_z_eval[ipt + 2*npts] = x2; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x13; + basis_xx_eval[ipt + 1*npts] = x14; + basis_xx_eval[ipt + 2*npts] = x*(x11 + x15); + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x18; + basis_xy_eval[ipt + 1*npts] = x19; + basis_xy_eval[ipt + 2*npts] = x13; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x19; + basis_xz_eval[ipt + 1*npts] = x22; + basis_xz_eval[ipt + 2*npts] = x14; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = y*(x15 + x16); + basis_yy_eval[ipt + 1*npts] = x23; + basis_yy_eval[ipt + 2*npts] = x18; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x23; + basis_yz_eval[ipt + 1*npts] = x24; + basis_yz_eval[ipt + 2*npts] = x19; + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x24; + basis_zz_eval[ipt + 1*npts] = z*(x15 + x20); + basis_zz_eval[ipt + 2*npts] = x22; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x27*y; + basis_lapl_eval[ipt + 1*npts] = x27*z; + basis_lapl_eval[ipt + 2*npts] = x*x27; + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = x32*y; + basis_lapl_x_eval[ipt + 1*npts] = x32*z; + basis_lapl_x_eval[ipt + 2*npts] = x*(x*x33 + x29) + x26 + x3*x30 + x3*x31 + x33*x4; + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x20 + x31*x5 + x33*x6 + x35*x5 + x36 + y*(x33*y + x34); + basis_lapl_y_eval[ipt + 1*npts] = x37*z; + basis_lapl_y_eval[ipt + 2*npts] = x*x37; + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x39*y; + basis_lapl_z_eval[ipt + 1*npts] = x10*x33 + x16 + x30*x9 + x35*x9 + x36 + z*(x33*z + x38); + basis_lapl_z_eval[ipt + 2*npts] = x*x39; + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*y; + ang_eval_1 = radial_eval*z; + ang_eval_2 = radial_eval*x; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = x1; + dang_eval_y_0 = radial_eval + radial_eval_alpha*x6; + dang_eval_z_0 = x8; + dang_eval_x_1 = x2; + dang_eval_y_1 = x8; + dang_eval_z_1 = radial_eval + radial_eval_alpha*x10; + dang_eval_x_2 = radial_eval + radial_eval_alpha*x4; + dang_eval_y_2 = x1; + dang_eval_z_2 = x2; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_laplacian.hpp new file mode 100644 index 000000000..98ad52ecc --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l1_laplacian.hpp @@ -0,0 +1,222 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_1_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_1_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_laplacian_1( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = radial_eval_alpha*x; + const auto x1 = x0*y; + const auto x2 = x0*z; + const auto x3 = x*x; + const auto x4 = y*y; + const auto x5 = y*z; + const auto x6 = radial_eval_alpha*x5; + const auto x7 = z*z; + const auto x8 = radial_eval_alpha_squared*x3; + const auto x9 = radial_eval_alpha + x8; + const auto x10 = x9*y; + const auto x11 = x9*z; + const auto x12 = 3.0*radial_eval_alpha; + const auto x13 = radial_eval_alpha_squared*x4; + const auto x14 = radial_eval_alpha + x13; + const auto x15 = x*x14; + const auto x16 = radial_eval_alpha_squared*x*x5; + const auto x17 = radial_eval_alpha_squared*x7; + const auto x18 = radial_eval_alpha + x17; + const auto x19 = x*x18; + const auto x20 = x14*z; + const auto x21 = x18*y; + const auto x22 = 5.0*radial_eval_alpha + x13 + x17 + x8; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*y; + basis_eval[ipt + 1*npts] = radial_eval*z; + basis_eval[ipt + 2*npts] = radial_eval*x; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1; + basis_x_eval[ipt + 1*npts] = x2; + basis_x_eval[ipt + 2*npts] = radial_eval + radial_eval_alpha*x3; + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = radial_eval + radial_eval_alpha*x4; + basis_y_eval[ipt + 1*npts] = x6; + basis_y_eval[ipt + 2*npts] = x1; + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x6; + basis_z_eval[ipt + 1*npts] = radial_eval + radial_eval_alpha*x7; + basis_z_eval[ipt + 2*npts] = x2; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x22*y; + basis_lapl_eval[ipt + 1*npts] = x22*z; + basis_lapl_eval[ipt + 2*npts] = x*x22; + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + + + ang_eval_0 = radial_eval*y; + ang_eval_1 = radial_eval*z; + ang_eval_2 = radial_eval*x; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + + dang_eval_x_0 = x1; + dang_eval_y_0 = radial_eval + radial_eval_alpha*x4; + dang_eval_z_0 = x6; + dang_eval_x_1 = x2; + dang_eval_y_1 = x6; + dang_eval_z_1 = radial_eval + radial_eval_alpha*x7; + dang_eval_x_2 = radial_eval + radial_eval_alpha*x3; + dang_eval_y_2 = x1; + dang_eval_z_2 = x2; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2.hpp new file mode 100644 index 000000000..6a631ae2c --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2.hpp @@ -0,0 +1,160 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_2_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_2_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_spherical_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + const auto x0 = radial_eval*sqrt_3*y; + const auto x1 = 0.5*radial_eval; + const auto x2 = x*x; + const auto x3 = y*y; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = x*x0; + basis_eval[ipt + 1*npts] = x0*z; + basis_eval[ipt + 2*npts] = -x1*(x2 + x3 - 2.0*z*z); + basis_eval[ipt + 3*npts] = radial_eval*sqrt_3*x*z; + basis_eval[ipt + 4*npts] = sqrt_3*x1*(x2 - x3); + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = x*x0; + ang_eval_1 = x0*z; + ang_eval_2 = -x1*(x2 + x3 - 2.0*z*z); + ang_eval_3 = radial_eval*sqrt_3*x*z; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = sqrt_3*x1*(x2 - x3); + basis_eval[ipt + 4*npts] = ang_eval_0; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_gradient.hpp new file mode 100644 index 000000000..2f4e3dbbe --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_gradient.hpp @@ -0,0 +1,239 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_2_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_2_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_spherical_gradient_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = sqrt_3*y; + const auto x1 = radial_eval*x0; + const auto x2 = 0.5*radial_eval; + const auto x3 = x*x; + const auto x4 = y*y; + const auto x5 = z*z; + const auto x6 = -x3 - x4 + 2.0*x5; + const auto x7 = sqrt_3*z; + const auto x8 = x3 - x4; + const auto x9 = radial_eval + radial_eval_alpha*x3; + const auto x10 = radial_eval_alpha*x*x0*z; + const auto x11 = 0.5*x; + const auto x12 = 2.0*radial_eval; + const auto x13 = -x12; + const auto x14 = radial_eval_alpha*x6; + const auto x15 = x13 + x14; + const auto x16 = radial_eval_alpha*x8; + const auto x17 = sqrt_3*x; + const auto x18 = radial_eval + radial_eval_alpha*x4; + const auto x19 = radial_eval + radial_eval_alpha*x5; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = x*x1; + basis_eval[ipt + 1*npts] = x1*z; + basis_eval[ipt + 2*npts] = x2*x6; + basis_eval[ipt + 3*npts] = radial_eval*x*x7; + basis_eval[ipt + 4*npts] = sqrt_3*x2*x8; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x0*x9; + basis_x_eval[ipt + 1*npts] = x10; + basis_x_eval[ipt + 2*npts] = x11*x15; + basis_x_eval[ipt + 3*npts] = x7*x9; + basis_x_eval[ipt + 4*npts] = sqrt_3*x11*(x12 + x16); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x17*x18; + basis_y_eval[ipt + 1*npts] = x18*x7; + basis_y_eval[ipt + 2*npts] = 0.5*x15*y; + basis_y_eval[ipt + 3*npts] = x10; + basis_y_eval[ipt + 4*npts] = 0.5*x0*(x13 + x16); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x10; + basis_z_eval[ipt + 1*npts] = x0*x19; + basis_z_eval[ipt + 2*npts] = 0.5*z*(4.0*radial_eval + x14); + basis_z_eval[ipt + 3*npts] = x17*x19; + basis_z_eval[ipt + 4*npts] = 0.5*radial_eval_alpha*x7*x8; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = x*x1; + ang_eval_1 = x1*z; + ang_eval_2 = x2*x6; + ang_eval_3 = radial_eval*x*x7; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = sqrt_3*x2*x8; + basis_eval[ipt + 4*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x0*x9; + dang_eval_y_0 = x17*x18; + dang_eval_z_0 = x10; + dang_eval_x_1 = x10; + dang_eval_y_1 = x18*x7; + dang_eval_z_1 = x0*x19; + dang_eval_x_2 = x11*x15; + dang_eval_y_2 = 0.5*x15*y; + dang_eval_z_2 = 0.5*z*(4.0*radial_eval + x14); + dang_eval_x_3 = x7*x9; + dang_eval_y_3 = x10; + dang_eval_z_3 = x17*x19; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = sqrt_3*x11*(x12 + x16); + dang_eval_y_0 = 0.5*x0*(x13 + x16); + dang_eval_z_0 = 0.5*radial_eval_alpha*x7*x8; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_hessian.hpp new file mode 100644 index 000000000..5737b20ee --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_hessian.hpp @@ -0,0 +1,324 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_2_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_2_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_hessian_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = sqrt_3*y; + const auto x1 = x*x0; + const auto x2 = x0*z; + const auto x3 = 0.5*radial_eval; + const auto x4 = x*x; + const auto x5 = y*y; + const auto x6 = z*z; + const auto x7 = -x4 - x5 + 2.0*x6; + const auto x8 = sqrt_3*z; + const auto x9 = x*x8; + const auto x10 = x4 - x5; + const auto x11 = radial_eval + radial_eval_alpha*x4; + const auto x12 = radial_eval_alpha*x1*z; + const auto x13 = 0.5*x; + const auto x14 = 2.0*radial_eval; + const auto x15 = -x14; + const auto x16 = radial_eval_alpha*x7; + const auto x17 = x15 + x16; + const auto x18 = radial_eval_alpha*x10; + const auto x19 = sqrt_3*x; + const auto x20 = radial_eval_alpha*x5; + const auto x21 = radial_eval + x20; + const auto x22 = 0.5*y; + const auto x23 = radial_eval_alpha*x6; + const auto x24 = radial_eval + x23; + const auto x25 = 0.5*z; + const auto x26 = 4.0*radial_eval; + const auto x27 = 3.0*radial_eval_alpha; + const auto x28 = radial_eval_alpha_squared*x4; + const auto x29 = x27 + x28; + const auto x30 = radial_eval_alpha + x28; + const auto x31 = x2*x30; + const auto x32 = 4.0*radial_eval_alpha; + const auto x33 = x32*x4; + const auto x34 = x14 + x33; + const auto x35 = 0.5*sqrt_3; + const auto x36 = x10*x30; + const auto x37 = radial_eval_alpha_squared*x5; + const auto x38 = radial_eval_alpha + x37; + const auto x39 = x38*x9; + const auto x40 = radial_eval_alpha_squared*x7; + const auto x41 = radial_eval_alpha_squared*x6; + const auto x42 = radial_eval_alpha + x41; + const auto x43 = x1*x42; + const auto x44 = 2.0*radial_eval_alpha; + const auto x45 = x40 + x44; + const auto x46 = radial_eval_alpha_squared*x10; + const auto x47 = x27 + x37; + const auto x48 = x32*x5; + const auto x49 = x14 + x48; + const auto x50 = x27 + x41; + const auto x51 = 8.0*radial_eval_alpha*x6 + x42*x7; + const auto x52 = x10*x42; + const auto x53 = 7.0*radial_eval_alpha + x28 + x37 + x41; + const auto x54 = -x48; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x1; + basis_eval[ipt + 1*npts] = radial_eval*x2; + basis_eval[ipt + 2*npts] = x3*x7; + basis_eval[ipt + 3*npts] = radial_eval*x9; + basis_eval[ipt + 4*npts] = sqrt_3*x10*x3; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x0*x11; + basis_x_eval[ipt + 1*npts] = x12; + basis_x_eval[ipt + 2*npts] = x13*x17; + basis_x_eval[ipt + 3*npts] = x11*x8; + basis_x_eval[ipt + 4*npts] = sqrt_3*x13*(x14 + x18); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x19*x21; + basis_y_eval[ipt + 1*npts] = x21*x8; + basis_y_eval[ipt + 2*npts] = x17*x22; + basis_y_eval[ipt + 3*npts] = x12; + basis_y_eval[ipt + 4*npts] = 0.5*x0*(x15 + x18); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x12; + basis_z_eval[ipt + 1*npts] = x0*x24; + basis_z_eval[ipt + 2*npts] = x25*(x16 + x26); + basis_z_eval[ipt + 3*npts] = x19*x24; + basis_z_eval[ipt + 4*npts] = 0.5*radial_eval_alpha*x10*x8; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x1*x29; + basis_xx_eval[ipt + 1*npts] = x31; + basis_xx_eval[ipt + 2*npts] = 0.5*x30*x7 - 0.5*x34; + basis_xx_eval[ipt + 3*npts] = x29*x9; + basis_xx_eval[ipt + 4*npts] = x35*(x34 + x36); + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = sqrt_3*(radial_eval_alpha_squared*x4*x5 + x11 + x20); + basis_xy_eval[ipt + 1*npts] = x39; + basis_xy_eval[ipt + 2*npts] = x13*y*(-x32 + x40); + basis_xy_eval[ipt + 3*npts] = x31; + basis_xy_eval[ipt + 4*npts] = radial_eval_alpha_squared*x0*x10*x13; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x31; + basis_xz_eval[ipt + 1*npts] = x43; + basis_xz_eval[ipt + 2*npts] = x13*x45*z; + basis_xz_eval[ipt + 3*npts] = sqrt_3*(radial_eval_alpha_squared*x4*x6 + x11 + x23); + basis_xz_eval[ipt + 4*npts] = x13*x8*(x44 + x46); + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x1*x47; + basis_yy_eval[ipt + 1*npts] = x2*x47; + basis_yy_eval[ipt + 2*npts] = 0.5*x38*x7 - 0.5*x49; + basis_yy_eval[ipt + 3*npts] = x39; + basis_yy_eval[ipt + 4*npts] = x35*(x10*x38 - x49); + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x39; + basis_yz_eval[ipt + 1*npts] = sqrt_3*(radial_eval_alpha_squared*x5*x6 + x21 + x23); + basis_yz_eval[ipt + 2*npts] = x22*x45*z; + basis_yz_eval[ipt + 3*npts] = x43; + basis_yz_eval[ipt + 4*npts] = x0*x25*(-x44 + x46); + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x43; + basis_zz_eval[ipt + 1*npts] = x2*x50; + basis_zz_eval[ipt + 2*npts] = 0.5*x26 + 0.5*x51; + basis_zz_eval[ipt + 3*npts] = x50*x9; + basis_zz_eval[ipt + 4*npts] = x35*x52; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x1; + ang_eval_1 = radial_eval*x2; + ang_eval_2 = x3*x7; + ang_eval_3 = radial_eval*x9; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = sqrt_3*x10*x3; + basis_eval[ipt + 4*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x0*x11; + dang_eval_y_0 = x19*x21; + dang_eval_z_0 = x12; + dang_eval_x_1 = x12; + dang_eval_y_1 = x21*x8; + dang_eval_z_1 = x0*x24; + dang_eval_x_2 = x13*x17; + dang_eval_y_2 = x17*x22; + dang_eval_z_2 = x25*(x16 + x26); + dang_eval_x_3 = x11*x8; + dang_eval_y_3 = x12; + dang_eval_z_3 = x19*x24; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = sqrt_3*x13*(x14 + x18); + dang_eval_y_0 = 0.5*x0*(x15 + x18); + dang_eval_z_0 = 0.5*radial_eval_alpha*x10*x8; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_lapgrad.hpp new file mode 100644 index 000000000..3c3d62a0e --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_lapgrad.hpp @@ -0,0 +1,393 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_2_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_2_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_lapgrad_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = sqrt_3*y; + const auto x1 = x*x0; + const auto x2 = x0*z; + const auto x3 = 0.5*radial_eval; + const auto x4 = x*x; + const auto x5 = x4; + const auto x6 = y*y; + const auto x7 = x6; + const auto x8 = z*z; + const auto x9 = x8; + const auto x10 = -x5 - x7 + 2.0*x9; + const auto x11 = sqrt_3*z; + const auto x12 = x*x11; + const auto x13 = x5 - x7; + const auto x14 = radial_eval + radial_eval_alpha*x5; + const auto x15 = radial_eval_alpha*x1*z; + const auto x16 = 0.5*x; + const auto x17 = 2.0*radial_eval; + const auto x18 = -x17; + const auto x19 = radial_eval_alpha*x10; + const auto x20 = x18 + x19; + const auto x21 = radial_eval_alpha*x13; + const auto x22 = sqrt_3*x; + const auto x23 = radial_eval_alpha*x7; + const auto x24 = radial_eval + x23; + const auto x25 = 0.5*y; + const auto x26 = radial_eval_alpha*x9; + const auto x27 = radial_eval + x26; + const auto x28 = 0.5*z; + const auto x29 = 4.0*radial_eval; + const auto x30 = 3.0*radial_eval_alpha; + const auto x31 = radial_eval_alpha_squared*x5; + const auto x32 = x30 + x31; + const auto x33 = radial_eval_alpha + x31; + const auto x34 = x2*x33; + const auto x35 = 4.0*radial_eval_alpha; + const auto x36 = x35*x5; + const auto x37 = x17 + x36; + const auto x38 = 0.5*sqrt_3; + const auto x39 = x13*x33; + const auto x40 = radial_eval_alpha_squared*x7; + const auto x41 = radial_eval_alpha + x40; + const auto x42 = x12*x41; + const auto x43 = radial_eval_alpha_squared*x10; + const auto x44 = radial_eval_alpha_squared*x9; + const auto x45 = radial_eval_alpha + x44; + const auto x46 = x1*x45; + const auto x47 = 2.0*radial_eval_alpha; + const auto x48 = x43 + x47; + const auto x49 = radial_eval_alpha_squared*x13; + const auto x50 = x30 + x40; + const auto x51 = x35*x7; + const auto x52 = x17 + x51; + const auto x53 = x30 + x44; + const auto x54 = 8.0*radial_eval_alpha; + const auto x55 = x10*x45 + x54*x9; + const auto x56 = x13*x45; + const auto x57 = x40 + x44; + const auto x58 = 7.0*radial_eval_alpha + x31 + x57; + const auto x59 = -x51; + const auto x60 = radial_eval_alpha_squared*x; + const auto x61 = radial_eval_alpha_cubed*(x*x*x); + const auto x62 = 3.0*x60 + x61; + const auto x63 = radial_eval_alpha_cubed*x7 + radial_eval_alpha_squared; + const auto x64 = radial_eval_alpha_cubed*x9 + radial_eval_alpha_squared; + const auto x65 = 2.0*radial_eval_alpha_squared; + const auto x66 = x*x62 + 3.0*x33 + x35 + x4*x63 + x4*x64 + x5*x65 + x57; + const auto x67 = 4.0*x60*x7; + const auto x68 = 2.0*x; + const auto x69 = 6.0*x*x33 + x*x35 + x41*x68 + x45*x68; + const auto x70 = x13*x63; + const auto x71 = x13*x64; + const auto x72 = radial_eval_alpha_squared*y; + const auto x73 = radial_eval_alpha_cubed*(y*y*y); + const auto x74 = 3.0*x72 + x73; + const auto x75 = radial_eval_alpha_cubed*x5 + radial_eval_alpha_squared; + const auto x76 = x31 + x35; + const auto x77 = 3.0*x41 + x44 + x6*x64 + x6*x75 + x65*x7 + x74*y + x76; + const auto x78 = x35*y; + const auto x79 = 4.0*x5*x72; + const auto x80 = 2.0*y; + const auto x81 = x33*x80; + const auto x82 = 6.0*x41*y; + const auto x83 = x45*x80; + const auto x84 = x13*x75; + const auto x85 = radial_eval_alpha_squared*z; + const auto x86 = radial_eval_alpha_cubed*(z*z*z); + const auto x87 = 3.0*x85 + x86; + const auto x88 = x40 + 3.0*x45 + x63*x8 + x65*x9 + x75*x8 + x76 + x87*z; + const auto x89 = 4.0*z; + const auto x90 = radial_eval_alpha_squared*x89; + const auto x91 = x5*x90; + const auto x92 = -x7*x90; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x1; + basis_eval[ipt + 1*npts] = radial_eval*x2; + basis_eval[ipt + 2*npts] = x10*x3; + basis_eval[ipt + 3*npts] = radial_eval*x12; + basis_eval[ipt + 4*npts] = sqrt_3*x13*x3; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x0*x14; + basis_x_eval[ipt + 1*npts] = x15; + basis_x_eval[ipt + 2*npts] = x16*x20; + basis_x_eval[ipt + 3*npts] = x11*x14; + basis_x_eval[ipt + 4*npts] = sqrt_3*x16*(x17 + x21); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x22*x24; + basis_y_eval[ipt + 1*npts] = x11*x24; + basis_y_eval[ipt + 2*npts] = x20*x25; + basis_y_eval[ipt + 3*npts] = x15; + basis_y_eval[ipt + 4*npts] = 0.5*x0*(x18 + x21); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x15; + basis_z_eval[ipt + 1*npts] = x0*x27; + basis_z_eval[ipt + 2*npts] = x28*(x19 + x29); + basis_z_eval[ipt + 3*npts] = x22*x27; + basis_z_eval[ipt + 4*npts] = 0.5*radial_eval_alpha*x11*x13; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x1*x32; + basis_xx_eval[ipt + 1*npts] = x34; + basis_xx_eval[ipt + 2*npts] = 0.5*x10*x33 - 0.5*x37; + basis_xx_eval[ipt + 3*npts] = x12*x32; + basis_xx_eval[ipt + 4*npts] = x38*(x37 + x39); + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = sqrt_3*(radial_eval_alpha_squared*x5*x7 + x14 + x23); + basis_xy_eval[ipt + 1*npts] = x42; + basis_xy_eval[ipt + 2*npts] = x16*y*(-x35 + x43); + basis_xy_eval[ipt + 3*npts] = x34; + basis_xy_eval[ipt + 4*npts] = radial_eval_alpha_squared*x0*x13*x16; + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x34; + basis_xz_eval[ipt + 1*npts] = x46; + basis_xz_eval[ipt + 2*npts] = x16*x48*z; + basis_xz_eval[ipt + 3*npts] = sqrt_3*(radial_eval_alpha_squared*x5*x9 + x14 + x26); + basis_xz_eval[ipt + 4*npts] = x11*x16*(x47 + x49); + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x1*x50; + basis_yy_eval[ipt + 1*npts] = x2*x50; + basis_yy_eval[ipt + 2*npts] = 0.5*x10*x41 - 0.5*x52; + basis_yy_eval[ipt + 3*npts] = x42; + basis_yy_eval[ipt + 4*npts] = x38*(x13*x41 - x52); + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x42; + basis_yz_eval[ipt + 1*npts] = sqrt_3*(radial_eval_alpha_squared*x7*x9 + x24 + x26); + basis_yz_eval[ipt + 2*npts] = x25*x48*z; + basis_yz_eval[ipt + 3*npts] = x46; + basis_yz_eval[ipt + 4*npts] = x0*x28*(-x47 + x49); + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x46; + basis_zz_eval[ipt + 1*npts] = x2*x53; + basis_zz_eval[ipt + 2*npts] = 0.5*x29 + 0.5*x55; + basis_zz_eval[ipt + 3*npts] = x12*x53; + basis_zz_eval[ipt + 4*npts] = x38*x56; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x1*x58; + basis_lapl_eval[ipt + 1*npts] = x2*x58; + basis_lapl_eval[ipt + 2*npts] = 0.5*x10*x33 + 0.5*x10*x41 - 0.5*x36 + 0.5*x55 + 0.5*x59; + basis_lapl_eval[ipt + 3*npts] = x12*x58; + basis_lapl_eval[ipt + 4*npts] = x38*(x13*x41 + x36 + x39 + x56 + x59); + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = x0*x66; + basis_lapl_x_eval[ipt + 1*npts] = x2*(x*x63 + x*x64 + 7.0*x60 + x61); + basis_lapl_x_eval[ipt + 2*npts] = 4.0*radial_eval_alpha_squared*x*x9 + 0.5*x*x10*x63 + 0.5*x*x10*x64 + 0.5*x10*x62 - 0.5*x67 - 0.5*x69; + basis_lapl_x_eval[ipt + 3*npts] = x11*x66; + basis_lapl_x_eval[ipt + 4*npts] = x38*(x*x70 + x*x71 + x13*x62 - x67 + x69); + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x22*x77; + basis_lapl_y_eval[ipt + 1*npts] = x11*x77; + basis_lapl_y_eval[ipt + 2*npts] = 4.0*radial_eval_alpha_squared*x9*y + 0.5*x10*x64*y + 0.5*x10*x74 + 0.5*x10*x75*y - 0.5*x78 - 0.5*x79 - 0.5*x81 - 0.5*x82 - 0.5*x83; + basis_lapl_y_eval[ipt + 3*npts] = x12*(x64*y + 7.0*x72 + x73 + x75*y); + basis_lapl_y_eval[ipt + 4*npts] = x38*(x13*x74 + x71*y - x78 + x79 - x81 - x82 - x83 + x84*y); + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x1*(x63*z + x75*z + 7.0*x85 + x86); + basis_lapl_z_eval[ipt + 1*npts] = x0*x88; + basis_lapl_z_eval[ipt + 2*npts] = 0.5*x10*x63*z + 0.5*x10*x75*z + 0.5*x10*x87 + 0.5*x33*x89 + 0.5*x41*x89 + 6.0*x45*z + 0.5*x54*z - 0.5*x91 + 0.5*x92; + basis_lapl_z_eval[ipt + 3*npts] = x22*x88; + basis_lapl_z_eval[ipt + 4*npts] = x38*(x13*x87 + x70*z + x84*z + x91 + x92); + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x1; + ang_eval_1 = radial_eval*x2; + ang_eval_2 = x10*x3; + ang_eval_3 = radial_eval*x12; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = sqrt_3*x13*x3; + basis_eval[ipt + 4*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x0*x14; + dang_eval_y_0 = x22*x24; + dang_eval_z_0 = x15; + dang_eval_x_1 = x15; + dang_eval_y_1 = x11*x24; + dang_eval_z_1 = x0*x27; + dang_eval_x_2 = x16*x20; + dang_eval_y_2 = x20*x25; + dang_eval_z_2 = x28*(x19 + x29); + dang_eval_x_3 = x11*x14; + dang_eval_y_3 = x15; + dang_eval_z_3 = x22*x27; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = sqrt_3*x16*(x17 + x21); + dang_eval_y_0 = 0.5*x0*(x18 + x21); + dang_eval_z_0 = 0.5*radial_eval_alpha*x11*x13; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_laplacian.hpp new file mode 100644 index 000000000..7af7499e5 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l2_laplacian.hpp @@ -0,0 +1,284 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_2_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_2_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_laplacian_2( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = sqrt_3*y; + const auto x1 = x*x0; + const auto x2 = x0*z; + const auto x3 = 0.5*radial_eval; + const auto x4 = x*x; + const auto x5 = y*y; + const auto x6 = z*z; + const auto x7 = -x4 - x5 + 2.0*x6; + const auto x8 = sqrt_3*z; + const auto x9 = x*x8; + const auto x10 = x4 - x5; + const auto x11 = radial_eval + radial_eval_alpha*x4; + const auto x12 = radial_eval_alpha*x1*z; + const auto x13 = 0.5*x; + const auto x14 = 2.0*radial_eval; + const auto x15 = -x14; + const auto x16 = radial_eval_alpha*x7; + const auto x17 = x15 + x16; + const auto x18 = radial_eval_alpha*x10; + const auto x19 = sqrt_3*x; + const auto x20 = radial_eval_alpha*x5; + const auto x21 = radial_eval + x20; + const auto x22 = 0.5*y; + const auto x23 = radial_eval_alpha*x6; + const auto x24 = radial_eval + x23; + const auto x25 = 0.5*z; + const auto x26 = 4.0*radial_eval; + const auto x27 = 3.0*radial_eval_alpha; + const auto x28 = radial_eval_alpha_squared*x4; + const auto x29 = x27 + x28; + const auto x30 = radial_eval_alpha + x28; + const auto x31 = x2*x30; + const auto x32 = 4.0*radial_eval_alpha; + const auto x33 = x32*x4; + const auto x34 = x14 + x33; + const auto x35 = 0.5*sqrt_3; + const auto x36 = x10*x30; + const auto x37 = radial_eval_alpha_squared*x5; + const auto x38 = radial_eval_alpha + x37; + const auto x39 = x38*x9; + const auto x40 = radial_eval_alpha_squared*x7; + const auto x41 = radial_eval_alpha_squared*x6; + const auto x42 = radial_eval_alpha + x41; + const auto x43 = x1*x42; + const auto x44 = 2.0*radial_eval_alpha; + const auto x45 = x40 + x44; + const auto x46 = radial_eval_alpha_squared*x10; + const auto x47 = x27 + x37; + const auto x48 = x32*x5; + const auto x49 = x14 + x48; + const auto x50 = x27 + x41; + const auto x51 = 8.0*radial_eval_alpha*x6 + x42*x7; + const auto x52 = x10*x42; + const auto x53 = 7.0*radial_eval_alpha + x28 + x37 + x41; + const auto x54 = -x48; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x1; + basis_eval[ipt + 1*npts] = radial_eval*x2; + basis_eval[ipt + 2*npts] = x3*x7; + basis_eval[ipt + 3*npts] = radial_eval*x9; + basis_eval[ipt + 4*npts] = sqrt_3*x10*x3; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x0*x11; + basis_x_eval[ipt + 1*npts] = x12; + basis_x_eval[ipt + 2*npts] = x13*x17; + basis_x_eval[ipt + 3*npts] = x11*x8; + basis_x_eval[ipt + 4*npts] = sqrt_3*x13*(x14 + x18); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x19*x21; + basis_y_eval[ipt + 1*npts] = x21*x8; + basis_y_eval[ipt + 2*npts] = x17*x22; + basis_y_eval[ipt + 3*npts] = x12; + basis_y_eval[ipt + 4*npts] = 0.5*x0*(x15 + x18); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x12; + basis_z_eval[ipt + 1*npts] = x0*x24; + basis_z_eval[ipt + 2*npts] = x25*(x16 + x26); + basis_z_eval[ipt + 3*npts] = x19*x24; + basis_z_eval[ipt + 4*npts] = 0.5*radial_eval_alpha*x10*x8; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x1*x53; + basis_lapl_eval[ipt + 1*npts] = x2*x53; + basis_lapl_eval[ipt + 2*npts] = 0.5*x30*x7 - 0.5*x33 + 0.5*x38*x7 + 0.5*x51 + 0.5*x54; + basis_lapl_eval[ipt + 3*npts] = x53*x9; + basis_lapl_eval[ipt + 4*npts] = x35*(x10*x38 + x33 + x36 + x52 + x54); + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x1; + ang_eval_1 = radial_eval*x2; + ang_eval_2 = x3*x7; + ang_eval_3 = radial_eval*x9; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = sqrt_3*x10*x3; + basis_eval[ipt + 4*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x0*x11; + dang_eval_y_0 = x19*x21; + dang_eval_z_0 = x12; + dang_eval_x_1 = x12; + dang_eval_y_1 = x21*x8; + dang_eval_z_1 = x0*x24; + dang_eval_x_2 = x13*x17; + dang_eval_y_2 = x17*x22; + dang_eval_z_2 = x25*(x16 + x26); + dang_eval_x_3 = x11*x8; + dang_eval_y_3 = x12; + dang_eval_z_3 = x19*x24; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = sqrt_3*x13*(x14 + x18); + dang_eval_y_0 = 0.5*x0*(x15 + x18); + dang_eval_z_0 = 0.5*radial_eval_alpha*x10*x8; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3.hpp new file mode 100644 index 000000000..c3eb0bd71 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3.hpp @@ -0,0 +1,174 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_3_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_3_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_spherical_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + const auto x0 = 0.25*radial_eval; + const auto x1 = x0*y; + const auto x2 = x*x; + const auto x3 = 3.0*x2; + const auto x4 = y*y; + const auto x5 = -x4; + const auto x6 = radial_eval*z; + const auto x7 = z*z; + const auto x8 = -x2 - x4 + 4.0*x7; + const auto x9 = 0.5*x6; + const auto x10 = 3.0*x4; + const auto x11 = x*x0; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = sqrt_10*x1*(x3 + x5); + basis_eval[ipt + 1*npts] = sqrt_15*x*x6*y; + basis_eval[ipt + 2*npts] = sqrt_6*x1*x8; + basis_eval[ipt + 3*npts] = -x9*(x10 + x3 - 2.0*x7); + basis_eval[ipt + 4*npts] = sqrt_6*x11*x8; + basis_eval[ipt + 5*npts] = sqrt_15*x9*(x2 + x5); + basis_eval[ipt + 6*npts] = sqrt_10*x11*(-x10 + x2); + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = sqrt_10*x1*(x3 + x5); + ang_eval_1 = sqrt_15*x*x6*y; + ang_eval_2 = sqrt_6*x1*x8; + ang_eval_3 = -x9*(x10 + x3 - 2.0*x7); + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = sqrt_6*x11*x8; + ang_eval_1 = sqrt_15*x9*(x2 + x5); + ang_eval_2 = sqrt_10*x11*(-x10 + x2); + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_gradient.hpp new file mode 100644 index 000000000..2f43b037e --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_gradient.hpp @@ -0,0 +1,281 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_3_max_wg = 256; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_3_nwarp = 8; + + +void collocation_device_shell_to_task_kernel_spherical_gradient_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = 0.25*sqrt_10; + const auto x1 = radial_eval*y; + const auto x2 = x*x; + const auto x3 = 3.0*x2; + const auto x4 = y*y; + const auto x5 = -x4; + const auto x6 = x3 + x5; + const auto x7 = sqrt_15*z; + const auto x8 = x7*y; + const auto x9 = radial_eval*x; + const auto x10 = 0.25*sqrt_6; + const auto x11 = z*z; + const auto x12 = -4.0*x11; + const auto x13 = x12 + x4; + const auto x14 = -x13 - x2; + const auto x15 = 0.5*z; + const auto x16 = 3.0*x4; + const auto x17 = -2.0*x11; + const auto x18 = -x16 - x17 - x3; + const auto x19 = 0.5*sqrt_15; + const auto x20 = x19*z; + const auto x21 = x2 + x5; + const auto x22 = -x16; + const auto x23 = x2 + x22; + const auto x24 = x*y; + const auto x25 = x0*x24; + const auto x26 = 6.0*radial_eval; + const auto x27 = 2.0*radial_eval; + const auto x28 = -x27; + const auto x29 = radial_eval_alpha*x14; + const auto x30 = x10*x24*(x28 + x29); + const auto x31 = -x26; + const auto x32 = radial_eval_alpha*x18 + x31; + const auto x33 = radial_eval_alpha*x21; + const auto x34 = radial_eval*(x22 + x3); + const auto x35 = radial_eval_alpha*x0*z; + const auto x36 = x10*z; + const auto x37 = 8.0*radial_eval + x29; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = x0*x1*x6; + basis_eval[ipt + 1*npts] = x8*x9; + basis_eval[ipt + 2*npts] = x1*x10*x14; + basis_eval[ipt + 3*npts] = radial_eval*x15*x18; + basis_eval[ipt + 4*npts] = x10*x14*x9; + basis_eval[ipt + 5*npts] = radial_eval*x20*x21; + basis_eval[ipt + 6*npts] = x0*x23*x9; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x25*(radial_eval_alpha*x6 + x26); + basis_x_eval[ipt + 1*npts] = x8*(radial_eval + radial_eval_alpha*x2); + basis_x_eval[ipt + 2*npts] = x30; + basis_x_eval[ipt + 3*npts] = x*x15*x32; + basis_x_eval[ipt + 4*npts] = -x10*(radial_eval*(x13 + x3) - radial_eval_alpha*x14*x2); + basis_x_eval[ipt + 5*npts] = x*x20*(x27 + x33); + basis_x_eval[ipt + 6*npts] = x0*(radial_eval_alpha*x2*x23 + x34); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*(radial_eval_alpha*x4*x6 + x34); + basis_y_eval[ipt + 1*npts] = x*x7*(radial_eval + radial_eval_alpha*x4); + basis_y_eval[ipt + 2*npts] = -x10*(radial_eval*(x12 + x16 + x2) - radial_eval_alpha*x14*x4); + basis_y_eval[ipt + 3*npts] = x15*x32*y; + basis_y_eval[ipt + 4*npts] = x30; + basis_y_eval[ipt + 5*npts] = x20*y*(x28 + x33); + basis_y_eval[ipt + 6*npts] = x25*(radial_eval_alpha*x23 + x31); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x35*x6*y; + basis_z_eval[ipt + 1*npts] = sqrt_15*x24*(radial_eval + radial_eval_alpha*x11); + basis_z_eval[ipt + 2*npts] = x36*x37*y; + basis_z_eval[ipt + 3*npts] = -1.5*radial_eval*(x17 + x2 + x4) + 0.5*radial_eval_alpha*x11*x18; + basis_z_eval[ipt + 4*npts] = x*x36*x37; + basis_z_eval[ipt + 5*npts] = x19*x21*(radial_eval + radial_eval_alpha*x11); + basis_z_eval[ipt + 6*npts] = x*x23*x35; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = x0*x1*x6; + ang_eval_1 = x8*x9; + ang_eval_2 = x1*x10*x14; + ang_eval_3 = radial_eval*x15*x18; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x10*x14*x9; + ang_eval_1 = radial_eval*x20*x21; + ang_eval_2 = x0*x23*x9; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x25*(radial_eval_alpha*x6 + x26); + dang_eval_y_0 = x0*(radial_eval_alpha*x4*x6 + x34); + dang_eval_z_0 = x35*x6*y; + dang_eval_x_1 = x8*(radial_eval + radial_eval_alpha*x2); + dang_eval_y_1 = x*x7*(radial_eval + radial_eval_alpha*x4); + dang_eval_z_1 = sqrt_15*x24*(radial_eval + radial_eval_alpha*x11); + dang_eval_x_2 = x30; + dang_eval_y_2 = -x10*(radial_eval*(x12 + x16 + x2) - radial_eval_alpha*x14*x4); + dang_eval_z_2 = x36*x37*y; + dang_eval_x_3 = x*x15*x32; + dang_eval_y_3 = x15*x32*y; + dang_eval_z_3 = -1.5*radial_eval*(x17 + x2 + x4) + 0.5*radial_eval_alpha*x11*x18; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = -x10*(radial_eval*(x13 + x3) - radial_eval_alpha*x14*x2); + dang_eval_y_0 = x30; + dang_eval_z_0 = x*x36*x37; + dang_eval_x_1 = x*x20*(x27 + x33); + dang_eval_y_1 = x20*y*(x28 + x33); + dang_eval_z_1 = x19*x21*(radial_eval + radial_eval_alpha*x11); + dang_eval_x_2 = x0*(radial_eval_alpha*x2*x23 + x34); + dang_eval_y_2 = x25*(radial_eval_alpha*x23 + x31); + dang_eval_z_2 = x*x23*x35; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_hessian.hpp new file mode 100644 index 000000000..709ca845a --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_hessian.hpp @@ -0,0 +1,414 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_3_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_3_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_hessian_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = 0.25*sqrt_10; + const auto x1 = x0*y; + const auto x2 = x*x; + const auto x3 = 3.0*x2; + const auto x4 = y*y; + const auto x5 = -x4; + const auto x6 = x3 + x5; + const auto x7 = sqrt_15*z; + const auto x8 = x7*y; + const auto x9 = x*x8; + const auto x10 = 0.25*sqrt_6; + const auto x11 = x10*y; + const auto x12 = z*z; + const auto x13 = -4.0*x12; + const auto x14 = x13 + x4; + const auto x15 = -x14 - x2; + const auto x16 = 0.5*z; + const auto x17 = 3.0*x4; + const auto x18 = -2.0*x12; + const auto x19 = -x17 - x18 - x3; + const auto x20 = x*x10; + const auto x21 = 0.5*sqrt_15; + const auto x22 = x21*z; + const auto x23 = x2 + x5; + const auto x24 = x*x0; + const auto x25 = -x17; + const auto x26 = x2 + x25; + const auto x27 = x*x1; + const auto x28 = 6.0*radial_eval; + const auto x29 = radial_eval + radial_eval_alpha*x2; + const auto x30 = x*x11; + const auto x31 = 2.0*radial_eval; + const auto x32 = -x31; + const auto x33 = radial_eval_alpha*x15; + const auto x34 = x30*(x32 + x33); + const auto x35 = x*x16; + const auto x36 = -x28; + const auto x37 = radial_eval_alpha*x19 + x36; + const auto x38 = -x14 - x3; + const auto x39 = x15*x2; + const auto x40 = x*x22; + const auto x41 = radial_eval_alpha*x23; + const auto x42 = x31 + x41; + const auto x43 = x25 + x3; + const auto x44 = radial_eval*x43; + const auto x45 = x2*x26; + const auto x46 = x4*x6; + const auto x47 = radial_eval_alpha*x4; + const auto x48 = radial_eval + x47; + const auto x49 = -x13 - x17 - x2; + const auto x50 = x15*x4; + const auto x51 = x32 + x41; + const auto x52 = radial_eval_alpha*z; + const auto x53 = sqrt_15*y; + const auto x54 = radial_eval_alpha*x12; + const auto x55 = 8.0*radial_eval; + const auto x56 = x33 + x55; + const auto x57 = -x18 - x2 - x4; + const auto x58 = x12*x19; + const auto x59 = x12*x23; + const auto x60 = radial_eval_alpha_squared*x2; + const auto x61 = radial_eval_alpha + x60; + const auto x62 = x6*x61; + const auto x63 = 12.0*radial_eval_alpha; + const auto x64 = x2*x63; + const auto x65 = x28 + x64; + const auto x66 = 3.0*radial_eval_alpha; + const auto x67 = 4.0*radial_eval_alpha; + const auto x68 = x2*x67; + const auto x69 = x31 + x68; + const auto x70 = x15*x61; + const auto x71 = 2.0*radial_eval_alpha; + const auto x72 = x38*x71 + x70; + const auto x73 = x23*x61; + const auto x74 = x43*x71; + const auto x75 = x26*x61 + x74; + const auto x76 = 6.0*radial_eval_alpha; + const auto x77 = radial_eval_alpha*x43; + const auto x78 = radial_eval_alpha_squared*x46 + x77; + const auto x79 = radial_eval_alpha*x49 + radial_eval_alpha_squared*x50; + const auto x80 = radial_eval_alpha*x38 + radial_eval_alpha_squared*x39; + const auto x81 = radial_eval_alpha_squared*x45 + x77; + const auto x82 = x27*z; + const auto x83 = x30*z*(radial_eval_alpha_squared*x15 + x76); + const auto x84 = radial_eval_alpha_squared*x58 - x12*x76 + x36 + x57*x66; + const auto x85 = x10*z; + const auto x86 = 8.0*radial_eval_alpha; + const auto x87 = x12*x71; + const auto x88 = radial_eval_alpha_squared*x59; + const auto x89 = x0*z; + const auto x90 = radial_eval_alpha_squared*x4; + const auto x91 = radial_eval_alpha + x90; + const auto x92 = x6*x91 + x74; + const auto x93 = x15*x91; + const auto x94 = x49*x71 + x93; + const auto x95 = x4*x63; + const auto x96 = x28 + x95; + const auto x97 = x4*x67; + const auto x98 = x31 + x97; + const auto x99 = radial_eval_alpha_squared*x12; + const auto x100 = radial_eval_alpha + x99; + const auto x101 = x100*x6; + const auto x102 = 16.0*radial_eval_alpha*x12 + x100*x15; + const auto x103 = x102 + x55; + const auto x104 = x100*x19 + x57*x76; + const auto x105 = x23*(x100 + x71); + const auto x106 = x100*x26; + const auto x107 = -x95; + const auto x108 = -x97; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x1*x6; + basis_eval[ipt + 1*npts] = radial_eval*x9; + basis_eval[ipt + 2*npts] = radial_eval*x11*x15; + basis_eval[ipt + 3*npts] = radial_eval*x16*x19; + basis_eval[ipt + 4*npts] = radial_eval*x15*x20; + basis_eval[ipt + 5*npts] = radial_eval*x22*x23; + basis_eval[ipt + 6*npts] = radial_eval*x24*x26; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x27*(radial_eval_alpha*x6 + x28); + basis_x_eval[ipt + 1*npts] = x29*x8; + basis_x_eval[ipt + 2*npts] = x34; + basis_x_eval[ipt + 3*npts] = x35*x37; + basis_x_eval[ipt + 4*npts] = x10*(radial_eval*x38 + radial_eval_alpha*x39); + basis_x_eval[ipt + 5*npts] = x40*x42; + basis_x_eval[ipt + 6*npts] = x0*(radial_eval_alpha*x45 + x44); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*(radial_eval_alpha*x46 + x44); + basis_y_eval[ipt + 1*npts] = x*x48*x7; + basis_y_eval[ipt + 2*npts] = x10*(radial_eval*x49 + radial_eval_alpha*x50); + basis_y_eval[ipt + 3*npts] = x16*x37*y; + basis_y_eval[ipt + 4*npts] = x34; + basis_y_eval[ipt + 5*npts] = x22*x51*y; + basis_y_eval[ipt + 6*npts] = x27*(radial_eval_alpha*x26 + x36); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x1*x52*x6; + basis_z_eval[ipt + 1*npts] = x*x53*(radial_eval + x54); + basis_z_eval[ipt + 2*npts] = x11*x56*z; + basis_z_eval[ipt + 3*npts] = 1.5*radial_eval*x57 + 0.5*radial_eval_alpha*x58; + basis_z_eval[ipt + 4*npts] = x20*x56*z; + basis_z_eval[ipt + 5*npts] = x21*(radial_eval*x23 + radial_eval_alpha*x59); + basis_z_eval[ipt + 6*npts] = x24*x26*x52; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x1*(x62 + x65); + basis_xx_eval[ipt + 1*npts] = x9*(x60 + x66); + basis_xx_eval[ipt + 2*npts] = x11*(x15*x61 - x69); + basis_xx_eval[ipt + 3*npts] = x16*(x19*x61 - x65); + basis_xx_eval[ipt + 4*npts] = x20*(x36 + x72); + basis_xx_eval[ipt + 5*npts] = x22*(x69 + x73); + basis_xx_eval[ipt + 6*npts] = x24*(x28 + x75); + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x24*(x28 + x4*x76 + x78); + basis_xy_eval[ipt + 1*npts] = x7*(radial_eval_alpha_squared*x2*x4 + x29 + x47); + basis_xy_eval[ipt + 2*npts] = x20*(x32 - x4*x71 + x79); + basis_xy_eval[ipt + 3*npts] = x35*y*(radial_eval_alpha_squared*x19 - x63); + basis_xy_eval[ipt + 4*npts] = x11*(-x2*x71 + x32 + x80); + basis_xy_eval[ipt + 5*npts] = radial_eval_alpha_squared*x23*x40*y; + basis_xy_eval[ipt + 6*npts] = x1*(-x2*x76 + x36 + x81); + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x82*(radial_eval_alpha_squared*x6 + x76); + basis_xz_eval[ipt + 1*npts] = x53*(radial_eval_alpha_squared*x12*x2 + x29 + x54); + basis_xz_eval[ipt + 2*npts] = x83; + basis_xz_eval[ipt + 3*npts] = 0.5*x*x84; + basis_xz_eval[ipt + 4*npts] = x85*(x2*x86 + x55 + x80); + basis_xz_eval[ipt + 5*npts] = x*x21*(x42 + x87 + x88); + basis_xz_eval[ipt + 6*npts] = x81*x89; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x1*(x36 + x92); + basis_yy_eval[ipt + 1*npts] = x9*(x66 + x90); + basis_yy_eval[ipt + 2*npts] = x11*(x36 + x94); + basis_yy_eval[ipt + 3*npts] = x16*(x19*x91 - x96); + basis_yy_eval[ipt + 4*npts] = x20*(x15*x91 - x98); + basis_yy_eval[ipt + 5*npts] = x22*(x23*x91 - x98); + basis_yy_eval[ipt + 6*npts] = x24*(x26*x91 - x96); + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x78*x89; + basis_yz_eval[ipt + 1*npts] = sqrt_15*x*(radial_eval_alpha_squared*x12*x4 + x48 + x54); + basis_yz_eval[ipt + 2*npts] = x85*(x4*x86 + x55 + x79); + basis_yz_eval[ipt + 3*npts] = 0.5*x84*y; + basis_yz_eval[ipt + 4*npts] = x83; + basis_yz_eval[ipt + 5*npts] = x21*y*(x51 - x87 + x88); + basis_yz_eval[ipt + 6*npts] = x82*(radial_eval_alpha_squared*x26 - x76); + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x1*x101; + basis_zz_eval[ipt + 1*npts] = x9*(x66 + x99); + basis_zz_eval[ipt + 2*npts] = x103*x11; + basis_zz_eval[ipt + 3*npts] = x16*(12.0*radial_eval + x104); + basis_zz_eval[ipt + 4*npts] = x103*x20; + basis_zz_eval[ipt + 5*npts] = x105*x22; + basis_zz_eval[ipt + 6*npts] = x106*x24; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x1*x6; + ang_eval_1 = radial_eval*x9; + ang_eval_2 = radial_eval*x11*x15; + ang_eval_3 = radial_eval*x16*x19; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x15*x20; + ang_eval_1 = radial_eval*x22*x23; + ang_eval_2 = radial_eval*x24*x26; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x27*(radial_eval_alpha*x6 + x28); + dang_eval_y_0 = x0*(radial_eval_alpha*x46 + x44); + dang_eval_z_0 = x1*x52*x6; + dang_eval_x_1 = x29*x8; + dang_eval_y_1 = x*x48*x7; + dang_eval_z_1 = x*x53*(radial_eval + x54); + dang_eval_x_2 = x34; + dang_eval_y_2 = x10*(radial_eval*x49 + radial_eval_alpha*x50); + dang_eval_z_2 = x11*x56*z; + dang_eval_x_3 = x35*x37; + dang_eval_y_3 = x16*x37*y; + dang_eval_z_3 = 1.5*radial_eval*x57 + 0.5*radial_eval_alpha*x58; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x10*(radial_eval*x38 + radial_eval_alpha*x39); + dang_eval_y_0 = x34; + dang_eval_z_0 = x20*x56*z; + dang_eval_x_1 = x40*x42; + dang_eval_y_1 = x22*x51*y; + dang_eval_z_1 = x21*(radial_eval*x23 + radial_eval_alpha*x59); + dang_eval_x_2 = x0*(radial_eval_alpha*x45 + x44); + dang_eval_y_2 = x27*(radial_eval_alpha*x26 + x36); + dang_eval_z_2 = x24*x26*x52; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_lapgrad.hpp new file mode 100644 index 000000000..a8edeb732 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_lapgrad.hpp @@ -0,0 +1,521 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_3_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_3_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_lapgrad_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = 0.25*sqrt_10; + const auto x1 = x0*y; + const auto x2 = x*x; + const auto x3 = x2; + const auto x4 = 3.0*x3; + const auto x5 = y*y; + const auto x6 = x5; + const auto x7 = -x6; + const auto x8 = x4 + x7; + const auto x9 = sqrt_15*z; + const auto x10 = x9*y; + const auto x11 = x*x10; + const auto x12 = 0.25*sqrt_6; + const auto x13 = x12*y; + const auto x14 = z*z; + const auto x15 = x14; + const auto x16 = -4.0*x15; + const auto x17 = x16 + x6; + const auto x18 = -x17 - x3; + const auto x19 = 0.5*z; + const auto x20 = 3.0*x6; + const auto x21 = -2.0*x15; + const auto x22 = -x20 - x21 - x4; + const auto x23 = x*x12; + const auto x24 = 0.5*sqrt_15; + const auto x25 = x24*z; + const auto x26 = x3 + x7; + const auto x27 = x*x0; + const auto x28 = -x20; + const auto x29 = x28 + x3; + const auto x30 = x*x1; + const auto x31 = 6.0*radial_eval; + const auto x32 = radial_eval + radial_eval_alpha*x3; + const auto x33 = x*x13; + const auto x34 = 2.0*radial_eval; + const auto x35 = -x34; + const auto x36 = radial_eval_alpha*x18; + const auto x37 = x33*(x35 + x36); + const auto x38 = x*x19; + const auto x39 = -x31; + const auto x40 = radial_eval_alpha*x22 + x39; + const auto x41 = -x17 - x4; + const auto x42 = x18*x3; + const auto x43 = x*x25; + const auto x44 = radial_eval_alpha*x26; + const auto x45 = x34 + x44; + const auto x46 = x28 + x4; + const auto x47 = radial_eval*x46; + const auto x48 = x29*x3; + const auto x49 = x6*x8; + const auto x50 = x*x9; + const auto x51 = radial_eval_alpha*x6; + const auto x52 = radial_eval + x51; + const auto x53 = -x16 - x20 - x3; + const auto x54 = x18*x6; + const auto x55 = x35 + x44; + const auto x56 = radial_eval_alpha*z; + const auto x57 = sqrt_15*y; + const auto x58 = x*x57; + const auto x59 = radial_eval_alpha*x15; + const auto x60 = 8.0*radial_eval; + const auto x61 = x36 + x60; + const auto x62 = -x21 - x3 - x6; + const auto x63 = x15*x22; + const auto x64 = x15*x26; + const auto x65 = radial_eval_alpha_squared*x3; + const auto x66 = radial_eval_alpha + x65; + const auto x67 = x66*x8; + const auto x68 = 12.0*radial_eval_alpha; + const auto x69 = x3*x68; + const auto x70 = x31 + x69; + const auto x71 = 3.0*radial_eval_alpha; + const auto x72 = 4.0*radial_eval_alpha; + const auto x73 = x3*x72; + const auto x74 = x34 + x73; + const auto x75 = x18*x66; + const auto x76 = 2.0*radial_eval_alpha; + const auto x77 = x41*x76 + x75; + const auto x78 = x26*x66; + const auto x79 = x46*x76; + const auto x80 = x29*x66 + x79; + const auto x81 = 6.0*radial_eval_alpha; + const auto x82 = radial_eval_alpha*x46; + const auto x83 = radial_eval_alpha_squared*x49 + x82; + const auto x84 = x3*x6; + const auto x85 = radial_eval_alpha*x53 + radial_eval_alpha_squared*x54; + const auto x86 = radial_eval_alpha*x41 + radial_eval_alpha_squared*x42; + const auto x87 = radial_eval_alpha_squared*x48 + x82; + const auto x88 = x30*z; + const auto x89 = x15*x3; + const auto x90 = x33*z*(radial_eval_alpha_squared*x18 + x81); + const auto x91 = radial_eval_alpha_squared*x63 - x15*x81 + x39 + x62*x71; + const auto x92 = x12*z; + const auto x93 = 8.0*radial_eval_alpha; + const auto x94 = x15*x76; + const auto x95 = radial_eval_alpha_squared*x64; + const auto x96 = x0*z; + const auto x97 = radial_eval_alpha_squared*x6; + const auto x98 = radial_eval_alpha + x97; + const auto x99 = x79 + x8*x98; + const auto x100 = x18*x98; + const auto x101 = x100 + x53*x76; + const auto x102 = x6*x68; + const auto x103 = x102 + x31; + const auto x104 = x6*x72; + const auto x105 = x104 + x34; + const auto x106 = x15*x6; + const auto x107 = radial_eval_alpha_squared*x15; + const auto x108 = radial_eval_alpha + x107; + const auto x109 = x108*x8; + const auto x110 = 16.0*radial_eval_alpha*x15; + const auto x111 = x108*x18 + x110; + const auto x112 = x111 + x60; + const auto x113 = x108*x22 + x62*x81; + const auto x114 = x108*x26; + const auto x115 = x114 + x26*x76; + const auto x116 = x108*x29; + const auto x117 = x107 + x97; + const auto x118 = -x73; + const auto x119 = -x102; + const auto x120 = -x69; + const auto x121 = x119 + x120; + const auto x122 = -x104; + const auto x123 = x122 + x26*x98 + x73 + x78; + const auto x124 = 3.0*radial_eval_alpha_squared; + const auto x125 = x*(radial_eval_alpha_cubed*(x*x) + x124); + const auto x126 = radial_eval_alpha_cubed*x6 + radial_eval_alpha_squared; + const auto x127 = x126*x8; + const auto x128 = radial_eval_alpha_cubed*x15 + radial_eval_alpha_squared; + const auto x129 = x128*x8; + const auto x130 = 2.0*x; + const auto x131 = radial_eval_alpha_squared*x130; + const auto x132 = 6.0*x; + const auto x133 = 24.0*radial_eval_alpha; + const auto x134 = x*x133 + 18.0*x*x66 + x108*x132 + x132*x98; + const auto x135 = 4.0*radial_eval_alpha_squared; + const auto x136 = x*x93; + const auto x137 = 16.0*radial_eval_alpha_squared; + const auto x138 = x132*x66; + const auto x139 = x130*x98; + const auto x140 = x108*x130; + const auto x141 = x126*x18; + const auto x142 = x128*x18; + const auto x143 = x125*x18; + const auto x144 = 12.0*radial_eval_alpha_squared; + const auto x145 = x110 - x135*x84; + const auto x146 = x126*x26; + const auto x147 = x128*x26; + const auto x148 = x46*x98; + const auto x149 = x46*x66; + const auto x150 = x126*x29; + const auto x151 = x128*x29; + const auto x152 = x144*x84; + const auto x153 = x108*x46 + x119 + x69; + const auto x154 = y*(radial_eval_alpha_cubed*(y*y) + x124); + const auto x155 = radial_eval_alpha_cubed*x3 + radial_eval_alpha_squared; + const auto x156 = x155*x8; + const auto x157 = x65 + x81; + const auto x158 = x154*x18; + const auto x159 = x155*x18; + const auto x160 = x133*y; + const auto x161 = 6.0*y; + const auto x162 = x161*x66; + const auto x163 = 18.0*x98*y; + const auto x164 = x108*x161; + const auto x165 = 2.0*y; + const auto x166 = radial_eval_alpha_squared*x165; + const auto x167 = -x108*x165 - x161*x98 - x165*x66 - x93*y; + const auto x168 = x155*x26; + const auto x169 = x155*x29; + const auto x170 = x144*z; + const auto x171 = 2.0*radial_eval_alpha_squared*z; + const auto x172 = x171*x46; + const auto x173 = z*(radial_eval_alpha_cubed*(z*z) + x124); + const auto x174 = x135*z; + const auto x175 = 8.0*z; + const auto x176 = 24.0*x108*z + x141*z + x159*z + x173*x18 + x175*x66 + x175*x98 + 32.0*x56; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x1*x8; + basis_eval[ipt + 1*npts] = radial_eval*x11; + basis_eval[ipt + 2*npts] = radial_eval*x13*x18; + basis_eval[ipt + 3*npts] = radial_eval*x19*x22; + basis_eval[ipt + 4*npts] = radial_eval*x18*x23; + basis_eval[ipt + 5*npts] = radial_eval*x25*x26; + basis_eval[ipt + 6*npts] = radial_eval*x27*x29; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x30*(radial_eval_alpha*x8 + x31); + basis_x_eval[ipt + 1*npts] = x10*x32; + basis_x_eval[ipt + 2*npts] = x37; + basis_x_eval[ipt + 3*npts] = x38*x40; + basis_x_eval[ipt + 4*npts] = x12*(radial_eval*x41 + radial_eval_alpha*x42); + basis_x_eval[ipt + 5*npts] = x43*x45; + basis_x_eval[ipt + 6*npts] = x0*(radial_eval_alpha*x48 + x47); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*(radial_eval_alpha*x49 + x47); + basis_y_eval[ipt + 1*npts] = x50*x52; + basis_y_eval[ipt + 2*npts] = x12*(radial_eval*x53 + radial_eval_alpha*x54); + basis_y_eval[ipt + 3*npts] = x19*x40*y; + basis_y_eval[ipt + 4*npts] = x37; + basis_y_eval[ipt + 5*npts] = x25*x55*y; + basis_y_eval[ipt + 6*npts] = x30*(radial_eval_alpha*x29 + x39); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x1*x56*x8; + basis_z_eval[ipt + 1*npts] = x58*(radial_eval + x59); + basis_z_eval[ipt + 2*npts] = x13*x61*z; + basis_z_eval[ipt + 3*npts] = 1.5*radial_eval*x62 + 0.5*radial_eval_alpha*x63; + basis_z_eval[ipt + 4*npts] = x23*x61*z; + basis_z_eval[ipt + 5*npts] = x24*(radial_eval*x26 + radial_eval_alpha*x64); + basis_z_eval[ipt + 6*npts] = x27*x29*x56; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x1*(x67 + x70); + basis_xx_eval[ipt + 1*npts] = x11*(x65 + x71); + basis_xx_eval[ipt + 2*npts] = x13*(x18*x66 - x74); + basis_xx_eval[ipt + 3*npts] = x19*(x22*x66 - x70); + basis_xx_eval[ipt + 4*npts] = x23*(x39 + x77); + basis_xx_eval[ipt + 5*npts] = x25*(x74 + x78); + basis_xx_eval[ipt + 6*npts] = x27*(x31 + x80); + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x27*(x31 + x6*x81 + x83); + basis_xy_eval[ipt + 1*npts] = x9*(radial_eval_alpha_squared*x84 + x32 + x51); + basis_xy_eval[ipt + 2*npts] = x23*(x35 - x6*x76 + x85); + basis_xy_eval[ipt + 3*npts] = x38*y*(radial_eval_alpha_squared*x22 - x68); + basis_xy_eval[ipt + 4*npts] = x13*(-x3*x76 + x35 + x86); + basis_xy_eval[ipt + 5*npts] = radial_eval_alpha_squared*x26*x43*y; + basis_xy_eval[ipt + 6*npts] = x1*(-x3*x81 + x39 + x87); + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x88*(radial_eval_alpha_squared*x8 + x81); + basis_xz_eval[ipt + 1*npts] = x57*(radial_eval_alpha_squared*x89 + x32 + x59); + basis_xz_eval[ipt + 2*npts] = x90; + basis_xz_eval[ipt + 3*npts] = 0.5*x*x91; + basis_xz_eval[ipt + 4*npts] = x92*(x3*x93 + x60 + x86); + basis_xz_eval[ipt + 5*npts] = x*x24*(x45 + x94 + x95); + basis_xz_eval[ipt + 6*npts] = x87*x96; + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x1*(x39 + x99); + basis_yy_eval[ipt + 1*npts] = x11*(x71 + x97); + basis_yy_eval[ipt + 2*npts] = x13*(x101 + x39); + basis_yy_eval[ipt + 3*npts] = x19*(-x103 + x22*x98); + basis_yy_eval[ipt + 4*npts] = x23*(-x105 + x18*x98); + basis_yy_eval[ipt + 5*npts] = x25*(-x105 + x26*x98); + basis_yy_eval[ipt + 6*npts] = x27*(-x103 + x29*x98); + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x83*x96; + basis_yz_eval[ipt + 1*npts] = sqrt_15*x*(radial_eval_alpha_squared*x106 + x52 + x59); + basis_yz_eval[ipt + 2*npts] = x92*(x6*x93 + x60 + x85); + basis_yz_eval[ipt + 3*npts] = 0.5*x91*y; + basis_yz_eval[ipt + 4*npts] = x90; + basis_yz_eval[ipt + 5*npts] = x24*y*(x55 - x94 + x95); + basis_yz_eval[ipt + 6*npts] = x88*(radial_eval_alpha_squared*x29 - x81); + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x1*x109; + basis_zz_eval[ipt + 1*npts] = x11*(x107 + x71); + basis_zz_eval[ipt + 2*npts] = x112*x13; + basis_zz_eval[ipt + 3*npts] = x19*(12.0*radial_eval + x113); + basis_zz_eval[ipt + 4*npts] = x112*x23; + basis_zz_eval[ipt + 5*npts] = x115*x25; + basis_zz_eval[ipt + 6*npts] = x116*x27; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x1*(x109 + x67 + x69 + x99); + basis_lapl_eval[ipt + 1*npts] = x11*(9.0*radial_eval_alpha + x117 + x65); + basis_lapl_eval[ipt + 2*npts] = x13*(x101 + x111 + x118 + x75); + basis_lapl_eval[ipt + 3*npts] = x19*(x113 + x121 + x22*x66 + x22*x98); + basis_lapl_eval[ipt + 4*npts] = x23*(x100 + x111 + x122 + x77); + basis_lapl_eval[ipt + 5*npts] = x25*(x115 + x123); + basis_lapl_eval[ipt + 6*npts] = x27*(x116 + x119 + x29*x98 + x80); + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = x1*(x*x127 + x*x129 + x125*x8 + x131*x46 + x134); + basis_lapl_x_eval[ipt + 1*npts] = x10*(x*x125 + x117 + x126*x2 + x128*x2 + x135*x3 + 3.0*x66 + x81); + basis_lapl_x_eval[ipt + 2*npts] = x13*(x*x137*x15 + x*x141 + x*x142 + x131*x53 - x136 - x138 - x139 - x140 + x143); + basis_lapl_x_eval[ipt + 3*npts] = x19*(6.0*radial_eval_alpha_squared*x*x62 + x*x126*x22 + x*x128*x22 - x*x144*x6 + x125*x22 - x134); + basis_lapl_x_eval[ipt + 4*npts] = x12*(x*x143 + x108*x41 + x120 + x122 + x137*x89 + x141*x2 + x142*x2 + x145 + 3.0*x41*x66 + x41*x98); + basis_lapl_x_eval[ipt + 5*npts] = x25*(-x*x135*x6 + x*x146 + x*x147 + x125*x26 + x131*x26 + x136 + x138 + x139 + x140); + basis_lapl_x_eval[ipt + 6*npts] = x0*(x*x125*x29 + x148 + 3.0*x149 + x150*x2 + x151*x2 - x152 + x153); + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x0*(x129*x5 + 3.0*x148 + x149 + x152 + x153 + x154*x8*y + x156*x5); + basis_lapl_y_eval[ipt + 1*npts] = x50*(x107 + x128*x5 + x135*x6 + x154*y + x155*x5 + x157 + 3.0*x98); + basis_lapl_y_eval[ipt + 2*npts] = x12*(x106*x137 + x108*x53 + x118 + x119 + x142*x5 + x145 + x158*y + x159*x5 + x53*x66 + 3.0*x53*x98); + basis_lapl_y_eval[ipt + 3*npts] = -x19*(-6.0*radial_eval_alpha_squared*x62*y - x128*x22*y + x144*x3*y - x154*x22 - x155*x22*y + x160 + x162 + x163 + x164); + basis_lapl_y_eval[ipt + 4*npts] = x23*(x137*x15*y + x142*y + x158 + x159*y + x166*x41 + x167); + basis_lapl_y_eval[ipt + 5*npts] = x25*(x135*x3*y + x147*y + x154*x26 + x166*x26 + x167 + x168*y); + basis_lapl_y_eval[ipt + 6*npts] = x27*(x151*y + x154*x29 - x160 - x162 - x163 - x164 + x166*x46 + x169*y); + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x1*(x127*z + x156*z + x170*x3 + x172 + x173*x8); + basis_lapl_z_eval[ipt + 1*npts] = x58*(3.0*x108 + x126*x14 + x135*x15 + x14*x155 + x157 + x173*z + x97); + basis_lapl_z_eval[ipt + 2*npts] = x13*(x171*x53 - x174*x3 + x176); + basis_lapl_z_eval[ipt + 3*npts] = -0.5*x106*x144 + 4.5*x108*x62 + 0.5*x121 + 0.5*x126*x14*x22 + 0.5*x133*x15 + 0.5*x14*x155*x22 - 0.5*x144*x89 + 0.5*x173*x22*z + 1.5*x62*x66 + 1.5*x62*x98; + basis_lapl_z_eval[ipt + 4*npts] = x23*(x171*x41 - x174*x6 + x176); + basis_lapl_z_eval[ipt + 5*npts] = x24*(-x106*x135 + 3.0*x114 + x123 + x135*x89 + x14*x146 + x14*x168 + x173*x26*z); + basis_lapl_z_eval[ipt + 6*npts] = x27*(x150*z + x169*z - x170*x6 + x172 + x173*x29); + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x1*x8; + ang_eval_1 = radial_eval*x11; + ang_eval_2 = radial_eval*x13*x18; + ang_eval_3 = radial_eval*x19*x22; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x18*x23; + ang_eval_1 = radial_eval*x25*x26; + ang_eval_2 = radial_eval*x27*x29; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x30*(radial_eval_alpha*x8 + x31); + dang_eval_y_0 = x0*(radial_eval_alpha*x49 + x47); + dang_eval_z_0 = x1*x56*x8; + dang_eval_x_1 = x10*x32; + dang_eval_y_1 = x50*x52; + dang_eval_z_1 = x58*(radial_eval + x59); + dang_eval_x_2 = x37; + dang_eval_y_2 = x12*(radial_eval*x53 + radial_eval_alpha*x54); + dang_eval_z_2 = x13*x61*z; + dang_eval_x_3 = x38*x40; + dang_eval_y_3 = x19*x40*y; + dang_eval_z_3 = 1.5*radial_eval*x62 + 0.5*radial_eval_alpha*x63; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x12*(radial_eval*x41 + radial_eval_alpha*x42); + dang_eval_y_0 = x37; + dang_eval_z_0 = x23*x61*z; + dang_eval_x_1 = x43*x45; + dang_eval_y_1 = x25*x55*y; + dang_eval_z_1 = x24*(radial_eval*x26 + radial_eval_alpha*x64); + dang_eval_x_2 = x0*(radial_eval_alpha*x48 + x47); + dang_eval_y_2 = x30*(radial_eval_alpha*x29 + x39); + dang_eval_z_2 = x27*x29*x56; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_laplacian.hpp new file mode 100644 index 000000000..28cc5a7a9 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l3_laplacian.hpp @@ -0,0 +1,364 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_3_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_3_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_laplacian_3( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = 0.25*sqrt_10; + const auto x1 = x0*y; + const auto x2 = x*x; + const auto x3 = 3.0*x2; + const auto x4 = y*y; + const auto x5 = -x4; + const auto x6 = x3 + x5; + const auto x7 = sqrt_15*z; + const auto x8 = x7*y; + const auto x9 = x*x8; + const auto x10 = 0.25*sqrt_6; + const auto x11 = x10*y; + const auto x12 = z*z; + const auto x13 = -4.0*x12; + const auto x14 = x13 + x4; + const auto x15 = -x14 - x2; + const auto x16 = 0.5*z; + const auto x17 = 3.0*x4; + const auto x18 = -2.0*x12; + const auto x19 = -x17 - x18 - x3; + const auto x20 = x*x10; + const auto x21 = 0.5*sqrt_15; + const auto x22 = x21*z; + const auto x23 = x2 + x5; + const auto x24 = x*x0; + const auto x25 = -x17; + const auto x26 = x2 + x25; + const auto x27 = x*x1; + const auto x28 = 6.0*radial_eval; + const auto x29 = radial_eval + radial_eval_alpha*x2; + const auto x30 = x*x11; + const auto x31 = 2.0*radial_eval; + const auto x32 = -x31; + const auto x33 = radial_eval_alpha*x15; + const auto x34 = x30*(x32 + x33); + const auto x35 = x*x16; + const auto x36 = -x28; + const auto x37 = radial_eval_alpha*x19 + x36; + const auto x38 = -x14 - x3; + const auto x39 = x15*x2; + const auto x40 = x*x22; + const auto x41 = radial_eval_alpha*x23; + const auto x42 = x31 + x41; + const auto x43 = x25 + x3; + const auto x44 = radial_eval*x43; + const auto x45 = x2*x26; + const auto x46 = x4*x6; + const auto x47 = radial_eval_alpha*x4; + const auto x48 = radial_eval + x47; + const auto x49 = -x13 - x17 - x2; + const auto x50 = x15*x4; + const auto x51 = x32 + x41; + const auto x52 = radial_eval_alpha*z; + const auto x53 = sqrt_15*y; + const auto x54 = radial_eval_alpha*x12; + const auto x55 = 8.0*radial_eval; + const auto x56 = x33 + x55; + const auto x57 = -x18 - x2 - x4; + const auto x58 = x12*x19; + const auto x59 = x12*x23; + const auto x60 = radial_eval_alpha_squared*x2; + const auto x61 = radial_eval_alpha + x60; + const auto x62 = x6*x61; + const auto x63 = 12.0*radial_eval_alpha; + const auto x64 = x2*x63; + const auto x65 = x28 + x64; + const auto x66 = 3.0*radial_eval_alpha; + const auto x67 = 4.0*radial_eval_alpha; + const auto x68 = x2*x67; + const auto x69 = x31 + x68; + const auto x70 = x15*x61; + const auto x71 = 2.0*radial_eval_alpha; + const auto x72 = x38*x71 + x70; + const auto x73 = x23*x61; + const auto x74 = x43*x71; + const auto x75 = x26*x61 + x74; + const auto x76 = 6.0*radial_eval_alpha; + const auto x77 = radial_eval_alpha*x43; + const auto x78 = radial_eval_alpha_squared*x46 + x77; + const auto x79 = radial_eval_alpha*x49 + radial_eval_alpha_squared*x50; + const auto x80 = radial_eval_alpha*x38 + radial_eval_alpha_squared*x39; + const auto x81 = radial_eval_alpha_squared*x45 + x77; + const auto x82 = x27*z; + const auto x83 = x30*z*(radial_eval_alpha_squared*x15 + x76); + const auto x84 = radial_eval_alpha_squared*x58 - x12*x76 + x36 + x57*x66; + const auto x85 = x10*z; + const auto x86 = 8.0*radial_eval_alpha; + const auto x87 = x12*x71; + const auto x88 = radial_eval_alpha_squared*x59; + const auto x89 = x0*z; + const auto x90 = radial_eval_alpha_squared*x4; + const auto x91 = radial_eval_alpha + x90; + const auto x92 = x6*x91 + x74; + const auto x93 = x15*x91; + const auto x94 = x49*x71 + x93; + const auto x95 = x4*x63; + const auto x96 = x28 + x95; + const auto x97 = x4*x67; + const auto x98 = x31 + x97; + const auto x99 = radial_eval_alpha_squared*x12; + const auto x100 = radial_eval_alpha + x99; + const auto x101 = x100*x6; + const auto x102 = 16.0*radial_eval_alpha*x12 + x100*x15; + const auto x103 = x102 + x55; + const auto x104 = x100*x19 + x57*x76; + const auto x105 = x23*(x100 + x71); + const auto x106 = x100*x26; + const auto x107 = -x95; + const auto x108 = -x97; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x1*x6; + basis_eval[ipt + 1*npts] = radial_eval*x9; + basis_eval[ipt + 2*npts] = radial_eval*x11*x15; + basis_eval[ipt + 3*npts] = radial_eval*x16*x19; + basis_eval[ipt + 4*npts] = radial_eval*x15*x20; + basis_eval[ipt + 5*npts] = radial_eval*x22*x23; + basis_eval[ipt + 6*npts] = radial_eval*x24*x26; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x27*(radial_eval_alpha*x6 + x28); + basis_x_eval[ipt + 1*npts] = x29*x8; + basis_x_eval[ipt + 2*npts] = x34; + basis_x_eval[ipt + 3*npts] = x35*x37; + basis_x_eval[ipt + 4*npts] = x10*(radial_eval*x38 + radial_eval_alpha*x39); + basis_x_eval[ipt + 5*npts] = x40*x42; + basis_x_eval[ipt + 6*npts] = x0*(radial_eval_alpha*x45 + x44); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x0*(radial_eval_alpha*x46 + x44); + basis_y_eval[ipt + 1*npts] = x*x48*x7; + basis_y_eval[ipt + 2*npts] = x10*(radial_eval*x49 + radial_eval_alpha*x50); + basis_y_eval[ipt + 3*npts] = x16*x37*y; + basis_y_eval[ipt + 4*npts] = x34; + basis_y_eval[ipt + 5*npts] = x22*x51*y; + basis_y_eval[ipt + 6*npts] = x27*(radial_eval_alpha*x26 + x36); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x1*x52*x6; + basis_z_eval[ipt + 1*npts] = x*x53*(radial_eval + x54); + basis_z_eval[ipt + 2*npts] = x11*x56*z; + basis_z_eval[ipt + 3*npts] = 1.5*radial_eval*x57 + 0.5*radial_eval_alpha*x58; + basis_z_eval[ipt + 4*npts] = x20*x56*z; + basis_z_eval[ipt + 5*npts] = x21*(radial_eval*x23 + radial_eval_alpha*x59); + basis_z_eval[ipt + 6*npts] = x24*x26*x52; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x1*(x101 + x62 + x64 + x92); + basis_lapl_eval[ipt + 1*npts] = x9*(9.0*radial_eval_alpha + x60 + x90 + x99); + basis_lapl_eval[ipt + 2*npts] = x11*(x102 - x68 + x70 + x94); + basis_lapl_eval[ipt + 3*npts] = x16*(x104 + x107 + x19*x61 + x19*x91 - x64); + basis_lapl_eval[ipt + 4*npts] = x20*(x102 + x108 + x72 + x93); + basis_lapl_eval[ipt + 5*npts] = x22*(x105 + x108 + x23*x91 + x68 + x73); + basis_lapl_eval[ipt + 6*npts] = x24*(x106 + x107 + x26*x91 + x75); + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x1*x6; + ang_eval_1 = radial_eval*x9; + ang_eval_2 = radial_eval*x11*x15; + ang_eval_3 = radial_eval*x16*x19; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = radial_eval*x15*x20; + ang_eval_1 = radial_eval*x22*x23; + ang_eval_2 = radial_eval*x24*x26; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x27*(radial_eval_alpha*x6 + x28); + dang_eval_y_0 = x0*(radial_eval_alpha*x46 + x44); + dang_eval_z_0 = x1*x52*x6; + dang_eval_x_1 = x29*x8; + dang_eval_y_1 = x*x48*x7; + dang_eval_z_1 = x*x53*(radial_eval + x54); + dang_eval_x_2 = x34; + dang_eval_y_2 = x10*(radial_eval*x49 + radial_eval_alpha*x50); + dang_eval_z_2 = x11*x56*z; + dang_eval_x_3 = x35*x37; + dang_eval_y_3 = x16*x37*y; + dang_eval_z_3 = 1.5*radial_eval*x57 + 0.5*radial_eval_alpha*x58; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = x10*(radial_eval*x38 + radial_eval_alpha*x39); + dang_eval_y_0 = x34; + dang_eval_z_0 = x20*x56*z; + dang_eval_x_1 = x40*x42; + dang_eval_y_1 = x22*x51*y; + dang_eval_z_1 = x21*(radial_eval*x23 + radial_eval_alpha*x59); + dang_eval_x_2 = x0*(radial_eval_alpha*x45 + x44); + dang_eval_y_2 = x27*(radial_eval_alpha*x26 + x36); + dang_eval_z_2 = x24*x26*x52; + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4.hpp new file mode 100644 index 000000000..0be6a8e38 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4.hpp @@ -0,0 +1,187 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_4_max_wg = 512; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_4_nwarp = 16; + + +void collocation_device_shell_to_task_kernel_spherical_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + } + + + // Common Subexpressions + const auto x0 = 0.5*radial_eval*x*y; + const auto x1 = x*x; + const auto x2 = y*y; + const auto x3 = -x2; + const auto x4 = 0.25*radial_eval; + const auto x5 = x4*z; + const auto x6 = x5*y; + const auto x7 = 3.0*x1; + const auto x8 = z*z; + const auto x9 = 3.0*x2; + const auto x10 = -x7 + 4.0*x8 - x9; + const auto x11 = 0.125*radial_eval; + const auto x12 = x*x*x*x; + const auto x13 = y*y*y*y; + const auto x14 = 6.0*x1*x2; + const auto x15 = x1*x8; + const auto x16 = x2*x8; + const auto x17 = x*x5; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = sqrt_35*x0*(x1 + x3); + basis_eval[ipt + 1*npts] = sqrt_70*x6*(x3 + x7); + basis_eval[ipt + 2*npts] = -sqrt_5*x0*(x1 + x2 - 6.0*x8); + basis_eval[ipt + 3*npts] = sqrt_10*x10*x6; + basis_eval[ipt + 4*npts] = x11*(3.0*x12 + 3.0*x13 + x14 - 24.0*x15 - 24.0*x16 + 8.0*(z*z*z*z)); + basis_eval[ipt + 5*npts] = sqrt_10*x10*x17; + basis_eval[ipt + 6*npts] = -sqrt_5*x4*(x12 - x13 - 6.0*x15 + 6.0*x16); + basis_eval[ipt + 7*npts] = sqrt_70*x17*(x1 - x9); + basis_eval[ipt + 8*npts] = sqrt_35*x11*(x12 + x13 - x14); + + + + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = sqrt_35*x0*(x1 + x3); + ang_eval_1 = sqrt_70*x6*(x3 + x7); + ang_eval_2 = -sqrt_5*x0*(x1 + x2 - 6.0*x8); + ang_eval_3 = sqrt_10*x10*x6; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x11*(3.0*x12 + 3.0*x13 + x14 - 24.0*x15 - 24.0*x16 + 8.0*(z*z*z*z)); + ang_eval_1 = sqrt_10*x10*x17; + ang_eval_2 = -sqrt_5*x4*(x12 - x13 - 6.0*x15 + 6.0*x16); + ang_eval_3 = sqrt_70*x17*(x1 - x9); + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = sqrt_35*x11*(x12 + x13 - x14); + basis_eval[ipt + 8*npts] = ang_eval_0; + + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_gradient.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_gradient.hpp new file mode 100644 index 000000000..6bcd8ec09 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_gradient.hpp @@ -0,0 +1,326 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_gradient_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_gradient_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + } + + radial_eval_alpha *= -2; + + // Common Subexpressions + const auto x0 = 0.5*y; + const auto x1 = sqrt_35*x0; + const auto x2 = radial_eval*x; + const auto x3 = x*x; + const auto x4 = y*y; + const auto x5 = -x4; + const auto x6 = x3 + x5; + const auto x7 = 0.25*z; + const auto x8 = sqrt_70*x7; + const auto x9 = radial_eval*y; + const auto x10 = 3.0*x3; + const auto x11 = x10 + x5; + const auto x12 = sqrt_5*x0; + const auto x13 = z*z; + const auto x14 = -6.0*x13; + const auto x15 = x14 + x4; + const auto x16 = -x15 - x3; + const auto x17 = sqrt_10*x7; + const auto x18 = -4.0*x13; + const auto x19 = 3.0*x4; + const auto x20 = x18 + x19; + const auto x21 = -x10 - x20; + const auto x22 = 0.125*radial_eval; + const auto x23 = x*x*x*x; + const auto x24 = y*y*y*y; + const auto x25 = 6.0*x3*x4; + const auto x26 = x13*x3; + const auto x27 = x13*x4; + const auto x28 = 3.0*x23 + 3.0*x24 + x25 - 24.0*x26 - 24.0*x27 + 8.0*(z*z*z*z); + const auto x29 = 0.25*sqrt_5; + const auto x30 = -x23 + x24 + 6.0*x26 - 6.0*x27; + const auto x31 = -x19; + const auto x32 = x3 + x31; + const auto x33 = x23 + x24 - x25; + const auto x34 = radial_eval*x11; + const auto x35 = x*y; + const auto x36 = x35*x8; + const auto x37 = 6.0*radial_eval; + const auto x38 = -x37; + const auto x39 = x17*x35*(radial_eval_alpha*x21 + x38); + const auto x40 = 12.0*radial_eval; + const auto x41 = x*x*x; + const auto x42 = radial_eval_alpha*x; + const auto x43 = 4.0*radial_eval; + const auto x44 = 3.0*x; + const auto x45 = radial_eval*(x10 + x31); + const auto x46 = 0.125*sqrt_35; + const auto x47 = 0.5*x; + const auto x48 = radial_eval*x32; + const auto x49 = y*y*y; + const auto x50 = radial_eval_alpha*y; + const auto x51 = 3.0*y; + const auto x52 = 0.25*y; + const auto x53 = -radial_eval*(x10 - 12.0*x13 + x19) + radial_eval_alpha*x13*x21; + const auto x54 = 3.0*z; + const auto x55 = radial_eval_alpha*z; + const auto x56 = 0.25*x; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = x1*x2*x6; + basis_eval[ipt + 1*npts] = x11*x8*x9; + basis_eval[ipt + 2*npts] = x12*x16*x2; + basis_eval[ipt + 3*npts] = x17*x21*x9; + basis_eval[ipt + 4*npts] = x22*x28; + basis_eval[ipt + 5*npts] = x17*x2*x21; + basis_eval[ipt + 6*npts] = radial_eval*x29*x30; + basis_eval[ipt + 7*npts] = x2*x32*x8; + basis_eval[ipt + 8*npts] = sqrt_35*x22*x33; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1*(radial_eval_alpha*x3*x6 + x34); + basis_x_eval[ipt + 1*npts] = x36*(radial_eval_alpha*x11 + x37); + basis_x_eval[ipt + 2*npts] = -x12*(radial_eval*(x10 + x15) - radial_eval_alpha*x16*x3); + basis_x_eval[ipt + 3*npts] = x39; + basis_x_eval[ipt + 4*npts] = 0.125*x28*x42 + 0.125*x40*(-4.0*x*x13 + x*x4 + x41); + basis_x_eval[ipt + 5*npts] = -x17*(radial_eval*(x20 + 9.0*x3) - radial_eval_alpha*x21*x3); + basis_x_eval[ipt + 6*npts] = x29*(x30*x42 + x43*(x13*x44 - x41)); + basis_x_eval[ipt + 7*npts] = x8*(radial_eval_alpha*x3*x32 + x45); + basis_x_eval[ipt + 8*npts] = x46*(x33*x42 - x43*(x4*x44 - x41)); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = sqrt_35*x47*(radial_eval_alpha*x4*x6 + x48); + basis_y_eval[ipt + 1*npts] = x8*(radial_eval_alpha*x11*x4 + x45); + basis_y_eval[ipt + 2*npts] = -sqrt_5*x47*(radial_eval*(x14 + x19 + x3) - radial_eval_alpha*x16*x4); + basis_y_eval[ipt + 3*npts] = -x17*(radial_eval*(x10 + x18 + 9.0*x4) - radial_eval_alpha*x21*x4); + basis_y_eval[ipt + 4*npts] = 0.125*x28*x50 + 0.125*x40*(-4.0*x13*y + x3*y + x49); + basis_y_eval[ipt + 5*npts] = x39; + basis_y_eval[ipt + 6*npts] = x29*(x30*x50 - x43*(x13*x51 - x49)); + basis_y_eval[ipt + 7*npts] = x36*(radial_eval_alpha*x32 + x38); + basis_y_eval[ipt + 8*npts] = x46*(x33*x50 - x43*(x3*x51 - x49)); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x1*x42*x6*z; + basis_z_eval[ipt + 1*npts] = sqrt_70*x52*(radial_eval_alpha*x11*x13 + x34); + basis_z_eval[ipt + 2*npts] = x*x12*z*(radial_eval_alpha*x16 + x40); + basis_z_eval[ipt + 3*npts] = sqrt_10*x52*x53; + basis_z_eval[ipt + 4*npts] = -2.0*radial_eval*(x3*x54 + x4*x54 - 2.0*z*z*z) + 0.125*x28*x55; + basis_z_eval[ipt + 5*npts] = sqrt_10*x53*x56; + basis_z_eval[ipt + 6*npts] = x29*z*(radial_eval_alpha*x30 + x40*x6); + basis_z_eval[ipt + 7*npts] = sqrt_70*x56*(radial_eval_alpha*x13*x32 + x48); + basis_z_eval[ipt + 8*npts] = x33*x46*x55; + + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = x1*x2*x6; + ang_eval_1 = x11*x8*x9; + ang_eval_2 = x12*x16*x2; + ang_eval_3 = x17*x21*x9; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x22*x28; + ang_eval_1 = x17*x2*x21; + ang_eval_2 = radial_eval*x29*x30; + ang_eval_3 = x2*x32*x8; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = sqrt_35*x22*x33; + basis_eval[ipt + 8*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x1*(radial_eval_alpha*x3*x6 + x34); + dang_eval_y_0 = sqrt_35*x47*(radial_eval_alpha*x4*x6 + x48); + dang_eval_z_0 = x1*x42*x6*z; + dang_eval_x_1 = x36*(radial_eval_alpha*x11 + x37); + dang_eval_y_1 = x8*(radial_eval_alpha*x11*x4 + x45); + dang_eval_z_1 = sqrt_70*x52*(radial_eval_alpha*x11*x13 + x34); + dang_eval_x_2 = -x12*(radial_eval*(x10 + x15) - radial_eval_alpha*x16*x3); + dang_eval_y_2 = -sqrt_5*x47*(radial_eval*(x14 + x19 + x3) - radial_eval_alpha*x16*x4); + dang_eval_z_2 = x*x12*z*(radial_eval_alpha*x16 + x40); + dang_eval_x_3 = x39; + dang_eval_y_3 = -x17*(radial_eval*(x10 + x18 + 9.0*x4) - radial_eval_alpha*x21*x4); + dang_eval_z_3 = sqrt_10*x52*x53; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = 0.125*x28*x42 + 0.125*x40*(-4.0*x*x13 + x*x4 + x41); + dang_eval_y_0 = 0.125*x28*x50 + 0.125*x40*(-4.0*x13*y + x3*y + x49); + dang_eval_z_0 = -2.0*radial_eval*(x3*x54 + x4*x54 - 2.0*z*z*z) + 0.125*x28*x55; + dang_eval_x_1 = -x17*(radial_eval*(x20 + 9.0*x3) - radial_eval_alpha*x21*x3); + dang_eval_y_1 = x39; + dang_eval_z_1 = sqrt_10*x53*x56; + dang_eval_x_2 = x29*(x30*x42 + x43*(x13*x44 - x41)); + dang_eval_y_2 = x29*(x30*x50 - x43*(x13*x51 - x49)); + dang_eval_z_2 = x29*z*(radial_eval_alpha*x30 + x40*x6); + dang_eval_x_3 = x8*(radial_eval_alpha*x3*x32 + x45); + dang_eval_y_3 = x36*(radial_eval_alpha*x32 + x38); + dang_eval_z_3 = sqrt_70*x56*(radial_eval_alpha*x13*x32 + x48); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x46*(x33*x42 - x43*(x4*x44 - x41)); + dang_eval_y_0 = x46*(x33*x50 - x43*(x3*x51 - x49)); + dang_eval_z_0 = x33*x46*x55; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_hessian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_hessian.hpp new file mode 100644 index 000000000..3afede426 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_hessian.hpp @@ -0,0 +1,519 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_hessian_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_hessian_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = 0.5*sqrt_35; + const auto x1 = x0*y; + const auto x2 = x*x1; + const auto x3 = x*x; + const auto x4 = x3; + const auto x5 = y*y; + const auto x6 = x5; + const auto x7 = -x6; + const auto x8 = x4 + x7; + const auto x9 = 0.25*sqrt_70; + const auto x10 = x9*z; + const auto x11 = x10*y; + const auto x12 = 3.0*x4; + const auto x13 = x12 + x7; + const auto x14 = 0.5*sqrt_5; + const auto x15 = x14*y; + const auto x16 = x*x15; + const auto x17 = z*z; + const auto x18 = x17; + const auto x19 = -6.0*x18; + const auto x20 = x19 + x6; + const auto x21 = -x20 - x4; + const auto x22 = 0.25*sqrt_10; + const auto x23 = x22*z; + const auto x24 = x23*y; + const auto x25 = -4.0*x18; + const auto x26 = 3.0*x6; + const auto x27 = x25 + x26; + const auto x28 = -x12 - x27; + const auto x29 = 0.125*radial_eval; + const auto x30 = x*x*x*x; + const auto x31 = y*y*y*y; + const auto x32 = 6.0*x4*x6; + const auto x33 = x18*x4; + const auto x34 = x18*x6; + const auto x35 = 3.0*x30 + 3.0*x31 + x32 - 24.0*x33 - 24.0*x34 + 8.0*(z*z*z*z); + const auto x36 = x*x23; + const auto x37 = 0.25*sqrt_5; + const auto x38 = -x30 + x31 + 6.0*x33 - 6.0*x34; + const auto x39 = x*x10; + const auto x40 = -x26; + const auto x41 = x4 + x40; + const auto x42 = x30 + x31 - x32; + const auto x43 = radial_eval*x13; + const auto x44 = x4*x8; + const auto x45 = x*x11; + const auto x46 = 6.0*radial_eval; + const auto x47 = radial_eval_alpha*x13; + const auto x48 = x46 + x47; + const auto x49 = -x12 - x20; + const auto x50 = x21*x4; + const auto x51 = -x46; + const auto x52 = x*x24*(radial_eval_alpha*x28 + x51); + const auto x53 = 12.0*radial_eval; + const auto x54 = x*x*x; + const auto x55 = 4.0*x; + const auto x56 = x*x6 - x18*x55 + x54; + const auto x57 = radial_eval_alpha*x; + const auto x58 = 9.0*x4; + const auto x59 = -x27 - x58; + const auto x60 = x28*x4; + const auto x61 = 4.0*radial_eval; + const auto x62 = 3.0*x; + const auto x63 = x18*x62 - x54; + const auto x64 = x12 + x40; + const auto x65 = radial_eval*x64; + const auto x66 = x4*x41; + const auto x67 = radial_eval_alpha*x66 + x65; + const auto x68 = 0.125*sqrt_35; + const auto x69 = x54 - x6*x62; + const auto x70 = x*x0; + const auto x71 = radial_eval*x41; + const auto x72 = x6*x8; + const auto x73 = x13*x6; + const auto x74 = radial_eval_alpha*x73 + x65; + const auto x75 = x*x14; + const auto x76 = x19 + x26; + const auto x77 = -x4 - x76; + const auto x78 = x21*x6; + const auto x79 = 9.0*x6; + const auto x80 = x12 + x25; + const auto x81 = -x79 - x80; + const auto x82 = x28*x6; + const auto x83 = y*y*y; + const auto x84 = 4.0*y; + const auto x85 = -x18*x84 + x4*y + x83; + const auto x86 = radial_eval_alpha*y; + const auto x87 = 3.0*y; + const auto x88 = -x18*x87 + x83; + const auto x89 = radial_eval_alpha*x41; + const auto x90 = x51 + x89; + const auto x91 = -x4*x87 + x83; + const auto x92 = x1*z; + const auto x93 = x9*y; + const auto x94 = x13*x18; + const auto x95 = x22*y; + const auto x96 = -12.0*x18; + const auto x97 = x26 + x96; + const auto x98 = -x12 - x97; + const auto x99 = x18*x28; + const auto x100 = radial_eval*x98 + radial_eval_alpha*x99; + const auto x101 = 3.0*z; + const auto x102 = -x101*x4 - x101*x6 + 2.0*(z*z*z); + const auto x103 = radial_eval_alpha*z; + const auto x104 = x37*z; + const auto x105 = x53*x8; + const auto x106 = x18*x41; + const auto x107 = 2.0*radial_eval_alpha; + const auto x108 = x107*x13; + const auto x109 = radial_eval_alpha + radial_eval_alpha_squared*x4; + const auto x110 = x108 + x109*x8; + const auto x111 = x109*x13; + const auto x112 = 12.0*radial_eval_alpha; + const auto x113 = x112*x4; + const auto x114 = x113 + x46; + const auto x115 = x107*x49 + x109*x21; + const auto x116 = x109*x35 + x53*(x6 + x80) + 24.0*x56*x57; + const auto x117 = -18.0*radial_eval; + const auto x118 = x109*x28; + const auto x119 = x107*x59 + x118; + const auto x120 = -x4; + const auto x121 = 8.0*x57; + const auto x122 = x109*x38 + x121*x63 + x53*(x120 + x18); + const auto x123 = x107*x64; + const auto x124 = x109*x41 + x123; + const auto x125 = x105 + x109*x42 + x121*x69; + const auto x126 = radial_eval_alpha*x3; + const auto x127 = radial_eval_alpha*x5; + const auto x128 = 6.0*radial_eval_alpha; + const auto x129 = x128*x6; + const auto x130 = radial_eval_alpha*x64; + const auto x131 = 24.0*radial_eval; + const auto x132 = x*x131; + const auto x133 = x132*y; + const auto x134 = 12.0*x57; + const auto x135 = 12.0*x86; + const auto x136 = radial_eval_alpha_squared*x; + const auto x137 = x136*y; + const auto x138 = -x128*x4 + x51; + const auto x139 = radial_eval_alpha*x55; + const auto x140 = radial_eval_alpha*x84; + const auto x141 = x*x93; + const auto x142 = x128*x18; + const auto x143 = -x142; + const auto x144 = x*x95*(radial_eval_alpha*x98 + radial_eval_alpha_squared*x99 + x143 + x51); + const auto x145 = 96.0*radial_eval*z; + const auto x146 = 12.0*x103; + const auto x147 = radial_eval_alpha*x17; + const auto x148 = 4.0*radial_eval_alpha; + const auto x149 = x147*x64; + const auto x150 = x68*z; + const auto x151 = x107*x41; + const auto x152 = radial_eval_alpha + radial_eval_alpha_squared*x6; + const auto x153 = x151 + x152*x8; + const auto x154 = x123 + x13*x152; + const auto x155 = x107*x77 + x152*x21; + const auto x156 = x152*x28; + const auto x157 = x107*x81 + x156; + const auto x158 = x152*x35 + x53*(x27 + x4) + 24.0*x85*x86; + const auto x159 = x112*x6; + const auto x160 = x159 + x46; + const auto x161 = 8.0*x86; + const auto x162 = x152*x38 + x161*x88 - x53*(x18 - x6); + const auto x163 = x152*x42 + x161*x91 + x53*(x120 + x6); + const auto x164 = radial_eval_alpha_squared*y; + const auto x165 = radial_eval_alpha + radial_eval_alpha_squared*x18; + const auto x166 = x165*x8; + const auto x167 = x108 + x13*x165; + const auto x168 = 24.0*radial_eval_alpha*x18 + x165*x21; + const auto x169 = x107*x98 + x165*x28; + const auto x170 = x131 + x169; + const auto x171 = -48.0*radial_eval*(-2.0*x18 + x4 + x6) + 32.0*x102*x103 + x165*x35; + const auto x172 = x105 + 24.0*x147*x8 + x165*x38; + const auto x173 = x151 + x165*x41; + const auto x174 = x165*x42; + const auto x175 = -x159; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x2*x8; + basis_eval[ipt + 1*npts] = radial_eval*x11*x13; + basis_eval[ipt + 2*npts] = radial_eval*x16*x21; + basis_eval[ipt + 3*npts] = radial_eval*x24*x28; + basis_eval[ipt + 4*npts] = x29*x35; + basis_eval[ipt + 5*npts] = radial_eval*x28*x36; + basis_eval[ipt + 6*npts] = radial_eval*x37*x38; + basis_eval[ipt + 7*npts] = radial_eval*x39*x41; + basis_eval[ipt + 8*npts] = sqrt_35*x29*x42; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1*(radial_eval_alpha*x44 + x43); + basis_x_eval[ipt + 1*npts] = x45*x48; + basis_x_eval[ipt + 2*npts] = x15*(radial_eval*x49 + radial_eval_alpha*x50); + basis_x_eval[ipt + 3*npts] = x52; + basis_x_eval[ipt + 4*npts] = 0.125*x35*x57 + 0.125*x53*x56; + basis_x_eval[ipt + 5*npts] = x23*(radial_eval*x59 + radial_eval_alpha*x60); + basis_x_eval[ipt + 6*npts] = x37*(x38*x57 + x61*x63); + basis_x_eval[ipt + 7*npts] = x10*x67; + basis_x_eval[ipt + 8*npts] = x68*(x42*x57 + x61*x69); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x70*(radial_eval_alpha*x72 + x71); + basis_y_eval[ipt + 1*npts] = x10*x74; + basis_y_eval[ipt + 2*npts] = x75*(radial_eval*x77 + radial_eval_alpha*x78); + basis_y_eval[ipt + 3*npts] = x23*(radial_eval*x81 + radial_eval_alpha*x82); + basis_y_eval[ipt + 4*npts] = 0.125*x35*x86 + 0.125*x53*x85; + basis_y_eval[ipt + 5*npts] = x52; + basis_y_eval[ipt + 6*npts] = x37*(x38*x86 + x61*x88); + basis_y_eval[ipt + 7*npts] = x45*x90; + basis_y_eval[ipt + 8*npts] = x68*(x42*x86 + x61*x91); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x57*x8*x92; + basis_z_eval[ipt + 1*npts] = x93*(radial_eval_alpha*x94 + x43); + basis_z_eval[ipt + 2*npts] = x16*z*(radial_eval_alpha*x21 + x53); + basis_z_eval[ipt + 3*npts] = x100*x95; + basis_z_eval[ipt + 4*npts] = 2.0*radial_eval*x102 + 0.125*x103*x35; + basis_z_eval[ipt + 5*npts] = x*x100*x22; + basis_z_eval[ipt + 6*npts] = x104*(radial_eval_alpha*x38 + x105); + basis_z_eval[ipt + 7*npts] = x*x9*(radial_eval_alpha*x106 + x71); + basis_z_eval[ipt + 8*npts] = x103*x42*x68; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x2*(x110 + x46); + basis_xx_eval[ipt + 1*npts] = x11*(x111 + x114); + basis_xx_eval[ipt + 2*npts] = x16*(x115 + x51); + basis_xx_eval[ipt + 3*npts] = x24*(x109*x28 - x114); + basis_xx_eval[ipt + 4*npts] = 0.125*x116; + basis_xx_eval[ipt + 5*npts] = x36*(x117 + x119); + basis_xx_eval[ipt + 6*npts] = x122*x37; + basis_xx_eval[ipt + 7*npts] = x39*(x124 + x46); + basis_xx_eval[ipt + 8*npts] = x125*x68; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x0*(radial_eval_alpha_squared*x4*x6*x8 + x126*x41 + x127*x13 + x65); + basis_xy_eval[ipt + 1*npts] = x39*(radial_eval_alpha_squared*x73 + x129 + x130 + x46); + basis_xy_eval[ipt + 2*npts] = x14*(-radial_eval*(x12 + x76) + radial_eval_alpha_squared*x21*x4*x6 + x126*x77 + x127*x49); + basis_xy_eval[ipt + 3*npts] = x36*(radial_eval_alpha*x81 + radial_eval_alpha_squared*x82 - x129 + x51); + basis_xy_eval[ipt + 4*npts] = 0.125*x133 + 0.125*x134*x85 + 0.125*x135*x56 + 0.125*x137*x35; + basis_xy_eval[ipt + 5*npts] = x24*(radial_eval_alpha*x59 + radial_eval_alpha_squared*x60 + x138); + basis_xy_eval[ipt + 6*npts] = x37*(x137*x38 + x139*x88 + x140*x63); + basis_xy_eval[ipt + 7*npts] = x11*(radial_eval_alpha_squared*x66 + x130 + x138); + basis_xy_eval[ipt + 8*npts] = x68*(-x133 + x137*x42 + x139*x91 + x140*x69); + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x92*(radial_eval_alpha_squared*x44 + x47); + basis_xz_eval[ipt + 1*npts] = x141*(radial_eval_alpha_squared*x94 + x142 + x48); + basis_xz_eval[ipt + 2*npts] = x15*z*(radial_eval_alpha*x49 + radial_eval_alpha_squared*x50 + x113 + x53); + basis_xz_eval[ipt + 3*npts] = x144; + basis_xz_eval[ipt + 4*npts] = -0.125*x*x145 + 2.0*x102*x57 + 0.125*x136*x35*z + 0.125*x146*x56; + basis_xz_eval[ipt + 5*npts] = x22*(-radial_eval*(x58 + x97) + radial_eval_alpha_squared*x18*x28*x4 + x126*x98 + x147*x59); + basis_xz_eval[ipt + 6*npts] = x104*(x132 + x134*x8 + x136*x38 + x148*x63); + basis_xz_eval[ipt + 7*npts] = x9*(radial_eval_alpha_squared*x18*x4*x41 + x149 + x67); + basis_xz_eval[ipt + 8*npts] = x150*(x136*x42 + x148*x69); + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x2*(x153 + x51); + basis_yy_eval[ipt + 1*npts] = x11*(x154 + x51); + basis_yy_eval[ipt + 2*npts] = x16*(x155 + x51); + basis_yy_eval[ipt + 3*npts] = x24*(x117 + x157); + basis_yy_eval[ipt + 4*npts] = 0.125*x158; + basis_yy_eval[ipt + 5*npts] = x36*(x152*x28 - x160); + basis_yy_eval[ipt + 6*npts] = x162*x37; + basis_yy_eval[ipt + 7*npts] = x39*(x152*x41 - x160); + basis_yy_eval[ipt + 8*npts] = x163*x68; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x70*z*(radial_eval_alpha_squared*x72 + x89); + basis_yz_eval[ipt + 1*npts] = x9*(radial_eval_alpha_squared*x13*x18*x6 + x149 + x74); + basis_yz_eval[ipt + 2*npts] = x75*z*(radial_eval_alpha*x77 + radial_eval_alpha_squared*x78 + x159 + x53); + basis_yz_eval[ipt + 3*npts] = x22*(-radial_eval*(x12 + x79 + x96) + radial_eval_alpha_squared*x18*x28*x6 + x127*x98 + x147*x81); + basis_yz_eval[ipt + 4*npts] = 2.0*x102*x86 - 0.125*x145*y + 0.125*x146*x85 + 0.125*x164*x35*z; + basis_yz_eval[ipt + 5*npts] = x144; + basis_yz_eval[ipt + 6*npts] = x104*(-x131*y + x135*x8 + x148*x88 + x164*x38); + basis_yz_eval[ipt + 7*npts] = x141*(radial_eval_alpha_squared*x106 + x143 + x90); + basis_yz_eval[ipt + 8*npts] = x150*(x148*x91 + x164*x42); + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x166*x2; + basis_zz_eval[ipt + 1*npts] = x11*x167; + basis_zz_eval[ipt + 2*npts] = x16*(x168 + x53); + basis_zz_eval[ipt + 3*npts] = x170*x24; + basis_zz_eval[ipt + 4*npts] = 0.125*x171; + basis_zz_eval[ipt + 5*npts] = x170*x36; + basis_zz_eval[ipt + 6*npts] = x172*x37; + basis_zz_eval[ipt + 7*npts] = x173*x39; + basis_zz_eval[ipt + 8*npts] = x174*x68; + + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x2*x8; + ang_eval_1 = radial_eval*x11*x13; + ang_eval_2 = radial_eval*x16*x21; + ang_eval_3 = radial_eval*x24*x28; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x29*x35; + ang_eval_1 = radial_eval*x28*x36; + ang_eval_2 = radial_eval*x37*x38; + ang_eval_3 = radial_eval*x39*x41; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = sqrt_35*x29*x42; + basis_eval[ipt + 8*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x1*(radial_eval_alpha*x44 + x43); + dang_eval_y_0 = x70*(radial_eval_alpha*x72 + x71); + dang_eval_z_0 = x57*x8*x92; + dang_eval_x_1 = x45*x48; + dang_eval_y_1 = x10*x74; + dang_eval_z_1 = x93*(radial_eval_alpha*x94 + x43); + dang_eval_x_2 = x15*(radial_eval*x49 + radial_eval_alpha*x50); + dang_eval_y_2 = x75*(radial_eval*x77 + radial_eval_alpha*x78); + dang_eval_z_2 = x16*z*(radial_eval_alpha*x21 + x53); + dang_eval_x_3 = x52; + dang_eval_y_3 = x23*(radial_eval*x81 + radial_eval_alpha*x82); + dang_eval_z_3 = x100*x95; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = 0.125*x35*x57 + 0.125*x53*x56; + dang_eval_y_0 = 0.125*x35*x86 + 0.125*x53*x85; + dang_eval_z_0 = 2.0*radial_eval*x102 + 0.125*x103*x35; + dang_eval_x_1 = x23*(radial_eval*x59 + radial_eval_alpha*x60); + dang_eval_y_1 = x52; + dang_eval_z_1 = x*x100*x22; + dang_eval_x_2 = x37*(x38*x57 + x61*x63); + dang_eval_y_2 = x37*(x38*x86 + x61*x88); + dang_eval_z_2 = x104*(radial_eval_alpha*x38 + x105); + dang_eval_x_3 = x10*x67; + dang_eval_y_3 = x45*x90; + dang_eval_z_3 = x*x9*(radial_eval_alpha*x106 + x71); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x68*(x42*x57 + x61*x69); + dang_eval_y_0 = x68*(x42*x86 + x61*x91); + dang_eval_z_0 = x103*x42*x68; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_lapgrad.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_lapgrad.hpp new file mode 100644 index 000000000..e11fb857f --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_lapgrad.hpp @@ -0,0 +1,670 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_lapgrad_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_lapgrad_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + auto* __restrict__ basis_lapl_x_eval = task->d3bflapl_x + shoff; + auto* __restrict__ basis_lapl_y_eval = task->d3bflapl_y + shoff; + auto* __restrict__ basis_lapl_z_eval = task->d3bflapl_z + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + double radial_eval_alpha_cubed = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + radial_eval_alpha_cubed += a * a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + radial_eval_alpha_cubed *= -8; + + // Common Subexpressions + const auto x0 = 0.5*sqrt_35; + const auto x1 = x0*y; + const auto x2 = x*x1; + const auto x3 = x*x; + const auto x4 = x3; + const auto x5 = y*y; + const auto x6 = x5; + const auto x7 = -x6; + const auto x8 = x4 + x7; + const auto x9 = 0.25*sqrt_70; + const auto x10 = x9*z; + const auto x11 = x10*y; + const auto x12 = 3.0*x4; + const auto x13 = x12 + x7; + const auto x14 = 0.5*sqrt_5; + const auto x15 = x14*y; + const auto x16 = x*x15; + const auto x17 = z*z; + const auto x18 = x17; + const auto x19 = -6.0*x18; + const auto x20 = x19 + x6; + const auto x21 = -x20 - x4; + const auto x22 = 0.25*sqrt_10; + const auto x23 = x22*z; + const auto x24 = x23*y; + const auto x25 = -4.0*x18; + const auto x26 = 3.0*x6; + const auto x27 = x25 + x26; + const auto x28 = -x12 - x27; + const auto x29 = 0.125*radial_eval; + const auto x30 = x*x*x*x; + const auto x31 = y*y*y*y; + const auto x32 = x4*x6; + const auto x33 = 6.0*x32; + const auto x34 = x18*x4; + const auto x35 = x18*x6; + const auto x36 = 3.0*x30 + 3.0*x31 + x33 - 24.0*x34 - 24.0*x35 + 8.0*(z*z*z*z); + const auto x37 = x*x23; + const auto x38 = 0.25*sqrt_5; + const auto x39 = -x30 + x31 + 6.0*x34 - 6.0*x35; + const auto x40 = x*x10; + const auto x41 = -x26; + const auto x42 = x4 + x41; + const auto x43 = x30 + x31 - x33; + const auto x44 = radial_eval*x13; + const auto x45 = x4*x8; + const auto x46 = x*x11; + const auto x47 = 6.0*radial_eval; + const auto x48 = radial_eval_alpha*x13; + const auto x49 = x47 + x48; + const auto x50 = -x12 - x20; + const auto x51 = x21*x4; + const auto x52 = -x47; + const auto x53 = x*x24*(radial_eval_alpha*x28 + x52); + const auto x54 = 12.0*radial_eval; + const auto x55 = x*x*x; + const auto x56 = 4.0*x; + const auto x57 = x*x6 - x18*x56 + x55; + const auto x58 = radial_eval_alpha*x; + const auto x59 = 9.0*x4; + const auto x60 = -x27 - x59; + const auto x61 = x28*x4; + const auto x62 = 4.0*radial_eval; + const auto x63 = 3.0*x; + const auto x64 = x18*x63 - x55; + const auto x65 = x12 + x41; + const auto x66 = radial_eval*x65; + const auto x67 = x4*x42; + const auto x68 = radial_eval_alpha*x67 + x66; + const auto x69 = 0.125*sqrt_35; + const auto x70 = x55 - x6*x63; + const auto x71 = x*x0; + const auto x72 = radial_eval*x42; + const auto x73 = x6*x8; + const auto x74 = x13*x6; + const auto x75 = radial_eval_alpha*x74 + x66; + const auto x76 = x*x14; + const auto x77 = x19 + x26; + const auto x78 = -x4 - x77; + const auto x79 = x21*x6; + const auto x80 = 9.0*x6; + const auto x81 = x12 + x25; + const auto x82 = -x80 - x81; + const auto x83 = x28*x6; + const auto x84 = y*y*y; + const auto x85 = 4.0*y; + const auto x86 = -x18*x85 + x4*y + x84; + const auto x87 = radial_eval_alpha*y; + const auto x88 = 3.0*y; + const auto x89 = -x18*x88 + x84; + const auto x90 = radial_eval_alpha*x42; + const auto x91 = x52 + x90; + const auto x92 = -x4*x88 + x84; + const auto x93 = x1*z; + const auto x94 = x9*y; + const auto x95 = x13*x18; + const auto x96 = x22*y; + const auto x97 = -12.0*x18; + const auto x98 = x26 + x97; + const auto x99 = -x12 - x98; + const auto x100 = x18*x28; + const auto x101 = radial_eval*x99 + radial_eval_alpha*x100; + const auto x102 = z*z*z; + const auto x103 = 3.0*z; + const auto x104 = 2.0*x102 - x103*x4 - x103*x6; + const auto x105 = radial_eval_alpha*z; + const auto x106 = x*x22; + const auto x107 = x38*z; + const auto x108 = x54*x8; + const auto x109 = x*x9; + const auto x110 = x18*x42; + const auto x111 = 2.0*radial_eval_alpha; + const auto x112 = x111*x13; + const auto x113 = radial_eval_alpha + radial_eval_alpha_squared*x4; + const auto x114 = x113*x8; + const auto x115 = x112 + x114; + const auto x116 = x113*x13; + const auto x117 = 12.0*radial_eval_alpha; + const auto x118 = x117*x4; + const auto x119 = x118 + x47; + const auto x120 = x111*x50 + x113*x21; + const auto x121 = x6 + x81; + const auto x122 = x113*x36 + x121*x54 + 24.0*x57*x58; + const auto x123 = -18.0*radial_eval; + const auto x124 = x113*x28; + const auto x125 = x111*x60 + x124; + const auto x126 = -x4; + const auto x127 = x126 + x18; + const auto x128 = 8.0*x58; + const auto x129 = x113*x39 + x127*x54 + x128*x64; + const auto x130 = x111*x65; + const auto x131 = x113*x42 + x130; + const auto x132 = x108 + x113*x43 + x128*x70; + const auto x133 = radial_eval_alpha*x3; + const auto x134 = radial_eval_alpha*x5; + const auto x135 = 6.0*radial_eval_alpha; + const auto x136 = x135*x6; + const auto x137 = radial_eval_alpha*x65; + const auto x138 = -x12 - x77; + const auto x139 = 24.0*radial_eval; + const auto x140 = x*x139; + const auto x141 = x140*y; + const auto x142 = 12.0*x58; + const auto x143 = 12.0*x87; + const auto x144 = radial_eval_alpha_squared*x; + const auto x145 = x144*y; + const auto x146 = -x135*x4 + x52; + const auto x147 = radial_eval_alpha*x56; + const auto x148 = radial_eval_alpha*x85; + const auto x149 = x*x94; + const auto x150 = x135*x18; + const auto x151 = -x150; + const auto x152 = x*x96*(radial_eval_alpha*x99 + radial_eval_alpha_squared*x100 + x151 + x52); + const auto x153 = 96.0*radial_eval*z; + const auto x154 = 12.0*x105; + const auto x155 = x144*z; + const auto x156 = -x59 - x98; + const auto x157 = radial_eval_alpha*x17; + const auto x158 = x142*x8; + const auto x159 = 4.0*radial_eval_alpha; + const auto x160 = x157*x65; + const auto x161 = x69*z; + const auto x162 = x111*x42; + const auto x163 = radial_eval_alpha + radial_eval_alpha_squared*x6; + const auto x164 = x163*x8; + const auto x165 = x162 + x164; + const auto x166 = x13*x163 + x130; + const auto x167 = x111*x78 + x163*x21; + const auto x168 = x163*x28; + const auto x169 = x111*x82 + x168; + const auto x170 = x27 + x4; + const auto x171 = x163*x36 + x170*x54 + 24.0*x86*x87; + const auto x172 = x117*x6; + const auto x173 = x172 + x47; + const auto x174 = -x18 + x6; + const auto x175 = 8.0*x87; + const auto x176 = x163*x39 + x174*x54 + x175*x89; + const auto x177 = x126 + x6; + const auto x178 = x163*x43 + x175*x92 + x177*x54; + const auto x179 = -x12 - x80 - x97; + const auto x180 = radial_eval_alpha_squared*y; + const auto x181 = x180*z; + const auto x182 = x143*x8; + const auto x183 = radial_eval_alpha + radial_eval_alpha_squared*x18; + const auto x184 = x183*x8; + const auto x185 = x13*x183; + const auto x186 = x112 + x185; + const auto x187 = 24.0*radial_eval_alpha*x18; + const auto x188 = x183*x21 + x187; + const auto x189 = x111*x99 + x183*x28; + const auto x190 = x139 + x189; + const auto x191 = 2.0*x18 - x4 - x6; + const auto x192 = 48.0*radial_eval*x191 + 32.0*x104*x105 + x183*x36; + const auto x193 = x108 + 24.0*x157*x8 + x183*x39; + const auto x194 = x183*x42; + const auto x195 = x162 + x194; + const auto x196 = x183*x43; + const auto x197 = x118 + x166; + const auto x198 = x116 + x197; + const auto x199 = -x118; + const auto x200 = -x172; + const auto x201 = x163*x42; + const auto x202 = x131 + x200; + const auto x203 = x201 + x202; + const auto x204 = radial_eval_alpha_cubed*x55 + radial_eval_alpha_squared*x63; + const auto x205 = radial_eval_alpha_cubed*x6 + radial_eval_alpha_squared; + const auto x206 = x205*x8; + const auto x207 = radial_eval_alpha_cubed*x18 + radial_eval_alpha_squared; + const auto x208 = x207*x8; + const auto x209 = 2.0*radial_eval_alpha_squared; + const auto x210 = x209*x3; + const auto x211 = 36.0*x58; + const auto x212 = 18.0*x*x113; + const auto x213 = 6.0*x; + const auto x214 = x163*x213; + const auto x215 = x183*x213; + const auto x216 = 2.0*x144; + const auto x217 = x13*x205; + const auto x218 = x13*x207; + const auto x219 = x205*x21; + const auto x220 = x207*x21; + const auto x221 = 24.0*radial_eval_alpha_squared; + const auto x222 = x111*x138 + x187; + const auto x223 = x205*x28; + const auto x224 = x207*x28; + const auto x225 = x204*x28; + const auto x226 = 48.0*x58; + const auto x227 = x226*x6; + const auto x228 = 24.0*x145; + const auto x229 = x205*x36; + const auto x230 = x207*x36; + const auto x231 = 36.0*radial_eval_alpha; + const auto x232 = x111*x156; + const auto x233 = 12.0*radial_eval_alpha_squared; + const auto x234 = x233*x32; + const auto x235 = -x234; + const auto x236 = x200 + x235; + const auto x237 = 8.0*x145; + const auto x238 = 24.0*x17; + const auto x239 = x205*x39; + const auto x240 = x207*x39; + const auto x241 = x163*x65; + const auto x242 = x113*x65; + const auto x243 = x205*x42; + const auto x244 = x207*x42; + const auto x245 = x118 + x130 + x183*x65; + const auto x246 = x205*x43; + const auto x247 = x207*x43; + const auto x248 = radial_eval_alpha_cubed*x84 + radial_eval_alpha_squared*x88; + const auto x249 = radial_eval_alpha_cubed*x4 + radial_eval_alpha_squared; + const auto x250 = x249*x8; + const auto x251 = x209*x5; + const auto x252 = x13*x249; + const auto x253 = x21*x249; + const auto x254 = x248*x28; + const auto x255 = x249*x28; + const auto x256 = x111*x179 + x199; + const auto x257 = 48.0*x87; + const auto x258 = x257*x4; + const auto x259 = 36.0*x87; + const auto x260 = x249*x36; + const auto x261 = 2.0*x180; + const auto x262 = 6.0*y; + const auto x263 = -x113*x262 - 18.0*x163*y - x183*x262 - x259; + const auto x264 = x249*x39; + const auto x265 = x249*x42; + const auto x266 = x249*x43; + const auto x267 = x209*z; + const auto x268 = radial_eval_alpha_cubed*x102 + radial_eval_alpha_squared*x103; + const auto x269 = x17*x209; + const auto x270 = x269*x65; + const auto x271 = x233*x34; + const auto x272 = 12.0*z; + const auto x273 = 36.0*z; + const auto x274 = 48.0*radial_eval_alpha*x18 + x113*x99 + x163*x99 + x17*x223 + x17*x255 + 3.0*x183*x99 + x268*x28*z; + const auto x275 = 192.0*x105; + const auto x276 = -x233*x35; + const auto x277 = 48.0*x105; + const auto x278 = 8.0*x181; + const auto x279 = 8.0*x155; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x2*x8; + basis_eval[ipt + 1*npts] = radial_eval*x11*x13; + basis_eval[ipt + 2*npts] = radial_eval*x16*x21; + basis_eval[ipt + 3*npts] = radial_eval*x24*x28; + basis_eval[ipt + 4*npts] = x29*x36; + basis_eval[ipt + 5*npts] = radial_eval*x28*x37; + basis_eval[ipt + 6*npts] = radial_eval*x38*x39; + basis_eval[ipt + 7*npts] = radial_eval*x40*x42; + basis_eval[ipt + 8*npts] = sqrt_35*x29*x43; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1*(radial_eval_alpha*x45 + x44); + basis_x_eval[ipt + 1*npts] = x46*x49; + basis_x_eval[ipt + 2*npts] = x15*(radial_eval*x50 + radial_eval_alpha*x51); + basis_x_eval[ipt + 3*npts] = x53; + basis_x_eval[ipt + 4*npts] = 0.125*x36*x58 + 0.125*x54*x57; + basis_x_eval[ipt + 5*npts] = x23*(radial_eval*x60 + radial_eval_alpha*x61); + basis_x_eval[ipt + 6*npts] = x38*(x39*x58 + x62*x64); + basis_x_eval[ipt + 7*npts] = x10*x68; + basis_x_eval[ipt + 8*npts] = x69*(x43*x58 + x62*x70); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x71*(radial_eval_alpha*x73 + x72); + basis_y_eval[ipt + 1*npts] = x10*x75; + basis_y_eval[ipt + 2*npts] = x76*(radial_eval*x78 + radial_eval_alpha*x79); + basis_y_eval[ipt + 3*npts] = x23*(radial_eval*x82 + radial_eval_alpha*x83); + basis_y_eval[ipt + 4*npts] = 0.125*x36*x87 + 0.125*x54*x86; + basis_y_eval[ipt + 5*npts] = x53; + basis_y_eval[ipt + 6*npts] = x38*(x39*x87 + x62*x89); + basis_y_eval[ipt + 7*npts] = x46*x91; + basis_y_eval[ipt + 8*npts] = x69*(x43*x87 + x62*x92); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x58*x8*x93; + basis_z_eval[ipt + 1*npts] = x94*(radial_eval_alpha*x95 + x44); + basis_z_eval[ipt + 2*npts] = x16*z*(radial_eval_alpha*x21 + x54); + basis_z_eval[ipt + 3*npts] = x101*x96; + basis_z_eval[ipt + 4*npts] = 2.0*radial_eval*x104 + 0.125*x105*x36; + basis_z_eval[ipt + 5*npts] = x101*x106; + basis_z_eval[ipt + 6*npts] = x107*(radial_eval_alpha*x39 + x108); + basis_z_eval[ipt + 7*npts] = x109*(radial_eval_alpha*x110 + x72); + basis_z_eval[ipt + 8*npts] = x105*x43*x69; + + // Evaluate second derivative of bfn wrt xx + basis_xx_eval[ipt + 0*npts] = x2*(x115 + x47); + basis_xx_eval[ipt + 1*npts] = x11*(x116 + x119); + basis_xx_eval[ipt + 2*npts] = x16*(x120 + x52); + basis_xx_eval[ipt + 3*npts] = x24*(x113*x28 - x119); + basis_xx_eval[ipt + 4*npts] = 0.125*x122; + basis_xx_eval[ipt + 5*npts] = x37*(x123 + x125); + basis_xx_eval[ipt + 6*npts] = x129*x38; + basis_xx_eval[ipt + 7*npts] = x40*(x131 + x47); + basis_xx_eval[ipt + 8*npts] = x132*x69; + + // Evaluate second derivative of bfn wrt xy + basis_xy_eval[ipt + 0*npts] = x0*(radial_eval_alpha_squared*x4*x6*x8 + x13*x134 + x133*x42 + x66); + basis_xy_eval[ipt + 1*npts] = x40*(radial_eval_alpha_squared*x74 + x136 + x137 + x47); + basis_xy_eval[ipt + 2*npts] = x14*(radial_eval*x138 + radial_eval_alpha_squared*x21*x4*x6 + x133*x78 + x134*x50); + basis_xy_eval[ipt + 3*npts] = x37*(radial_eval_alpha*x82 + radial_eval_alpha_squared*x83 - x136 + x52); + basis_xy_eval[ipt + 4*npts] = 0.125*x141 + 0.125*x142*x86 + 0.125*x143*x57 + 0.125*x145*x36; + basis_xy_eval[ipt + 5*npts] = x24*(radial_eval_alpha*x60 + radial_eval_alpha_squared*x61 + x146); + basis_xy_eval[ipt + 6*npts] = x38*(x145*x39 + x147*x89 + x148*x64); + basis_xy_eval[ipt + 7*npts] = x11*(radial_eval_alpha_squared*x67 + x137 + x146); + basis_xy_eval[ipt + 8*npts] = x69*(-x141 + x145*x43 + x147*x92 + x148*x70); + + // Evaluate second derivative of bfn wrt xz + basis_xz_eval[ipt + 0*npts] = x93*(radial_eval_alpha_squared*x45 + x48); + basis_xz_eval[ipt + 1*npts] = x149*(radial_eval_alpha_squared*x95 + x150 + x49); + basis_xz_eval[ipt + 2*npts] = x15*z*(radial_eval_alpha*x50 + radial_eval_alpha_squared*x51 + x118 + x54); + basis_xz_eval[ipt + 3*npts] = x152; + basis_xz_eval[ipt + 4*npts] = -0.125*x*x153 + 2.0*x104*x58 + 0.125*x154*x57 + 0.125*x155*x36; + basis_xz_eval[ipt + 5*npts] = x22*(radial_eval*x156 + radial_eval_alpha_squared*x18*x28*x4 + x133*x99 + x157*x60); + basis_xz_eval[ipt + 6*npts] = x107*(x140 + x144*x39 + x158 + x159*x64); + basis_xz_eval[ipt + 7*npts] = x9*(radial_eval_alpha_squared*x18*x4*x42 + x160 + x68); + basis_xz_eval[ipt + 8*npts] = x161*(x144*x43 + x159*x70); + + // Evaluate second derivative of bfn wrt yy + basis_yy_eval[ipt + 0*npts] = x2*(x165 + x52); + basis_yy_eval[ipt + 1*npts] = x11*(x166 + x52); + basis_yy_eval[ipt + 2*npts] = x16*(x167 + x52); + basis_yy_eval[ipt + 3*npts] = x24*(x123 + x169); + basis_yy_eval[ipt + 4*npts] = 0.125*x171; + basis_yy_eval[ipt + 5*npts] = x37*(x163*x28 - x173); + basis_yy_eval[ipt + 6*npts] = x176*x38; + basis_yy_eval[ipt + 7*npts] = x40*(x163*x42 - x173); + basis_yy_eval[ipt + 8*npts] = x178*x69; + + // Evaluate second derivative of bfn wrt yz + basis_yz_eval[ipt + 0*npts] = x71*z*(radial_eval_alpha_squared*x73 + x90); + basis_yz_eval[ipt + 1*npts] = x9*(radial_eval_alpha_squared*x13*x18*x6 + x160 + x75); + basis_yz_eval[ipt + 2*npts] = x76*z*(radial_eval_alpha*x78 + radial_eval_alpha_squared*x79 + x172 + x54); + basis_yz_eval[ipt + 3*npts] = x22*(radial_eval*x179 + radial_eval_alpha_squared*x18*x28*x6 + x134*x99 + x157*x82); + basis_yz_eval[ipt + 4*npts] = 2.0*x104*x87 - 0.125*x153*y + 0.125*x154*x86 + 0.125*x181*x36; + basis_yz_eval[ipt + 5*npts] = x152; + basis_yz_eval[ipt + 6*npts] = x107*(-x139*y + x159*x89 + x180*x39 + x182); + basis_yz_eval[ipt + 7*npts] = x149*(radial_eval_alpha_squared*x110 + x151 + x91); + basis_yz_eval[ipt + 8*npts] = x161*(x159*x92 + x180*x43); + + // Evaluate second derivative of bfn wrt zz + basis_zz_eval[ipt + 0*npts] = x184*x2; + basis_zz_eval[ipt + 1*npts] = x11*x186; + basis_zz_eval[ipt + 2*npts] = x16*(x188 + x54); + basis_zz_eval[ipt + 3*npts] = x190*x24; + basis_zz_eval[ipt + 4*npts] = 0.125*x192; + basis_zz_eval[ipt + 5*npts] = x190*x37; + basis_zz_eval[ipt + 6*npts] = x193*x38; + basis_zz_eval[ipt + 7*npts] = x195*x40; + basis_zz_eval[ipt + 8*npts] = x196*x69; + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x2*(x115 + x165 + x184); + basis_lapl_eval[ipt + 1*npts] = x11*(x186 + x198); + basis_lapl_eval[ipt + 2*npts] = x16*(x120 + x167 + x188); + basis_lapl_eval[ipt + 3*npts] = x24*(x124 + x169 + x189 + x199); + basis_lapl_eval[ipt + 4*npts] = 0.125*x122 + 0.125*x171 + 0.125*x192; + basis_lapl_eval[ipt + 5*npts] = x37*(x125 + x168 + x189 + x200); + basis_lapl_eval[ipt + 6*npts] = x38*(x129 + x176 + x193); + basis_lapl_eval[ipt + 7*npts] = x40*(x195 + x203); + basis_lapl_eval[ipt + 8*npts] = x69*(x132 + x178 + x196); + + // Evaluate Laplacian gradient of bfn (dx) + basis_lapl_x_eval[ipt + 0*npts] = x1*(x*x204*x8 + 3.0*x116 + x185 + x197 + x206*x3 + x208*x3 + x210*x42); + basis_lapl_x_eval[ipt + 1*npts] = x11*(x*x217 + x*x218 + x13*x204 + x13*x216 + x211 + x212 + x214 + x215 + x216*x65); + basis_lapl_x_eval[ipt + 2*npts] = x15*(x*x204*x21 + 3.0*x113*x50 + x163*x50 + x183*x50 + x199 + x210*x78 + x219*x3 + x220*x3 + x221*x34 + x222); + basis_lapl_x_eval[ipt + 3*npts] = x24*(x*x223 + x*x224 - x211 - x212 - x214 - x215 + x216*x82 + x216*x99 + x225); + basis_lapl_x_eval[ipt + 4*npts] = 0.125*x*x229 + 0.125*x*x230 + 4.0*x104*x155 + 4.5*x113*x57 + 0.125*x121*x211 + 0.125*x142*x170 + 1.5*x163*x57 - 24.0*x18*x58 + 1.5*x183*x57 + 0.125*x191*x226 + 0.125*x204*x36 + 0.125*x227 + 0.125*x228*x86; + basis_lapl_x_eval[ipt + 5*npts] = x23*(x*x225 + 3.0*x113*x60 + x163*x60 + x183*x60 + x210*x99 + x223*x3 + x224*x3 - x231*x4 + x232 + x236); + basis_lapl_x_eval[ipt + 6*npts] = x38*(x*x239 + x*x240 + 12.0*x113*x64 + x127*x211 + x142*x174 + x144*x238*x8 + x158 + 4.0*x163*x64 + x18*x226 + 4.0*x183*x64 + x204*x39 + x237*x89); + basis_lapl_x_eval[ipt + 7*npts] = x10*(x*x204*x42 + x209*x67 + x236 + x241 + 3.0*x242 + x243*x3 + x244*x3 + x245); + basis_lapl_x_eval[ipt + 8*npts] = x69*(x*x246 + x*x247 + 12.0*x113*x70 + x142*x177 + 4.0*x163*x70 + 4.0*x183*x70 + x204*x43 + x211*x8 - x227 + x237*x92); + // Evaluate Laplacian gradient of bfn (dy) + basis_lapl_y_eval[ipt + 0*npts] = x71*(x13*x251 + x194 + 3.0*x201 + x202 + x208*x5 + x248*x8*y + x250*x5); + basis_lapl_y_eval[ipt + 1*npts] = x10*(x13*x248*y + x200 + x209*x74 + x218*x5 + x234 + 3.0*x241 + x242 + x245 + x252*x5); + basis_lapl_y_eval[ipt + 2*npts] = x76*(x113*x78 + 3.0*x163*x78 + x183*x78 + x200 + x21*x248*y + x220*x5 + x221*x35 + x222 + x251*x50 + x253*x5); + basis_lapl_y_eval[ipt + 3*npts] = x23*(x113*x82 + 3.0*x163*x82 + x183*x82 + x224*x5 - x231*x6 + x235 + x251*x99 + x254*y + x255*x5 + x256); + basis_lapl_y_eval[ipt + 4*npts] = 4.0*x104*x181 + 1.5*x113*x86 + 0.125*x121*x143 + 4.5*x163*x86 + 0.125*x170*x259 - 24.0*x18*x87 + 1.5*x183*x86 + 0.125*x191*x257 + 0.125*x228*x57 + 0.125*x230*y + 0.125*x248*x36 + 0.125*x258 + 0.125*x260*y; + basis_lapl_y_eval[ipt + 5*npts] = x37*(x224*y + x254 + x255*y + x261*x60 + x261*x99 + x263); + basis_lapl_y_eval[ipt + 6*npts] = x38*(4.0*x113*x89 + x127*x143 + 12.0*x163*x89 + x174*x259 - x18*x257 + x180*x238*x8 + x182 + 4.0*x183*x89 + x237*x64 + x240*y + x248*x39 + x264*y); + basis_lapl_y_eval[ipt + 7*npts] = x40*(x244*y + x248*x42 + x261*x42 + x261*x65 + x263 + x265*y); + basis_lapl_y_eval[ipt + 8*npts] = x69*(4.0*x113*x92 + 12.0*x163*x92 + x177*x259 + x182 + 4.0*x183*x92 + x237*x70 + x247*y + x248*x43 - x258 + x266*y); + // Evaluate Laplacian gradient of bfn (dz) + basis_lapl_z_eval[ipt + 0*npts] = x2*(x13*x267 + x206*z + x250*z + x267*x42 + x268*x8); + basis_lapl_z_eval[ipt + 1*npts] = x94*(x13*x268*z + x17*x217 + x17*x252 + 3.0*x185 + x198 + x270 + x271); + basis_lapl_z_eval[ipt + 2*npts] = x16*(72.0*x105 + x113*x272 + x163*x272 + x183*x273 + x21*x268 + x219*z + x253*z + x267*x50 + x267*x78); + basis_lapl_z_eval[ipt + 3*npts] = x96*(x256 + x269*x82 - x271 + x274); + basis_lapl_z_eval[ipt + 4*npts] = 2.0*x104*x113 + 2.0*x104*x163 + 6.0*x104*x183 + 18.0*x105*x191 + 0.125*x121*x154 + 0.125*x154*x170 + 3.0*x155*x57 + 3.0*x181*x86 + 0.125*x229*z + 0.125*x260*z + 0.125*x268*x36 - 0.125*x275*x4 - 0.125*x275*x6; + basis_lapl_z_eval[ipt + 5*npts] = x106*(x200 + x232 + x269*x60 + x274 + x276); + basis_lapl_z_eval[ipt + 6*npts] = x38*(36.0*x105*x8 + x114*x272 + x127*x154 + x154*x174 + x164*x272 + x184*x273 + x239*z + x264*z + x268*x39 + x277*x4 - x277*x6 + x278*x89 + x279*x64); + basis_lapl_z_eval[ipt + 7*npts] = x109*(x17*x243 + x17*x265 + 3.0*x194 + x203 + x268*x42*z + x270 + x276); + basis_lapl_z_eval[ipt + 8*npts] = x69*(x154*x177 + x154*x8 + x246*z + x266*z + x268*x43 + x278*x92 + x279*x70); + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x2*x8; + ang_eval_1 = radial_eval*x11*x13; + ang_eval_2 = radial_eval*x16*x21; + ang_eval_3 = radial_eval*x24*x28; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x29*x36; + ang_eval_1 = radial_eval*x28*x37; + ang_eval_2 = radial_eval*x38*x39; + ang_eval_3 = radial_eval*x40*x42; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = sqrt_35*x29*x43; + basis_eval[ipt + 8*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x1*(radial_eval_alpha*x45 + x44); + dang_eval_y_0 = x71*(radial_eval_alpha*x73 + x72); + dang_eval_z_0 = x58*x8*x93; + dang_eval_x_1 = x46*x49; + dang_eval_y_1 = x10*x75; + dang_eval_z_1 = x94*(radial_eval_alpha*x95 + x44); + dang_eval_x_2 = x15*(radial_eval*x50 + radial_eval_alpha*x51); + dang_eval_y_2 = x76*(radial_eval*x78 + radial_eval_alpha*x79); + dang_eval_z_2 = x16*z*(radial_eval_alpha*x21 + x54); + dang_eval_x_3 = x53; + dang_eval_y_3 = x23*(radial_eval*x82 + radial_eval_alpha*x83); + dang_eval_z_3 = x101*x96; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = 0.125*x36*x58 + 0.125*x54*x57; + dang_eval_y_0 = 0.125*x36*x87 + 0.125*x54*x86; + dang_eval_z_0 = 2.0*radial_eval*x104 + 0.125*x105*x36; + dang_eval_x_1 = x23*(radial_eval*x60 + radial_eval_alpha*x61); + dang_eval_y_1 = x53; + dang_eval_z_1 = x101*x106; + dang_eval_x_2 = x38*(x39*x58 + x62*x64); + dang_eval_y_2 = x38*(x39*x87 + x62*x89); + dang_eval_z_2 = x107*(radial_eval_alpha*x39 + x108); + dang_eval_x_3 = x10*x68; + dang_eval_y_3 = x46*x91; + dang_eval_z_3 = x109*(radial_eval_alpha*x110 + x72); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x69*(x43*x58 + x62*x70); + dang_eval_y_0 = x69*(x43*x87 + x62*x92); + dang_eval_z_0 = x105*x43*x69; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_laplacian.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_laplacian.hpp new file mode 100644 index 000000000..a99003d83 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation/collocation_shell_to_task_kernels_spherical_l4_laplacian.hpp @@ -0,0 +1,459 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "collocation_device_constants.hpp" +#include "device/xc_device_task.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "device/common/shell_to_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" +#include +#include +#include + +namespace GauXC { + +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_4_max_wg = 128; +static constexpr uint32_t collocation_device_shell_to_task_kernel_spherical_laplacian_4_nwarp = 4; + + +void collocation_device_shell_to_task_kernel_spherical_laplacian_4( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + auto& alpha = GauXC::sycl::local_mem(); + auto& coeff = GauXC::sycl::local_mem(); + + const uint32_t local_warp_id = it.get_local_id(2) / sycl::warp_size; + double* my_alpha = &alpha[local_warp_id][0]; + double* my_coeff = &coeff[local_warp_id][0]; + + for( auto ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + + + // Load Shell Data into registers / SM + const uint32_t nprim = shell->nprim(); + const double3 O = *reinterpret_cast(shell->O_data()); + + const int global_warp_id = (it.get_global_id(2)) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + // Read in coeffs/exps into SM on first warp + { + auto* coeff_gm = shell->coeff_data(); + auto* alpha_gm = shell->alpha_data(); + static_assert( detail::shell_nprim_max == sycl::warp_size ); + const int warp_rank = it.get_local_id(2) % sycl::warp_size; + my_alpha[warp_rank] = alpha_gm[warp_rank]; + my_coeff[warp_rank] = coeff_gm[warp_rank]; + } + + // Loop over tasks assigned to shells + // Place each task on a different warp + schedule across blocks + for( int itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const auto* __restrict__ points_x = task->points_x; + const auto* __restrict__ points_y = task->points_y; + const auto* __restrict__ points_z = task->points_z; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + + auto* __restrict__ basis_eval = task->bf + shoff; + auto* __restrict__ basis_x_eval = task->dbfx + shoff; + auto* __restrict__ basis_y_eval = task->dbfy + shoff; + auto* __restrict__ basis_z_eval = task->dbfz + shoff; + auto* __restrict__ basis_lapl_eval = task->d2bflapl + shoff; + + // Loop over points in task + // Assign each point to separate thread within the warp + for( int ipt = it.get_local_id(2) % sycl::warp_size; ipt < npts; ipt += sycl::warp_size ) { + //const double3 point = points[ipt]; + double3 point; + point.x = points_x[ipt]; + point.y = points_y[ipt]; + point.z = points_z[ipt]; + + + const auto x = point.x - O.x; + const auto y = point.y - O.y; + const auto z = point.z - O.z; + const auto rsq = x*x + y*y + z*z; + + // Evaluate radial part of bfn + double radial_eval = 0.; + double radial_eval_alpha = 0.; + double radial_eval_alpha_squared = 0.; + + for( uint32_t i = 0; i < nprim; ++i ) { + const auto a = my_alpha[i]; + const auto e = my_coeff[i] * ::sycl::exp( - a * rsq ); + + radial_eval += e; + radial_eval_alpha += a * e; + radial_eval_alpha_squared += a * a * e; + } + + radial_eval_alpha *= -2; + radial_eval_alpha_squared *= 4; + + // Common Subexpressions + const auto x0 = 0.5*sqrt_35; + const auto x1 = x0*y; + const auto x2 = x*x1; + const auto x3 = x*x; + const auto x4 = x3; + const auto x5 = y*y; + const auto x6 = x5; + const auto x7 = -x6; + const auto x8 = x4 + x7; + const auto x9 = 0.25*sqrt_70; + const auto x10 = x9*z; + const auto x11 = x10*y; + const auto x12 = 3.0*x4; + const auto x13 = x12 + x7; + const auto x14 = 0.5*sqrt_5; + const auto x15 = x14*y; + const auto x16 = x*x15; + const auto x17 = z*z; + const auto x18 = x17; + const auto x19 = -6.0*x18; + const auto x20 = x19 + x6; + const auto x21 = -x20 - x4; + const auto x22 = 0.25*sqrt_10; + const auto x23 = x22*z; + const auto x24 = x23*y; + const auto x25 = -4.0*x18; + const auto x26 = 3.0*x6; + const auto x27 = x25 + x26; + const auto x28 = -x12 - x27; + const auto x29 = 0.125*radial_eval; + const auto x30 = x*x*x*x; + const auto x31 = y*y*y*y; + const auto x32 = 6.0*x4*x6; + const auto x33 = x18*x4; + const auto x34 = x18*x6; + const auto x35 = 3.0*x30 + 3.0*x31 + x32 - 24.0*x33 - 24.0*x34 + 8.0*(z*z*z*z); + const auto x36 = x*x23; + const auto x37 = 0.25*sqrt_5; + const auto x38 = -x30 + x31 + 6.0*x33 - 6.0*x34; + const auto x39 = x*x10; + const auto x40 = -x26; + const auto x41 = x4 + x40; + const auto x42 = x30 + x31 - x32; + const auto x43 = radial_eval*x13; + const auto x44 = x4*x8; + const auto x45 = x*x11; + const auto x46 = 6.0*radial_eval; + const auto x47 = radial_eval_alpha*x13; + const auto x48 = x46 + x47; + const auto x49 = -x12 - x20; + const auto x50 = x21*x4; + const auto x51 = -x46; + const auto x52 = x*x24*(radial_eval_alpha*x28 + x51); + const auto x53 = 12.0*radial_eval; + const auto x54 = x*x*x; + const auto x55 = 4.0*x; + const auto x56 = x*x6 - x18*x55 + x54; + const auto x57 = radial_eval_alpha*x; + const auto x58 = 9.0*x4; + const auto x59 = -x27 - x58; + const auto x60 = x28*x4; + const auto x61 = 4.0*radial_eval; + const auto x62 = 3.0*x; + const auto x63 = x18*x62 - x54; + const auto x64 = x12 + x40; + const auto x65 = radial_eval*x64; + const auto x66 = x4*x41; + const auto x67 = radial_eval_alpha*x66 + x65; + const auto x68 = 0.125*sqrt_35; + const auto x69 = x54 - x6*x62; + const auto x70 = x*x0; + const auto x71 = radial_eval*x41; + const auto x72 = x6*x8; + const auto x73 = x13*x6; + const auto x74 = radial_eval_alpha*x73 + x65; + const auto x75 = x*x14; + const auto x76 = x19 + x26; + const auto x77 = -x4 - x76; + const auto x78 = x21*x6; + const auto x79 = 9.0*x6; + const auto x80 = x12 + x25; + const auto x81 = -x79 - x80; + const auto x82 = x28*x6; + const auto x83 = y*y*y; + const auto x84 = 4.0*y; + const auto x85 = -x18*x84 + x4*y + x83; + const auto x86 = radial_eval_alpha*y; + const auto x87 = 3.0*y; + const auto x88 = -x18*x87 + x83; + const auto x89 = radial_eval_alpha*x41; + const auto x90 = x51 + x89; + const auto x91 = -x4*x87 + x83; + const auto x92 = x1*z; + const auto x93 = x9*y; + const auto x94 = x13*x18; + const auto x95 = x22*y; + const auto x96 = -12.0*x18; + const auto x97 = x26 + x96; + const auto x98 = -x12 - x97; + const auto x99 = x18*x28; + const auto x100 = radial_eval*x98 + radial_eval_alpha*x99; + const auto x101 = 3.0*z; + const auto x102 = -x101*x4 - x101*x6 + 2.0*(z*z*z); + const auto x103 = radial_eval_alpha*z; + const auto x104 = x37*z; + const auto x105 = x53*x8; + const auto x106 = x18*x41; + const auto x107 = 2.0*radial_eval_alpha; + const auto x108 = x107*x13; + const auto x109 = radial_eval_alpha + radial_eval_alpha_squared*x4; + const auto x110 = x108 + x109*x8; + const auto x111 = x109*x13; + const auto x112 = 12.0*radial_eval_alpha; + const auto x113 = x112*x4; + const auto x114 = x113 + x46; + const auto x115 = x107*x49 + x109*x21; + const auto x116 = x109*x35 + x53*(x6 + x80) + 24.0*x56*x57; + const auto x117 = -18.0*radial_eval; + const auto x118 = x109*x28; + const auto x119 = x107*x59 + x118; + const auto x120 = -x4; + const auto x121 = 8.0*x57; + const auto x122 = x109*x38 + x121*x63 + x53*(x120 + x18); + const auto x123 = x107*x64; + const auto x124 = x109*x41 + x123; + const auto x125 = x105 + x109*x42 + x121*x69; + const auto x126 = radial_eval_alpha*x3; + const auto x127 = radial_eval_alpha*x5; + const auto x128 = 6.0*radial_eval_alpha; + const auto x129 = x128*x6; + const auto x130 = radial_eval_alpha*x64; + const auto x131 = 24.0*radial_eval; + const auto x132 = x*x131; + const auto x133 = x132*y; + const auto x134 = 12.0*x57; + const auto x135 = 12.0*x86; + const auto x136 = radial_eval_alpha_squared*x; + const auto x137 = x136*y; + const auto x138 = -x128*x4 + x51; + const auto x139 = radial_eval_alpha*x55; + const auto x140 = radial_eval_alpha*x84; + const auto x141 = x*x93; + const auto x142 = x128*x18; + const auto x143 = -x142; + const auto x144 = x*x95*(radial_eval_alpha*x98 + radial_eval_alpha_squared*x99 + x143 + x51); + const auto x145 = 96.0*radial_eval*z; + const auto x146 = 12.0*x103; + const auto x147 = radial_eval_alpha*x17; + const auto x148 = 4.0*radial_eval_alpha; + const auto x149 = x147*x64; + const auto x150 = x68*z; + const auto x151 = x107*x41; + const auto x152 = radial_eval_alpha + radial_eval_alpha_squared*x6; + const auto x153 = x151 + x152*x8; + const auto x154 = x123 + x13*x152; + const auto x155 = x107*x77 + x152*x21; + const auto x156 = x152*x28; + const auto x157 = x107*x81 + x156; + const auto x158 = x152*x35 + x53*(x27 + x4) + 24.0*x85*x86; + const auto x159 = x112*x6; + const auto x160 = x159 + x46; + const auto x161 = 8.0*x86; + const auto x162 = x152*x38 + x161*x88 - x53*(x18 - x6); + const auto x163 = x152*x42 + x161*x91 + x53*(x120 + x6); + const auto x164 = radial_eval_alpha_squared*y; + const auto x165 = radial_eval_alpha + radial_eval_alpha_squared*x18; + const auto x166 = x165*x8; + const auto x167 = x108 + x13*x165; + const auto x168 = 24.0*radial_eval_alpha*x18 + x165*x21; + const auto x169 = x107*x98 + x165*x28; + const auto x170 = x131 + x169; + const auto x171 = -48.0*radial_eval*(-2.0*x18 + x4 + x6) + 32.0*x102*x103 + x165*x35; + const auto x172 = x105 + 24.0*x147*x8 + x165*x38; + const auto x173 = x151 + x165*x41; + const auto x174 = x165*x42; + const auto x175 = -x159; + + + // Evaluate basis function + basis_eval[ipt + 0*npts] = radial_eval*x2*x8; + basis_eval[ipt + 1*npts] = radial_eval*x11*x13; + basis_eval[ipt + 2*npts] = radial_eval*x16*x21; + basis_eval[ipt + 3*npts] = radial_eval*x24*x28; + basis_eval[ipt + 4*npts] = x29*x35; + basis_eval[ipt + 5*npts] = radial_eval*x28*x36; + basis_eval[ipt + 6*npts] = radial_eval*x37*x38; + basis_eval[ipt + 7*npts] = radial_eval*x39*x41; + basis_eval[ipt + 8*npts] = sqrt_35*x29*x42; + + + + // Evaluate first derivative of bfn wrt x + basis_x_eval[ipt + 0*npts] = x1*(radial_eval_alpha*x44 + x43); + basis_x_eval[ipt + 1*npts] = x45*x48; + basis_x_eval[ipt + 2*npts] = x15*(radial_eval*x49 + radial_eval_alpha*x50); + basis_x_eval[ipt + 3*npts] = x52; + basis_x_eval[ipt + 4*npts] = 0.125*x35*x57 + 0.125*x53*x56; + basis_x_eval[ipt + 5*npts] = x23*(radial_eval*x59 + radial_eval_alpha*x60); + basis_x_eval[ipt + 6*npts] = x37*(x38*x57 + x61*x63); + basis_x_eval[ipt + 7*npts] = x10*x67; + basis_x_eval[ipt + 8*npts] = x68*(x42*x57 + x61*x69); + + // Evaluate first derivative of bfn wrt y + basis_y_eval[ipt + 0*npts] = x70*(radial_eval_alpha*x72 + x71); + basis_y_eval[ipt + 1*npts] = x10*x74; + basis_y_eval[ipt + 2*npts] = x75*(radial_eval*x77 + radial_eval_alpha*x78); + basis_y_eval[ipt + 3*npts] = x23*(radial_eval*x81 + radial_eval_alpha*x82); + basis_y_eval[ipt + 4*npts] = 0.125*x35*x86 + 0.125*x53*x85; + basis_y_eval[ipt + 5*npts] = x52; + basis_y_eval[ipt + 6*npts] = x37*(x38*x86 + x61*x88); + basis_y_eval[ipt + 7*npts] = x45*x90; + basis_y_eval[ipt + 8*npts] = x68*(x42*x86 + x61*x91); + + // Evaluate first derivative of bfn wrt z + basis_z_eval[ipt + 0*npts] = x57*x8*x92; + basis_z_eval[ipt + 1*npts] = x93*(radial_eval_alpha*x94 + x43); + basis_z_eval[ipt + 2*npts] = x16*z*(radial_eval_alpha*x21 + x53); + basis_z_eval[ipt + 3*npts] = x100*x95; + basis_z_eval[ipt + 4*npts] = 2.0*radial_eval*x102 + 0.125*x103*x35; + basis_z_eval[ipt + 5*npts] = x*x100*x22; + basis_z_eval[ipt + 6*npts] = x104*(radial_eval_alpha*x38 + x105); + basis_z_eval[ipt + 7*npts] = x*x9*(radial_eval_alpha*x106 + x71); + basis_z_eval[ipt + 8*npts] = x103*x42*x68; + + + // Evaluate Laplacian of bfn + basis_lapl_eval[ipt + 0*npts] = x2*(x110 + x153 + x166); + basis_lapl_eval[ipt + 1*npts] = x11*(x111 + x113 + x154 + x167); + basis_lapl_eval[ipt + 2*npts] = x16*(x115 + x155 + x168); + basis_lapl_eval[ipt + 3*npts] = x24*(-x113 + x118 + x157 + x169); + basis_lapl_eval[ipt + 4*npts] = 0.125*x116 + 0.125*x158 + 0.125*x171; + basis_lapl_eval[ipt + 5*npts] = x36*(x119 + x156 + x169 + x175); + basis_lapl_eval[ipt + 6*npts] = x37*(x122 + x162 + x172); + basis_lapl_eval[ipt + 7*npts] = x39*(x124 + x152*x41 + x173 + x175); + basis_lapl_eval[ipt + 8*npts] = x68*(x125 + x163 + x174); + + + + + +#if 0 + // Evaluate the angular part of bfn + + + + double ang_eval_0; + double ang_eval_1; + double ang_eval_2; + double ang_eval_3; + + + ang_eval_0 = radial_eval*x2*x8; + ang_eval_1 = radial_eval*x11*x13; + ang_eval_2 = radial_eval*x16*x21; + ang_eval_3 = radial_eval*x24*x28; + basis_eval[ipt + 0*npts] = ang_eval_0; + basis_eval[ipt + 1*npts] = ang_eval_1; + basis_eval[ipt + 2*npts] = ang_eval_2; + basis_eval[ipt + 3*npts] = ang_eval_3; + + ang_eval_0 = x29*x35; + ang_eval_1 = radial_eval*x28*x36; + ang_eval_2 = radial_eval*x37*x38; + ang_eval_3 = radial_eval*x39*x41; + basis_eval[ipt + 4*npts] = ang_eval_0; + basis_eval[ipt + 5*npts] = ang_eval_1; + basis_eval[ipt + 6*npts] = ang_eval_2; + basis_eval[ipt + 7*npts] = ang_eval_3; + + ang_eval_0 = sqrt_35*x29*x42; + basis_eval[ipt + 8*npts] = ang_eval_0; + + + double dang_eval_x_0, dang_eval_y_0, dang_eval_z_0; + double dang_eval_x_1, dang_eval_y_1, dang_eval_z_1; + double dang_eval_x_2, dang_eval_y_2, dang_eval_z_2; + double dang_eval_x_3, dang_eval_y_3, dang_eval_z_3; + + dang_eval_x_0 = x1*(radial_eval_alpha*x44 + x43); + dang_eval_y_0 = x70*(radial_eval_alpha*x72 + x71); + dang_eval_z_0 = x57*x8*x92; + dang_eval_x_1 = x45*x48; + dang_eval_y_1 = x10*x74; + dang_eval_z_1 = x93*(radial_eval_alpha*x94 + x43); + dang_eval_x_2 = x15*(radial_eval*x49 + radial_eval_alpha*x50); + dang_eval_y_2 = x75*(radial_eval*x77 + radial_eval_alpha*x78); + dang_eval_z_2 = x16*z*(radial_eval_alpha*x21 + x53); + dang_eval_x_3 = x52; + dang_eval_y_3 = x23*(radial_eval*x81 + radial_eval_alpha*x82); + dang_eval_z_3 = x100*x95; + basis_x_eval[ipt + 0*npts] = dang_eval_x_0; + basis_y_eval[ipt + 0*npts] = dang_eval_y_0; + basis_z_eval[ipt + 0*npts] = dang_eval_z_0; + basis_x_eval[ipt + 1*npts] = dang_eval_x_1; + basis_y_eval[ipt + 1*npts] = dang_eval_y_1; + basis_z_eval[ipt + 1*npts] = dang_eval_z_1; + basis_x_eval[ipt + 2*npts] = dang_eval_x_2; + basis_y_eval[ipt + 2*npts] = dang_eval_y_2; + basis_z_eval[ipt + 2*npts] = dang_eval_z_2; + basis_x_eval[ipt + 3*npts] = dang_eval_x_3; + basis_y_eval[ipt + 3*npts] = dang_eval_y_3; + basis_z_eval[ipt + 3*npts] = dang_eval_z_3; + + dang_eval_x_0 = 0.125*x35*x57 + 0.125*x53*x56; + dang_eval_y_0 = 0.125*x35*x86 + 0.125*x53*x85; + dang_eval_z_0 = 2.0*radial_eval*x102 + 0.125*x103*x35; + dang_eval_x_1 = x23*(radial_eval*x59 + radial_eval_alpha*x60); + dang_eval_y_1 = x52; + dang_eval_z_1 = x*x100*x22; + dang_eval_x_2 = x37*(x38*x57 + x61*x63); + dang_eval_y_2 = x37*(x38*x86 + x61*x88); + dang_eval_z_2 = x104*(radial_eval_alpha*x38 + x105); + dang_eval_x_3 = x10*x67; + dang_eval_y_3 = x45*x90; + dang_eval_z_3 = x*x9*(radial_eval_alpha*x106 + x71); + basis_x_eval[ipt + 4*npts] = dang_eval_x_0; + basis_y_eval[ipt + 4*npts] = dang_eval_y_0; + basis_z_eval[ipt + 4*npts] = dang_eval_z_0; + basis_x_eval[ipt + 5*npts] = dang_eval_x_1; + basis_y_eval[ipt + 5*npts] = dang_eval_y_1; + basis_z_eval[ipt + 5*npts] = dang_eval_z_1; + basis_x_eval[ipt + 6*npts] = dang_eval_x_2; + basis_y_eval[ipt + 6*npts] = dang_eval_y_2; + basis_z_eval[ipt + 6*npts] = dang_eval_z_2; + basis_x_eval[ipt + 7*npts] = dang_eval_x_3; + basis_y_eval[ipt + 7*npts] = dang_eval_y_3; + basis_z_eval[ipt + 7*npts] = dang_eval_z_3; + + dang_eval_x_0 = x68*(x42*x57 + x61*x69); + dang_eval_y_0 = x68*(x42*x86 + x61*x91); + dang_eval_z_0 = x103*x42*x68; + basis_x_eval[ipt + 8*npts] = dang_eval_x_0; + basis_y_eval[ipt + 8*npts] = dang_eval_y_0; + basis_z_eval[ipt + 8*npts] = dang_eval_z_0; + +#endif + } // Loop over points within task + } // Loop over tasks + + } // Loop over shells +} // end kernel + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_device.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_device.cxx new file mode 100644 index 000000000..3b4d5b792 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_device.cxx @@ -0,0 +1,679 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include +#include "device_specific/sycl_util.hpp" +#include "exceptions/sycl_exception.hpp" +#include +#include + +#include "device/common/collocation_device.hpp" +#include "device/sycl/kernels/collocation_masked_kernels.hpp" +#include "device/sycl/kernels/collocation_masked_combined_kernels.hpp" +#include "device/sycl/kernels/collocation_shell_to_task_kernels.hpp" + +#include "device_specific/sycl_device_constants.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" + +#include + +#define GAUXC_SYCL_MAX_L 4 + +namespace GauXC { + +namespace detail { + +template +void launch_shell_to_task( ::sycl::queue& queue, uint32_t ntask_average, + uint32_t nshells, KernelOp op, Args&&... args ) { + + static_assert( nthreads == nwarp * sycl::warp_size, + "Shell scratch must have one row per sub-group" ); + + if( !nshells ) return; + + const uint32_t nwarp_per_block = nthreads / sycl::warp_size; + const uint32_t n_task_blocks = util::div_ceil( ntask_average, nwarp_per_block ); + + ::sycl::range<3> global( nshells, 1, n_task_blocks * nthreads ); + ::sycl::range<3> local ( 1, 1, nthreads ); + + GAUXC_SYCL_ERROR( "Collocation Shell-To-Task Launch Failed", + queue.parallel_for( ::sycl::nd_range<3>(global, local), + [=](::sycl::nd_item<3>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + op( nshells, args... ); + }) + ); + +} + +} + + +template +void eval_collocation_masked( + size_t nshells, + size_t nbf, + size_t npts, + const Shell* shells_device, + const size_t* mask_device, + const size_t* offs_device, + const T* pts_device, + T* eval_device, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + const auto nmax_threads = util::sycl_kernel_max_threads_per_block( stream ); + const auto max_warps_per_thread_block = nmax_threads / sycl::warp_size; + + ::sycl::range<2> local( max_warps_per_thread_block, sycl::warp_size ); + ::sycl::range<2> global( + util::div_ceil( nshells, local[0] ) * local[0], + util::div_ceil( npts, local[1] ) * local[1] ); + + GAUXC_SYCL_ERROR( "Collocation Masked Launch Failed", + stream.parallel_for( ::sycl::nd_range<2>(global, local), + [=](::sycl::nd_item<2>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + collocation_device_masked_kernel( nshells, nbf, npts, + shells_device, mask_device, offs_device, pts_device, eval_device ); + }) + ); + +} + +template +void eval_collocation_masked( + size_t nshells, + size_t nbf, + size_t npts, + const Shell* shells_device, + const size_t* mask_device, + const size_t* offs_device, + const double* pts_device, + double* eval_device, + device_queue queue +); + + + + +template +void eval_collocation_masked_combined( + size_t ntasks, + size_t npts_max, + size_t nshells_max, + Shell* shells_device, + XCDeviceTask* device_tasks, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + const auto nmax_threads = util::sycl_kernel_max_threads_per_block( stream ); + const auto max_warps_per_thread_block = nmax_threads / sycl::warp_size; + + ::sycl::range<3> local( 1, max_warps_per_thread_block, sycl::warp_size ); + ::sycl::range<3> global( ntasks, + util::div_ceil( nshells_max, local[1] ) * local[1], + util::div_ceil( npts_max, local[2] ) * local[2] ); + + GAUXC_SYCL_ERROR( "Collocation Masked Combined Launch Failed", + stream.parallel_for( ::sycl::nd_range<3>(global, local), + [=](::sycl::nd_item<3>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + collocation_device_masked_combined_kernel( ntasks, + shells_device, device_tasks ); + }) + ); + +} + +template +void eval_collocation_masked_combined( + size_t ntasks, + size_t npts_max, + size_t nshells_max, + Shell* shells_device, + XCDeviceTask* device_tasks, + device_queue queue +); + + + + +template +void eval_collocation_masked_deriv1( + size_t nshells, + size_t nbf, + size_t npts, + const Shell* shells_device, + const size_t* mask_device, + const size_t* offs_device, + const T* pts_device, + T* eval_device, + T* deval_device_x, + T* deval_device_y, + T* deval_device_z, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + const auto nmax_threads = util::sycl_kernel_max_threads_per_block( stream ); + const auto max_warps_per_thread_block = nmax_threads / sycl::warp_size; + + ::sycl::range<2> local( max_warps_per_thread_block, sycl::warp_size ); + ::sycl::range<2> global( + util::div_ceil( nshells, local[0] ) * local[0], + util::div_ceil( npts, local[1] ) * local[1] ); + + GAUXC_SYCL_ERROR( "Collocation Masked Deriv1 Launch Failed", + stream.parallel_for( ::sycl::nd_range<2>(global, local), + [=](::sycl::nd_item<2>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + collocation_device_masked_kernel_deriv1( nshells, nbf, npts, + shells_device, mask_device, offs_device, pts_device, eval_device, + deval_device_x, deval_device_y, deval_device_z ); + }) + ); + +} + +template +void eval_collocation_masked_deriv1( + size_t nshells, + size_t nbf, + size_t npts, + const Shell* shells_device, + const size_t* mask_device, + const size_t* offs_device, + const double* pts_device, + double* eval_device, + double* deval_device_x, + double* deval_device_y, + double* deval_device_z, + device_queue queue +); + + + + +template +void eval_collocation_masked_combined_deriv1( + size_t ntasks, + size_t npts_max, + size_t nshells_max, + Shell* shells_device, + XCDeviceTask* device_tasks, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + const auto nmax_threads = util::sycl_kernel_max_threads_per_block( stream ); + + ::sycl::range<3> local( 1, nmax_threads / sycl::warp_size, sycl::warp_size ); + ::sycl::range<3> global( ntasks, + util::div_ceil( nshells_max, local[1] ) * local[1], + util::div_ceil( npts_max, local[2] ) * local[2] ); + + GAUXC_SYCL_ERROR( "Collocation Masked Combined Deriv1 Launch Failed", + stream.parallel_for( ::sycl::nd_range<3>(global, local), + [=](::sycl::nd_item<3>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + collocation_device_masked_combined_kernel_deriv1( ntasks, + shells_device, device_tasks ); + }) + ); + +} + +template +void eval_collocation_masked_combined_deriv1( + size_t ntasks, + size_t npts_max, + size_t nshells_max, + Shell* shells_device, + XCDeviceTask* device_tasks, + device_queue queue +); + + +template +void dispatch_shell_to_task_collocation( ::sycl::queue& queue, int32_t l, + bool pure, uint32_t ntask_average, uint32_t nshells, Args&&... args ) { + + switch(l) { + case 0: + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_0( nsh, a... ); + }, std::forward(args)... ); + break; + case 1: + if( pure ) + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_1( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_1( nsh, a... ); + }, std::forward(args)... ); + break; + case 2: + if( pure ) + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_2( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_2( nsh, a... ); + }, std::forward(args)... ); + break; + case 3: + if( pure ) + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_3( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_3( nsh, a... ); + }, std::forward(args)... ); + break; + case 4: + if( pure ) + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_4( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_4( nsh, a... ); + }, std::forward(args)... ); + break; + default: GAUXC_GENERIC_EXCEPTION("SYCL L_MAX = 4"); + } +} + + +void eval_collocation_shell_to_task( + uint32_t max_l, + AngularMomentumShellToTaskBatch* l_batched_shell_to_task, + XCDeviceTask* device_tasks, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + for( auto l = 0u; l <= max_l; ++l ) { + auto pure = l_batched_shell_to_task[l].pure; + auto shell_to_task_device = l_batched_shell_to_task[l].shell_to_task_device; + auto nshells = l_batched_shell_to_task[l].nshells_in_batch; + auto ntask_average = std::max(1ul, l_batched_shell_to_task[l].ntask_average); + dispatch_shell_to_task_collocation( stream, l, pure, ntask_average, + nshells, shell_to_task_device, device_tasks ); + } + +} + + +template +void dispatch_shell_to_task_collocation_gradient( ::sycl::queue& queue, int32_t l, + bool pure, uint32_t ntask_average, uint32_t nshells, Args&&... args ) { + + switch(l) { + case 0: + detail::launch_shell_to_task<512,16>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_gradient_0( nsh, a... ); + }, std::forward(args)... ); + break; + case 1: + if( pure ) + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_gradient_1( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_gradient_1( nsh, a... ); + }, std::forward(args)... ); + break; + case 2: + if( pure ) + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_gradient_2( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_gradient_2( nsh, a... ); + }, std::forward(args)... ); + break; + case 3: + if( pure ) + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_gradient_3( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_gradient_3( nsh, a... ); + }, std::forward(args)... ); + break; + case 4: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_gradient_4( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_gradient_4( nsh, a... ); + }, std::forward(args)... ); + break; + default: GAUXC_GENERIC_EXCEPTION("SYCL L_MAX = 4"); + } +} + + +void eval_collocation_shell_to_task_gradient( + uint32_t max_l, + AngularMomentumShellToTaskBatch* l_batched_shell_to_task, + XCDeviceTask* device_tasks, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + for( auto l = 0u; l <= max_l; ++l ) { + auto pure = l_batched_shell_to_task[l].pure; + auto shell_to_task_device = l_batched_shell_to_task[l].shell_to_task_device; + auto nshells = l_batched_shell_to_task[l].nshells_in_batch; + auto ntask_average = std::max(1ul, l_batched_shell_to_task[l].ntask_average); + dispatch_shell_to_task_collocation_gradient( stream, l, pure, ntask_average, + nshells, shell_to_task_device, device_tasks ); + } + +} + + +template +void dispatch_shell_to_task_collocation_hessian( ::sycl::queue& queue, int32_t l, + bool pure, uint32_t ntask_average, uint32_t nshells, Args&&... args ) { + + switch(l) { + case 0: + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_hessian_0( nsh, a... ); + }, std::forward(args)... ); + break; + case 1: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_hessian_1( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_hessian_1( nsh, a... ); + }, std::forward(args)... ); + break; + case 2: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_hessian_2( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_hessian_2( nsh, a... ); + }, std::forward(args)... ); + break; + case 3: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_hessian_3( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_hessian_3( nsh, a... ); + }, std::forward(args)... ); + break; + case 4: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_hessian_4( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_hessian_4( nsh, a... ); + }, std::forward(args)... ); + break; + default: GAUXC_GENERIC_EXCEPTION("SYCL L_MAX = 4"); + } +} + + +void eval_collocation_shell_to_task_hessian( + uint32_t max_l, + AngularMomentumShellToTaskBatch* l_batched_shell_to_task, + XCDeviceTask* device_tasks, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + for( auto l = 0u; l <= max_l; ++l ) { + auto pure = l_batched_shell_to_task[l].pure; + auto shell_to_task_device = l_batched_shell_to_task[l].shell_to_task_device; + auto nshells = l_batched_shell_to_task[l].nshells_in_batch; + auto ntask_average = std::max(1ul, l_batched_shell_to_task[l].ntask_average); + dispatch_shell_to_task_collocation_hessian( stream, l, pure, ntask_average, + nshells, shell_to_task_device, device_tasks ); + } + +} + + +template +void dispatch_shell_to_task_collocation_laplacian( ::sycl::queue& queue, int32_t l, + bool pure, uint32_t ntask_average, uint32_t nshells, Args&&... args ) { + + switch(l) { + case 0: + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_laplacian_0( nsh, a... ); + }, std::forward(args)... ); + break; + case 1: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_laplacian_1( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_laplacian_1( nsh, a... ); + }, std::forward(args)... ); + break; + case 2: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_laplacian_2( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_laplacian_2( nsh, a... ); + }, std::forward(args)... ); + break; + case 3: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_laplacian_3( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_laplacian_3( nsh, a... ); + }, std::forward(args)... ); + break; + case 4: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_laplacian_4( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_laplacian_4( nsh, a... ); + }, std::forward(args)... ); + break; + default: GAUXC_GENERIC_EXCEPTION("SYCL L_MAX = 4"); + } +} + + +void eval_collocation_shell_to_task_laplacian( + uint32_t max_l, + AngularMomentumShellToTaskBatch* l_batched_shell_to_task, + XCDeviceTask* device_tasks, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + for( auto l = 0u; l <= max_l; ++l ) { + auto pure = l_batched_shell_to_task[l].pure; + auto shell_to_task_device = l_batched_shell_to_task[l].shell_to_task_device; + auto nshells = l_batched_shell_to_task[l].nshells_in_batch; + auto ntask_average = std::max(1ul, l_batched_shell_to_task[l].ntask_average); + dispatch_shell_to_task_collocation_laplacian( stream, l, pure, ntask_average, + nshells, shell_to_task_device, device_tasks ); + } + +} + + +template +void dispatch_shell_to_task_collocation_lapgrad( ::sycl::queue& queue, int32_t l, + bool pure, uint32_t ntask_average, uint32_t nshells, Args&&... args ) { + + switch(l) { + case 0: + detail::launch_shell_to_task<256,8>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_lapgrad_0( nsh, a... ); + }, std::forward(args)... ); + break; + case 1: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_lapgrad_1( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_lapgrad_1( nsh, a... ); + }, std::forward(args)... ); + break; + case 2: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_lapgrad_2( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_lapgrad_2( nsh, a... ); + }, std::forward(args)... ); + break; + case 3: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_lapgrad_3( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_lapgrad_3( nsh, a... ); + }, std::forward(args)... ); + break; + case 4: + if( pure ) + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_spherical_lapgrad_4( nsh, a... ); + }, std::forward(args)... ); + else + detail::launch_shell_to_task<128,4>( queue, ntask_average, nshells, + [=](uint32_t nsh, auto... a) { + collocation_device_shell_to_task_kernel_cartesian_lapgrad_4( nsh, a... ); + }, std::forward(args)... ); + break; + default: GAUXC_GENERIC_EXCEPTION("SYCL L_MAX = 4"); + } +} + + +void eval_collocation_shell_to_task_lapgrad( + uint32_t max_l, + AngularMomentumShellToTaskBatch* l_batched_shell_to_task, + XCDeviceTask* device_tasks, + device_queue queue +) { + + ::sycl::queue& stream = queue.queue_as(); + + for( auto l = 0u; l <= max_l; ++l ) { + auto pure = l_batched_shell_to_task[l].pure; + auto shell_to_task_device = l_batched_shell_to_task[l].shell_to_task_device; + auto nshells = l_batched_shell_to_task[l].nshells_in_batch; + auto ntask_average = std::max(1ul, l_batched_shell_to_task[l].ntask_average); + dispatch_shell_to_task_collocation_lapgrad( stream, l, pure, ntask_average, + nshells, shell_to_task_device, device_tasks ); + } + +} + + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_masked_combined_kernels.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_masked_combined_kernels.hpp new file mode 100644 index 000000000..56954eff9 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_masked_combined_kernels.hpp @@ -0,0 +1,203 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include + +#include + +#include "device/sycl/kernels/collocation/collocation_angular_cartesian.hpp" +#include "device/sycl/kernels/collocation/collocation_angular_spherical_unnorm.hpp" +//#include "device/cuda/kernels/cuda_alg_variant_control.hpp" +#include "device/xc_device_task.hpp" +#include "device/sycl/kernels/sycl_launch.hpp" + +namespace GauXC { + + +template +void collocation_device_masked_combined_kernel( + size_t ntasks, + Shell* __restrict__ shells_device, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + const size_t tid_x = it.get_global_id(2); + const size_t tid_y = it.get_global_id(1); + + if( it.get_group(0) < ntasks ) { + + auto& task = device_tasks[ it.get_group(0) ]; + + const auto nshells = task.bfn_screening.nshells; + //const auto nbf = task.bfn_screening.nbe; + const auto npts = task.npts; + //const auto* __restrict__ pts_device = task.points; + const auto* __restrict__ pts_x_device = task.points_x; + const auto* __restrict__ pts_y_device = task.points_y; + const auto* __restrict__ pts_z_device = task.points_z; + const auto* __restrict__ mask_device = task.bfn_screening.shell_list; + const auto* __restrict__ offs_device = task.bfn_screening.shell_offs; + + auto* __restrict__ eval_device = task.bf; + + if( tid_x < npts and tid_y < nshells ) { + + const uint32_t ipt = tid_x; + const uint32_t ish = tid_y; + const uint32_t ibf = offs_device[ish]; + + const auto& shell = shells_device[mask_device[ish]]; + //const auto* pt = pts_device + 3*ipt; + const auto pt_x = pts_x_device[ipt]; + const auto pt_y = pts_y_device[ipt]; + const auto pt_z = pts_z_device[ipt]; + + + const auto* __restrict__ O = shell.O_data(); + const auto* __restrict__ alpha = shell.alpha_data(); + const auto* __restrict__ coeff = shell.coeff_data(); + + const auto xc = pt_x - O[0]; + const auto yc = pt_y - O[1]; + const auto zc = pt_z - O[2]; + + const auto rsq = xc*xc + yc*yc + zc*zc; + + const uint32_t nprim = shell.nprim(); + auto tmp = 0.; + for( uint32_t i = 0; i < nprim; ++i ) + tmp += coeff[i] * ::sycl::exp( - alpha[i] * rsq ); + + auto * bf_eval = eval_device + ibf*npts + ipt; + + const bool do_sph = shell.pure(); + if( do_sph ) + collocation_spherical_unnorm_angular( npts, shell.l(), tmp, xc, yc, zc, bf_eval ); + else + collocation_cartesian_angular( npts, shell.l(), tmp, xc, yc, zc, bf_eval ); + + } // shell / point idx check + + } // Batch idx check + +} + + + + + + + + + + + + + + + +template +void collocation_device_masked_combined_kernel_deriv1( + size_t ntasks, + Shell* __restrict__ shells_device, + XCDeviceTask* __restrict__ device_tasks +) { + + auto it = GauXC::sycl::this_item(); + // DBWY: These are factored into the loop for this optimization + const size_t tid_x = it.get_global_id(2); + const size_t tid_y = it.get_global_id(1); + + if( it.get_group(0) < ntasks ) { + + auto& task = device_tasks[ it.get_group(0) ]; + + const auto nshells = task.bfn_screening.nshells; + //const auto nbf = task.bfn_screening.nbe; + const auto npts = task.npts; + //const auto* __restrict__ pts_device = task.points; + const auto* __restrict__ pts_x_device = task.points_x; + const auto* __restrict__ pts_y_device = task.points_y; + const auto* __restrict__ pts_z_device = task.points_z; + const auto* __restrict__ mask_device = task.bfn_screening.shell_list; + const auto* __restrict__ offs_device = task.bfn_screening.shell_offs; + + auto* __restrict__ eval_device = task.bf; + auto* __restrict__ deval_device_x = task.dbfx; + auto* __restrict__ deval_device_y = task.dbfy; + auto* __restrict__ deval_device_z = task.dbfz; + + if( tid_y < nshells and tid_x < npts ) { + + const uint32_t ish = tid_y; + const uint32_t ipt = tid_x; + const uint32_t ibf = offs_device[ish]; + + const auto& shell = shells_device[mask_device[ish]]; + + //const auto* pt = pts_device + 3*ipt; + const auto pt_x = pts_x_device[ipt]; + const auto pt_y = pts_y_device[ipt]; + const auto pt_z = pts_z_device[ipt]; + + + const auto* __restrict__ O = shell.O_data(); + const auto* __restrict__ alpha = shell.alpha_data(); + const auto* __restrict__ coeff = shell.coeff_data(); + + const auto xc = pt_x - O[0]; + const auto yc = pt_y - O[1]; + const auto zc = pt_z - O[2]; + + const auto rsq = xc*xc + yc*yc + zc*zc; + + const uint32_t nprim = shell.nprim(); + auto tmp = 0.; + auto tmp_x = 0., tmp_y = 0., tmp_z = 0.; + for( uint32_t i = 0; i < nprim; ++i ) { + + const auto a = alpha[i]; + const auto e = coeff[i] * ::sycl::exp( - a * rsq ); + + const auto ae = 2. * a * e; + + tmp += e; + tmp_x -= ae * xc; + tmp_y -= ae * yc; + tmp_z -= ae * zc; + + } + + auto * bf_eval = eval_device + ibf*npts + ipt; + auto * dx_eval = deval_device_x + ibf*npts + ipt; + auto * dy_eval = deval_device_y + ibf*npts + ipt; + auto * dz_eval = deval_device_z + ibf*npts + ipt; + + const bool do_sph = shell.pure(); + if( do_sph ) + collocation_spherical_unnorm_angular_deriv1( npts, shell.l(), tmp, tmp_x, tmp_y, + tmp_z, xc, yc, zc, bf_eval, dx_eval, + dy_eval, dz_eval ); + else + collocation_cartesian_angular_deriv1( npts, shell.l(), tmp, tmp_x, tmp_y, tmp_z, + xc, yc, zc, bf_eval, dx_eval, + dy_eval, dz_eval ); + + } // shell / point idx check + } // Batch idx check + + +} + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_masked_kernels.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_masked_kernels.hpp new file mode 100644 index 000000000..3c26a4b9e --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_masked_kernels.hpp @@ -0,0 +1,163 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include + +#include + +#include "device/sycl/kernels/collocation/collocation_angular_cartesian.hpp" +#include "device/sycl/kernels/collocation/collocation_angular_spherical_unnorm.hpp" + +namespace GauXC { + + +template +void collocation_device_masked_kernel( + size_t nshells, + size_t nbf, + size_t npts, + const Shell* __restrict__ shells_device, + const size_t* __restrict__ mask_device, + const size_t* __restrict__ offs_device, + const T* __restrict__ pts_device, + T* __restrict__ eval_device +) { + + auto it = ::sycl::ext::oneapi::this_work_item::get_nd_item<2>(); + const size_t tid_x = it.get_global_id(1); + const size_t tid_y = it.get_global_id(0); + + if( tid_x < npts and tid_y < nshells ) { + + const size_t ipt = tid_x; + const size_t ish = tid_y; + + const size_t ibf = offs_device[ish]; + + const auto& shell = shells_device[mask_device[ish]]; + const auto* pt = pts_device + 3*ipt; + + + const auto* __restrict__ O = shell.O_data(); + const auto* __restrict__ alpha = shell.alpha_data(); + const auto* __restrict__ coeff = shell.coeff_data(); + + const auto xc = pt[0] - O[0]; + const auto yc = pt[1] - O[1]; + const auto zc = pt[2] - O[2]; + + const auto rsq = xc*xc + yc*yc + zc*zc; + + const size_t nprim = shell.nprim(); + auto tmp = 0.; + for( size_t i = 0; i < nprim; ++i ) + tmp += coeff[i] * ::sycl::exp( - alpha[i] * rsq ); + + auto * bf_eval = eval_device + ibf*npts + ipt; + + const bool do_sph = shell.pure(); + if( do_sph ) + collocation_spherical_unnorm_angular( npts, shell.l(), tmp, xc, yc, zc, bf_eval ); + else + collocation_cartesian_angular( npts, shell.l(), tmp, xc, yc, zc, bf_eval ); + + } + +} + + + + + + + + +template +void collocation_device_masked_kernel_deriv1( + size_t nshells, + size_t nbf, + size_t npts, + const Shell* __restrict__ shells_device, + const size_t* __restrict__ mask_device, + const size_t* __restrict__ offs_device, + const T* __restrict__ pts_device, + T* __restrict__ eval_device, + T* __restrict__ deval_device_x, + T* __restrict__ deval_device_y, + T* __restrict__ deval_device_z +) { + + auto it = ::sycl::ext::oneapi::this_work_item::get_nd_item<2>(); + const size_t tid_x = it.get_global_id(1); + const size_t tid_y = it.get_global_id(0); + + if( tid_x < npts and tid_y < nshells ) { + + const size_t ipt = tid_x; + const size_t ish = tid_y; + + const size_t ibf = offs_device[ish]; + + const auto& shell = shells_device[mask_device[ish]]; + const auto* pt = pts_device + 3*ipt; + + + const auto* __restrict__ O = shell.O_data(); + const auto* __restrict__ alpha = shell.alpha_data(); + const auto* __restrict__ coeff = shell.coeff_data(); + + const auto xc = pt[0] - O[0]; + const auto yc = pt[1] - O[1]; + const auto zc = pt[2] - O[2]; + + const auto rsq = xc*xc + yc*yc + zc*zc; + + const size_t nprim = shell.nprim(); + auto tmp = 0.; + auto tmp_x = 0., tmp_y = 0., tmp_z = 0.; + for( size_t i = 0; i < nprim; ++i ) { + + const auto a = alpha[i]; + const auto e = coeff[i] * ::sycl::exp( - a * rsq ); + + const auto ae = 2. * a * e; + + tmp += e; + tmp_x -= ae * xc; + tmp_y -= ae * yc; + tmp_z -= ae * zc; + + } + + auto * bf_eval = eval_device + ibf*npts + ipt; + auto * dx_eval = deval_device_x + ibf*npts + ipt; + auto * dy_eval = deval_device_y + ibf*npts + ipt; + auto * dz_eval = deval_device_z + ibf*npts + ipt; + + const bool do_sph = shell.pure(); + if( do_sph ) + collocation_spherical_unnorm_angular_deriv1( npts, shell.l(), tmp, tmp_x, tmp_y, tmp_z, + xc, yc, zc, bf_eval, dx_eval, + dy_eval, dz_eval ); + else + collocation_cartesian_angular_deriv1( npts, shell.l(), tmp, tmp_x, tmp_y, tmp_z, + xc, yc, zc, bf_eval, dx_eval, + dy_eval, dz_eval ); + + } + + +} + +} // namespace GauXC + diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_shell_to_task_kernels.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_shell_to_task_kernels.hpp new file mode 100644 index 000000000..e18494b80 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/collocation_shell_to_task_kernels.hpp @@ -0,0 +1,82 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + + +#include "collocation/collocation_shell_to_task_kernels_cartesian_l0.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l1.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l2.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l3.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l4.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_cartesian_l0_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l1_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l2_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l3_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l4_gradient.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_cartesian_l0_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l1_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l2_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l3_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l4_hessian.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_cartesian_l0_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l1_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l2_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l3_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l4_laplacian.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_cartesian_l0_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l1_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l2_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l3_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_cartesian_l4_lapgrad.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_spherical_l0.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l1.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l2.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l3.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l4.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_spherical_l0_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l1_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l2_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l3_gradient.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l4_gradient.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_spherical_l0_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l1_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l2_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l3_hessian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l4_hessian.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_spherical_l0_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l1_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l2_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l3_laplacian.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l4_laplacian.hpp" + + +#include "collocation/collocation_shell_to_task_kernels_spherical_l0_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l1_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l2_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l3_lapgrad.hpp" +#include "collocation/collocation_shell_to_task_kernels_spherical_l4_lapgrad.hpp" diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/grid_to_center.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/grid_to_center.cxx new file mode 100644 index 000000000..833c79371 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/grid_to_center.cxx @@ -0,0 +1,87 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include +#include "grid_to_center.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_util.hpp" +#include "device_specific/sycl_vector_types.hpp" + +namespace GauXC { + +void compute_grid_to_center_dist( int32_t npts, int32_t natoms, + const double* coords, const double* points_x, const double* points_y, + const double* points_z, double* dist, int32_t lddist, ::sycl::queue& stream ) { + + const uint32_t distance_thread_y = sycl::max_warps_per_thread_block / 2; + + ::sycl::range<2> local( distance_thread_y, sycl::warp_size ); + ::sycl::range<2> global( + util::div_ceil( npts, local[0] * distance_thread_y ) * local[0], + util::div_ceil( natoms, local[1] ) * local[1] ); + + GAUXC_SYCL_ERROR( "Grid-To-Center Launch Failed", + stream.submit([&](::sycl::handler& cgh) { + cgh.parallel_for( ::sycl::nd_range<2>(global, local), + [=](::sycl::nd_item<2>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + + auto it = ::sycl::ext::oneapi::this_work_item::get_nd_item<2>(); + auto& point_buffer = *::sycl::ext::oneapi::group_local_memory_for_overwrite( + ::sycl::ext::oneapi::this_work_item::get_work_group<2>() ); + + const auto tid_x = it.get_local_id(1); + const auto tid_y = it.get_local_id(0); + + double3 coord_reg{0.,0.,0.}; + + const int natoms_block = (natoms + sycl::warp_size-1) / sycl::warp_size; + const int coords_block = (npts + sycl::warp_size-1) / sycl::warp_size; + + const double3* coords_vec = (const double3*) coords; + + for (size_t j = it.get_group(1); j < (size_t)natoms_block; j += it.get_group_range(1)) { + const int iAtom = j * sycl::warp_size + tid_x; + // Load blocks into registers/local memory + if (iAtom < natoms) { + coord_reg = coords_vec[iAtom]; + } + for (size_t i = it.get_group(0); i < (size_t)coords_block; i += it.get_group_range(0)) { + const int iPt_load = i * sycl::warp_size + tid_x; + if (iPt_load < npts) { + point_buffer[tid_x].x = points_x[iPt_load]; + point_buffer[tid_x].y = points_y[iPt_load]; + point_buffer[tid_x].z = points_z[iPt_load]; + } + ::sycl::group_barrier( it.get_group() ); + + // do the computation + for (uint32_t k = tid_y; k < sycl::warp_size; k += sycl::warp_size/2) { + const int iPt_sm = k; + const int iPt = i * sycl::warp_size + iPt_sm; + const double rx = point_buffer[iPt_sm].x - coord_reg.x; + const double ry = point_buffer[iPt_sm].y - coord_reg.y; + const double rz = point_buffer[iPt_sm].z - coord_reg.z; + + if (iAtom < natoms and iPt < npts) { + dist[ iAtom + iPt * lddist ] = ::sycl::sqrt( rx*rx + ry*ry + rz*rz ); + } + } + ::sycl::group_barrier( it.get_group() ); + } + } + }); + }) + ); + +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/grid_to_center.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/grid_to_center.hpp new file mode 100644 index 000000000..e14816802 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/grid_to_center.hpp @@ -0,0 +1,21 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include + +namespace GauXC { + +void compute_grid_to_center_dist( int32_t npts, int32_t natoms, + const double* coords, const double* points_x, const double* points_y, + const double* points_z, double* dist, int32_t lddist, ::sycl::queue& stream ); + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/increment_exc_grad.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/increment_exc_grad.cxx new file mode 100644 index 000000000..e73109604 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/increment_exc_grad.cxx @@ -0,0 +1,997 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/increment_exc_grad.hpp" +#include "sycl_extensions.hpp" +#include "sycl_atomics.hpp" +#include "sycl_launch.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include +#include "device_specific/sycl_util.hpp" +#include + +namespace GauXC { + +template +void increment_exc_grad_lda_rks_kernel( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks, + double* __restrict__ EXC_GRAD +) { + + auto it = GauXC::sycl::this_item(); + + for( uint32_t ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + const int iCen = shell_to_task[ish].center_idx; + const uint32_t shsz = shell->size(); + + const int global_warp_id = + it.get_global_id(2) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + double g_acc_x(0), g_acc_y(0), g_acc_z(0); + for( uint32_t itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + const int iParent = task->iParent; + if constexpr( with_weight_derivatives ) { + if( iCen == iParent ) + continue; + } + double g_acc_x_task(0), g_acc_y_task(0), g_acc_z_task(0); + + const auto* __restrict__ basis_x_eval = task->dbfx + shoff; + const auto* __restrict__ basis_y_eval = task->dbfy + shoff; + const auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + const auto* __restrict__ xmat = task->zmat + shoff; + const auto* __restrict__ vrho = task->vrho; + + for( uint32_t ipt = it.get_local_id(2) % sycl::warp_size; + ipt < npts; + ipt += sycl::warp_size ) { + + const double vrho_i = vrho[ipt]; + for( uint32_t ibf = 0; ibf < shsz; ++ibf ) { + const double z_mu_i = vrho_i * xmat[ipt + ibf*npts]; + const double dbfx_mu_i = basis_x_eval[ipt + ibf*npts]; + const double dbfy_mu_i = basis_y_eval[ipt + ibf*npts]; + const double dbfz_mu_i = basis_z_eval[ipt + ibf*npts]; + + g_acc_x_task += z_mu_i * dbfx_mu_i; + g_acc_y_task += z_mu_i * dbfy_mu_i; + g_acc_z_task += z_mu_i * dbfz_mu_i; + } // Loop over bfns within a shell + + } // Loop over points + + g_acc_x += g_acc_x_task; + g_acc_y += g_acc_y_task; + g_acc_z += g_acc_z_task; + + //write to Parent atom with translational invariance + if constexpr( with_weight_derivatives ) { + atomic_add_device( EXC_GRAD + 3*iParent + 0, 2.0 * g_acc_x_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 1, 2.0 * g_acc_y_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 2, 2.0 * g_acc_z_task ); + } + + } // Loop over tasks assigned to shell + + constexpr auto warp_size = sycl::warp_size; + g_acc_x = -2. * sycl::warp_reduce_sum( g_acc_x ); + g_acc_y = -2. * sycl::warp_reduce_sum( g_acc_y ); + g_acc_z = -2. * sycl::warp_reduce_sum( g_acc_z ); + + if( (it.get_local_id(2) % sycl::warp_size) == 0 ) { + + atomic_add_device( EXC_GRAD + 3*iCen + 0, g_acc_x ); + atomic_add_device( EXC_GRAD + 3*iCen + 1, g_acc_y ); + atomic_add_device( EXC_GRAD + 3*iCen + 2, g_acc_z ); + } + + } // Loop over shells + +} + +template +void increment_exc_grad_lda_uks_kernel( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks, + double* __restrict__ EXC_GRAD +) { + + auto it = GauXC::sycl::this_item(); + + for( uint32_t ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + const int iCen = shell_to_task[ish].center_idx; + const uint32_t shsz = shell->size(); + + const int global_warp_id = + it.get_global_id(2) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + double g_acc_x(0), g_acc_y(0), g_acc_z(0); + for( uint32_t itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + const int iParent = task->iParent; + if constexpr( with_weight_derivatives ) { + if( iCen == iParent ) + continue; + } + double g_acc_x_task(0), g_acc_y_task(0), g_acc_z_task(0); + + const auto* __restrict__ basis_x_eval = task->dbfx + shoff; + const auto* __restrict__ basis_y_eval = task->dbfy + shoff; + const auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + const auto* __restrict__ xmatS = task->xmatS + shoff; + const auto* __restrict__ xmatZ = task->xmatZ + shoff; + const auto* __restrict__ vrhop = task->vrho_pos; + const auto* __restrict__ vrhom = task->vrho_neg; + + for( uint32_t ipt = it.get_local_id(2) % sycl::warp_size; + ipt < npts; + ipt += sycl::warp_size ) { + + const double vrhop_i = vrhop[ipt]; + const double vrhom_i = vrhom[ipt]; + + const auto vrhoS_i = 0.5 * (vrhop_i + vrhom_i); + const auto vrhoZ_i = 0.5 * (vrhop_i - vrhom_i); + for( uint32_t ibf = 0; ibf < shsz; ++ibf ) { + const double zS_mu_i = vrhoS_i * xmatS[ipt + ibf*npts]; + const double zZ_mu_i = vrhoZ_i * xmatZ[ipt + ibf*npts]; + const double dbfx_mu_i = basis_x_eval[ipt + ibf*npts]; + const double dbfy_mu_i = basis_y_eval[ipt + ibf*npts]; + const double dbfz_mu_i = basis_z_eval[ipt + ibf*npts]; + + g_acc_x_task += zS_mu_i * dbfx_mu_i; + g_acc_y_task += zS_mu_i * dbfy_mu_i; + g_acc_z_task += zS_mu_i * dbfz_mu_i; + g_acc_x_task += zZ_mu_i * dbfx_mu_i; + g_acc_y_task += zZ_mu_i * dbfy_mu_i; + g_acc_z_task += zZ_mu_i * dbfz_mu_i; + } // Loop over bfns within a shell + + } // Loop over points + + g_acc_x += g_acc_x_task; + g_acc_y += g_acc_y_task; + g_acc_z += g_acc_z_task; + + //write to Parent atom with translational invariance + if constexpr( with_weight_derivatives ) { + atomic_add_device( EXC_GRAD + 3*iParent + 0, 2.0 * g_acc_x_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 1, 2.0 * g_acc_y_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 2, 2.0 * g_acc_z_task ); + } + + } // Loop over tasks assigned to shell + + constexpr auto warp_size = sycl::warp_size; + g_acc_x = -2. * sycl::warp_reduce_sum( g_acc_x ); + g_acc_y = -2. * sycl::warp_reduce_sum( g_acc_y ); + g_acc_z = -2. * sycl::warp_reduce_sum( g_acc_z ); + + if( (it.get_local_id(2) % sycl::warp_size) == 0 ) { + atomic_add_device( EXC_GRAD + 3*iCen + 0, g_acc_x ); + atomic_add_device( EXC_GRAD + 3*iCen + 1, g_acc_y ); + atomic_add_device( EXC_GRAD + 3*iCen + 2, g_acc_z ); + } + + } // Loop over shells + +} + +void increment_exc_grad_lda( integrator_ks_scheme ks_scheme, size_t nshell, ShellToTaskDevice* shell_to_task, + XCDeviceTask* device_tasks, double* EXC_GRAD, bool with_weight_derivatives, device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + sycl::launch_dim threads(1024), blocks(1,1,nshell); + + switch(ks_scheme) { + case RKS: + if (with_weight_derivatives) { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_lda_rks_kernel", + [=](){ + increment_exc_grad_lda_rks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } else { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_lda_rks_kernel", + [=](){ + increment_exc_grad_lda_rks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } + break; + case UKS: + if (with_weight_derivatives) { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_lda_uks_kernel", + [=](){ + increment_exc_grad_lda_uks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } else { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_lda_uks_kernel", + [=](){ + increment_exc_grad_lda_uks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } + break; + default: GAUXC_GENERIC_EXCEPTION("LDA EXC GRAD + GKS NYI"); + } +} + + + + +template +void increment_exc_grad_gga_rks_kernel( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks, + double* __restrict__ EXC_GRAD +) { + + auto it = GauXC::sycl::this_item(); + + for( uint32_t ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + const int iCen = shell_to_task[ish].center_idx; + const uint32_t shsz = shell->size(); + + const int global_warp_id = + it.get_global_id(2) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + double g_acc_x(0), g_acc_y(0), g_acc_z(0); + for( uint32_t itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + const int iParent = task->iParent; + if constexpr( with_weight_derivatives ) { + if( iCen == iParent ) + continue; + } + double g_acc_x_task(0), g_acc_y_task(0), g_acc_z_task(0); + + const auto* __restrict__ basis_x_eval = task->dbfx + shoff; + const auto* __restrict__ basis_y_eval = task->dbfy + shoff; + const auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + const auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + const auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + const auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + const auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + const auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + const auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + const auto* __restrict__ xmat = task->zmat + shoff; + const auto* __restrict__ xmat_x = task->xmat_x + shoff; + const auto* __restrict__ xmat_y = task->xmat_y + shoff; + const auto* __restrict__ xmat_z = task->xmat_z + shoff; + + const auto* __restrict__ vrho = task->vrho; + const auto* __restrict__ vgamma = task->vgamma; + + const auto* __restrict__ den_x = task->dden_sx; + const auto* __restrict__ den_y = task->dden_sy; + const auto* __restrict__ den_z = task->dden_sz; + + for( uint32_t ipt = it.get_local_id(2) % sycl::warp_size; + ipt < npts; + ipt += sycl::warp_size ) { + + const double vrho_i = vrho[ipt]; + const double vgamma_i = vgamma[ipt]; + + const double denx_i = den_x[ipt]; + const double deny_i = den_y[ipt]; + const double denz_i = den_z[ipt]; + for( uint32_t ibf = 0; ibf < shsz; ++ibf ) { + const double z_mu_i = xmat[ipt + ibf*npts]; + const double dbfx_mu_i = basis_x_eval[ipt + ibf*npts]; + const double dbfy_mu_i = basis_y_eval[ipt + ibf*npts]; + const double dbfz_mu_i = basis_z_eval[ipt + ibf*npts]; + + g_acc_x_task += vrho_i * z_mu_i * dbfx_mu_i; + g_acc_y_task += vrho_i * z_mu_i * dbfy_mu_i; + g_acc_z_task += vrho_i * z_mu_i * dbfz_mu_i; + + const double zx = xmat_x[ipt + ibf*npts]; + const double zy = xmat_y[ipt + ibf*npts]; + const double zz = xmat_z[ipt + ibf*npts]; + + const double d11_xmat_term = denx_i * zx + deny_i * zy + denz_i * zz; + + const double d2bfxx = basis_xx_eval[ipt + ibf*npts]; + const double d2bfxy = basis_xy_eval[ipt + ibf*npts]; + const double d2bfxz = basis_xz_eval[ipt + ibf*npts]; + const double d2bfyy = basis_yy_eval[ipt + ibf*npts]; + const double d2bfyz = basis_yz_eval[ipt + ibf*npts]; + const double d2bfzz = basis_zz_eval[ipt + ibf*npts]; + + const double d2_term_x = d2bfxx*denx_i + d2bfxy*deny_i + d2bfxz*denz_i; + const double d2_term_y = d2bfxy*denx_i + d2bfyy*deny_i + d2bfyz*denz_i; + const double d2_term_z = d2bfxz*denx_i + d2bfyz*deny_i + d2bfzz*denz_i; + + g_acc_x_task += 2 * vgamma_i * ( z_mu_i * d2_term_x + dbfx_mu_i * d11_xmat_term ); + g_acc_y_task += 2 * vgamma_i * ( z_mu_i * d2_term_y + dbfy_mu_i * d11_xmat_term ); + g_acc_z_task += 2 * vgamma_i * ( z_mu_i * d2_term_z + dbfz_mu_i * d11_xmat_term ); + + } // Loop over bfns within a shell + + } // Loop over points + + g_acc_x += g_acc_x_task; + g_acc_y += g_acc_y_task; + g_acc_z += g_acc_z_task; + + //write to Parent atom with translational invariance + if constexpr( with_weight_derivatives ) { + atomic_add_device( EXC_GRAD + 3*iParent + 0, 2.0 * g_acc_x_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 1, 2.0 * g_acc_y_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 2, 2.0 * g_acc_z_task ); + } + + } // Loop over tasks assigned to shell + + constexpr auto warp_size = sycl::warp_size; + g_acc_x = -2. * sycl::warp_reduce_sum( g_acc_x ); + g_acc_y = -2. * sycl::warp_reduce_sum( g_acc_y ); + g_acc_z = -2. * sycl::warp_reduce_sum( g_acc_z ); + + if( (it.get_local_id(2) % sycl::warp_size) == 0 ) { + atomic_add_device( EXC_GRAD + 3*iCen + 0, g_acc_x ); + atomic_add_device( EXC_GRAD + 3*iCen + 1, g_acc_y ); + atomic_add_device( EXC_GRAD + 3*iCen + 2, g_acc_z ); + } + + } // Loop over shells + +} + +template +void increment_exc_grad_gga_uks_kernel( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks, + double* __restrict__ EXC_GRAD +) { + + auto it = GauXC::sycl::this_item(); + + for( uint32_t ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + const int iCen = shell_to_task[ish].center_idx; + const uint32_t shsz = shell->size(); + + const int global_warp_id = + it.get_global_id(2) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + double g_acc_x(0), g_acc_y(0), g_acc_z(0); + for( uint32_t itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + const int iParent = task->iParent; + if constexpr( with_weight_derivatives ) { + if( iCen == iParent ) + continue; + } + double g_acc_x_task(0), g_acc_y_task(0), g_acc_z_task(0); + + const auto* __restrict__ basis_x_eval = task->dbfx + shoff; + const auto* __restrict__ basis_y_eval = task->dbfy + shoff; + const auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + const auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + const auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + const auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + const auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + const auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + const auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + const auto* __restrict__ xmatS = task->xmatS + shoff; + const auto* __restrict__ xmatS_x = task->xmatS_x + shoff; + const auto* __restrict__ xmatS_y = task->xmatS_y + shoff; + const auto* __restrict__ xmatS_z = task->xmatS_z + shoff; + + const auto* __restrict__ xmatZ = task->xmatZ + shoff; + const auto* __restrict__ xmatZ_x = task->xmatZ_x + shoff; + const auto* __restrict__ xmatZ_y = task->xmatZ_y + shoff; + const auto* __restrict__ xmatZ_z = task->xmatZ_z + shoff; + + const auto* __restrict__ vrhop = task->vrho_pos; + const auto* __restrict__ vrhom = task->vrho_neg; + + const auto* __restrict__ vgamma_pp = task->vgamma_pp; + const auto* __restrict__ vgamma_pm = task->vgamma_pm; + const auto* __restrict__ vgamma_mm = task->vgamma_mm; + + const auto* __restrict__ dens_x = task->dden_sx; + const auto* __restrict__ dens_y = task->dden_sy; + const auto* __restrict__ dens_z = task->dden_sz; + + const auto* __restrict__ denz_x = task->dden_zx; + const auto* __restrict__ denz_y = task->dden_zy; + const auto* __restrict__ denz_z = task->dden_zz; + + for( uint32_t ipt = it.get_local_id(2) % sycl::warp_size; + ipt < npts; + ipt += sycl::warp_size ) { + + const double vrhop_i = vrhop[ipt]; + const double vrhom_i = vrhom[ipt]; + const double vrhoS_i = 0.5 * (vrhop_i + vrhom_i); + const double vrhoZ_i = 0.5 * (vrhop_i - vrhom_i); + + const double vgammapp_i = vgamma_pp[ipt]; + const double vgammapm_i = vgamma_pm[ipt]; + const double vgammamm_i = vgamma_mm[ipt]; + + const double denSx_i = dens_x[ipt]; + const double denSy_i = dens_y[ipt]; + const double denSz_i = dens_z[ipt]; + const double denZx_i = denz_x[ipt]; + const double denZy_i = denz_y[ipt]; + const double denZz_i = denz_z[ipt]; + + for( uint32_t ibf = 0; ibf < shsz; ++ibf ) { + const double xN = xmatS[ipt + ibf*npts]; + const double xZ = xmatZ[ipt + ibf*npts]; + const double dbfx_mu_i = basis_x_eval[ipt + ibf*npts]; + const double dbfy_mu_i = basis_y_eval[ipt + ibf*npts]; + const double dbfz_mu_i = basis_z_eval[ipt + ibf*npts]; + + g_acc_x_task += vrhoS_i * xN * dbfx_mu_i; + g_acc_y_task += vrhoS_i * xN * dbfy_mu_i; + g_acc_z_task += vrhoS_i * xN * dbfz_mu_i; + g_acc_x_task += vrhoZ_i * xZ * dbfx_mu_i; + g_acc_y_task += vrhoZ_i * xZ * dbfy_mu_i; + g_acc_z_task += vrhoZ_i * xZ * dbfz_mu_i; + + const double xNx = xmatS_x[ipt + ibf*npts]; + const double xNy = xmatS_y[ipt + ibf*npts]; + const double xNz = xmatS_z[ipt + ibf*npts]; + const double xZx = xmatZ_x[ipt + ibf*npts]; + const double xZy = xmatZ_y[ipt + ibf*npts]; + const double xZz = xmatZ_z[ipt + ibf*npts]; + + const double d11nn_xmat_term = denSx_i * xNx + denSy_i * xNy + denSz_i * xNz; + const double d11nz_xmat_term = denSx_i * xZx + denSy_i * xZy + denSz_i * xZz; + const double d11zn_xmat_term = denZx_i * xNx + denZy_i * xNy + denZz_i * xNz; + const double d11zz_xmat_term = denZx_i * xZx + denZy_i * xZy + denZz_i * xZz; + + const double d2bfxx = basis_xx_eval[ipt + ibf*npts]; + const double d2bfxy = basis_xy_eval[ipt + ibf*npts]; + const double d2bfxz = basis_xz_eval[ipt + ibf*npts]; + const double d2bfyy = basis_yy_eval[ipt + ibf*npts]; + const double d2bfyz = basis_yz_eval[ipt + ibf*npts]; + const double d2bfzz = basis_zz_eval[ipt + ibf*npts]; + + const double d2n_term_x = d2bfxx*denSx_i + d2bfxy*denSy_i + d2bfxz*denSz_i; + const double d2n_term_y = d2bfxy*denSx_i + d2bfyy*denSy_i + d2bfyz*denSz_i; + const double d2n_term_z = d2bfxz*denSx_i + d2bfyz*denSy_i + d2bfzz*denSz_i; + const double d2z_term_x = d2bfxx*denZx_i + d2bfxy*denZy_i + d2bfxz*denZz_i; + const double d2z_term_y = d2bfxy*denZx_i + d2bfyy*denZy_i + d2bfyz*denZz_i; + const double d2z_term_z = d2bfxz*denZx_i + d2bfyz*denZy_i + d2bfzz*denZz_i; + + g_acc_x_task += 0.5 * (vgammapp_i + vgammapm_i + vgammamm_i) * (d2n_term_x * xN + d11nn_xmat_term * dbfx_mu_i); + g_acc_x_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2z_term_x * xN + d11zn_xmat_term * dbfx_mu_i); + g_acc_x_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2n_term_x * xZ + d11nz_xmat_term * dbfx_mu_i); + g_acc_x_task += 0.5 * (vgammapp_i - vgammapm_i + vgammamm_i) * (d2z_term_x * xZ + d11zz_xmat_term * dbfx_mu_i); + + g_acc_y_task += 0.5 * (vgammapp_i + vgammapm_i + vgammamm_i) * (d2n_term_y * xN + d11nn_xmat_term * dbfy_mu_i); + g_acc_y_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2z_term_y * xN + d11zn_xmat_term * dbfy_mu_i); + g_acc_y_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2n_term_y * xZ + d11nz_xmat_term * dbfy_mu_i); + g_acc_y_task += 0.5 * (vgammapp_i - vgammapm_i + vgammamm_i) * (d2z_term_y * xZ + d11zz_xmat_term * dbfy_mu_i); + + g_acc_z_task += 0.5 * (vgammapp_i + vgammapm_i + vgammamm_i) * (d2n_term_z * xN + d11nn_xmat_term * dbfz_mu_i); + g_acc_z_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2z_term_z * xN + d11zn_xmat_term * dbfz_mu_i); + g_acc_z_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2n_term_z * xZ + d11nz_xmat_term * dbfz_mu_i); + g_acc_z_task += 0.5 * (vgammapp_i - vgammapm_i + vgammamm_i) * (d2z_term_z * xZ + d11zz_xmat_term * dbfz_mu_i); + + }// Loop over bfns within a shell + + } // Loop over points + + g_acc_x += g_acc_x_task; + g_acc_y += g_acc_y_task; + g_acc_z += g_acc_z_task; + + //write to Parent atom with translational invariance + if constexpr( with_weight_derivatives ) { + atomic_add_device( EXC_GRAD + 3*iParent + 0, 2.0 * g_acc_x_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 1, 2.0 * g_acc_y_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 2, 2.0 * g_acc_z_task ); + } + + } // Loop over tasks assigned to shell + + constexpr auto warp_size = sycl::warp_size; + g_acc_x = -2. * sycl::warp_reduce_sum( g_acc_x ); + g_acc_y = -2. * sycl::warp_reduce_sum( g_acc_y ); + g_acc_z = -2. * sycl::warp_reduce_sum( g_acc_z ); + + if( (it.get_local_id(2) % sycl::warp_size) == 0 ) { + atomic_add_device( EXC_GRAD + 3*iCen + 0, g_acc_x ); + atomic_add_device( EXC_GRAD + 3*iCen + 1, g_acc_y ); + atomic_add_device( EXC_GRAD + 3*iCen + 2, g_acc_z ); + } + + } // Loop over shells + +} + +void increment_exc_grad_gga( integrator_ks_scheme ks_scheme, size_t nshell, ShellToTaskDevice* shell_to_task, + XCDeviceTask* device_tasks, double* EXC_GRAD, bool with_weight_derivatives, device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + sycl::launch_dim threads(512), blocks(1,1,nshell); + + switch(ks_scheme) { + case RKS: + if (with_weight_derivatives) { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_gga_rks_kernel", + [=](){ + increment_exc_grad_gga_rks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } else { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_gga_rks_kernel", + [=](){ + increment_exc_grad_gga_rks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } + break; + case UKS: + if (with_weight_derivatives) { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_gga_uks_kernel", + [=](){ + increment_exc_grad_gga_uks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } else { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_gga_uks_kernel", + [=](){ + increment_exc_grad_gga_uks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } + break; + default: GAUXC_GENERIC_EXCEPTION("GGA EXC GRAD + GKS NYI"); + } +} + + + + +template +void increment_exc_grad_mgga_rks_kernel( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks, + double* __restrict__ EXC_GRAD +) { + + auto it = GauXC::sycl::this_item(); + + for( uint32_t ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + const int iCen = shell_to_task[ish].center_idx; + const uint32_t shsz = shell->size(); + + const int global_warp_id = + it.get_global_id(2) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + double g_acc_x(0), g_acc_y(0), g_acc_z(0); + for( uint32_t itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + const int iParent = task->iParent; + if constexpr( with_weight_derivatives ) { + if( iCen == iParent ) + continue; + } + double g_acc_x_task(0), g_acc_y_task(0), g_acc_z_task(0); + + const auto* __restrict__ basis_x_eval = task->dbfx + shoff; + const auto* __restrict__ basis_y_eval = task->dbfy + shoff; + const auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + const auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + const auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + const auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + const auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + const auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + const auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + const auto* __restrict__ xmat = task->zmat + shoff; + const auto* __restrict__ xmat_x = task->xmat_x + shoff; + const auto* __restrict__ xmat_y = task->xmat_y + shoff; + const auto* __restrict__ xmat_z = task->xmat_z + shoff; + + const auto* __restrict__ vrho = task->vrho; + const auto* __restrict__ vgamma = task->vgamma; + const auto* __restrict__ vtau = task->vtau; + + const auto* __restrict__ den_x = task->dden_sx; + const auto* __restrict__ den_y = task->dden_sy; + const auto* __restrict__ den_z = task->dden_sz; + + for( uint32_t ipt = it.get_local_id(2) % sycl::warp_size; + ipt < npts; + ipt += sycl::warp_size ) { + + const double vrho_i = vrho[ipt]; + const double vgamma_i = vgamma[ipt]; + const double vtau_i = 0.5 * vtau[ipt]; + + const double denx_i = den_x[ipt]; + const double deny_i = den_y[ipt]; + const double denz_i = den_z[ipt]; + for( uint32_t ibf = 0; ibf < shsz; ++ibf ) { + const double z_mu_i = xmat[ipt + ibf*npts]; + const double dbfx_mu_i = basis_x_eval[ipt + ibf*npts]; + const double dbfy_mu_i = basis_y_eval[ipt + ibf*npts]; + const double dbfz_mu_i = basis_z_eval[ipt + ibf*npts]; + + g_acc_x_task += vrho_i * z_mu_i * dbfx_mu_i; + g_acc_y_task += vrho_i * z_mu_i * dbfy_mu_i; + g_acc_z_task += vrho_i * z_mu_i * dbfz_mu_i; + + const double zx = xmat_x[ipt + ibf*npts]; + const double zy = xmat_y[ipt + ibf*npts]; + const double zz = xmat_z[ipt + ibf*npts]; + + const double d11_xmat_term = denx_i * zx + deny_i * zy + denz_i * zz; + + const double d2bfxx = basis_xx_eval[ipt + ibf*npts]; + const double d2bfxy = basis_xy_eval[ipt + ibf*npts]; + const double d2bfxz = basis_xz_eval[ipt + ibf*npts]; + const double d2bfyy = basis_yy_eval[ipt + ibf*npts]; + const double d2bfyz = basis_yz_eval[ipt + ibf*npts]; + const double d2bfzz = basis_zz_eval[ipt + ibf*npts]; + + { + const double d2_term_x = d2bfxx*denx_i + d2bfxy*deny_i + d2bfxz*denz_i; + const double d2_term_y = d2bfxy*denx_i + d2bfyy*deny_i + d2bfyz*denz_i; + const double d2_term_z = d2bfxz*denx_i + d2bfyz*deny_i + d2bfzz*denz_i; + + g_acc_x_task += 2 * vgamma_i * ( z_mu_i * d2_term_x + dbfx_mu_i * d11_xmat_term ); + g_acc_y_task += 2 * vgamma_i * ( z_mu_i * d2_term_y + dbfy_mu_i * d11_xmat_term ); + g_acc_z_task += 2 * vgamma_i * ( z_mu_i * d2_term_z + dbfz_mu_i * d11_xmat_term ); + } + + { + const double d2_term_x = d2bfxx*zx + d2bfxy*zy + d2bfxz*zz; + const double d2_term_y = d2bfxy*zx + d2bfyy*zy + d2bfyz*zz; + const double d2_term_z = d2bfxz*zx + d2bfyz*zy + d2bfzz*zz; + + g_acc_x_task += vtau_i * d2_term_x; + g_acc_y_task += vtau_i * d2_term_y; + g_acc_z_task += vtau_i * d2_term_z; + } + + } // Loop over bfns within a shell + + } // Loop over points + + g_acc_x += g_acc_x_task; + g_acc_y += g_acc_y_task; + g_acc_z += g_acc_z_task; + + //write to Parent atom with translational invariance + if constexpr( with_weight_derivatives ) { + atomic_add_device( EXC_GRAD + 3*iParent + 0, 2.0 * g_acc_x_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 1, 2.0 * g_acc_y_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 2, 2.0 * g_acc_z_task ); + } + + } // Loop over tasks assigned to shell + + constexpr auto warp_size = sycl::warp_size; + g_acc_x = -2. * sycl::warp_reduce_sum( g_acc_x ); + g_acc_y = -2. * sycl::warp_reduce_sum( g_acc_y ); + g_acc_z = -2. * sycl::warp_reduce_sum( g_acc_z ); + + if( (it.get_local_id(2) % sycl::warp_size) == 0 ) { + atomic_add_device( EXC_GRAD + 3*iCen + 0, g_acc_x ); + atomic_add_device( EXC_GRAD + 3*iCen + 1, g_acc_y ); + atomic_add_device( EXC_GRAD + 3*iCen + 2, g_acc_z ); + } + + } // Loop over shells + +} + +template +void increment_exc_grad_mgga_uks_kernel( + uint32_t nshell, + ShellToTaskDevice* __restrict__ shell_to_task, + XCDeviceTask* __restrict__ device_tasks, + double* __restrict__ EXC_GRAD +) { + + auto it = GauXC::sycl::this_item(); + for( uint32_t ish = it.get_group(0); ish < nshell; ish += it.get_group_range(0) ) { + const uint32_t ntasks = shell_to_task[ish].ntask; + const auto shell = shell_to_task[ish].shell_device; + const auto task_idx = shell_to_task[ish].task_idx_device; + const auto task_shell_offs = shell_to_task[ish].task_shell_offs_device; + const int iCen = shell_to_task[ish].center_idx; + const uint32_t shsz = shell->size(); + + const int global_warp_id = + it.get_global_id(2) / sycl::warp_size; + const int nwarp_global = std::max(it.get_global_range(2) / sycl::warp_size,1); + + double g_acc_x(0), g_acc_y(0), g_acc_z(0); + for( uint32_t itask = global_warp_id; itask < ntasks; itask += nwarp_global ) { + + const auto* task = device_tasks + task_idx[itask]; + const uint32_t npts = task->npts; + const size_t shoff = task_shell_offs[itask] * npts; + const int iParent = task->iParent; + if constexpr( with_weight_derivatives ) { + if( iCen == iParent ) + continue; + } + double g_acc_x_task(0), g_acc_y_task(0), g_acc_z_task(0); + + const auto* __restrict__ basis_x_eval = task->dbfx + shoff; + const auto* __restrict__ basis_y_eval = task->dbfy + shoff; + const auto* __restrict__ basis_z_eval = task->dbfz + shoff; + + const auto* __restrict__ basis_xx_eval = task->d2bfxx + shoff; + const auto* __restrict__ basis_xy_eval = task->d2bfxy + shoff; + const auto* __restrict__ basis_xz_eval = task->d2bfxz + shoff; + const auto* __restrict__ basis_yy_eval = task->d2bfyy + shoff; + const auto* __restrict__ basis_yz_eval = task->d2bfyz + shoff; + const auto* __restrict__ basis_zz_eval = task->d2bfzz + shoff; + + const auto* __restrict__ xmatS = task->xmatS + shoff; + const auto* __restrict__ xmatS_x = task->xmatS_x + shoff; + const auto* __restrict__ xmatS_y = task->xmatS_y + shoff; + const auto* __restrict__ xmatS_z = task->xmatS_z + shoff; + + const auto* __restrict__ xmatZ = task->xmatZ + shoff; + const auto* __restrict__ xmatZ_x = task->xmatZ_x + shoff; + const auto* __restrict__ xmatZ_y = task->xmatZ_y + shoff; + const auto* __restrict__ xmatZ_z = task->xmatZ_z + shoff; + + const auto* __restrict__ vrhop = task->vrho_pos; + const auto* __restrict__ vrhom = task->vrho_neg; + const auto* __restrict__ vtaup = task->vtau_pos; + const auto* __restrict__ vtaum = task->vtau_neg; + + const auto* __restrict__ vgamma_pp = task->vgamma_pp; + const auto* __restrict__ vgamma_pm = task->vgamma_pm; + const auto* __restrict__ vgamma_mm = task->vgamma_mm; + + const auto* __restrict__ dens_x = task->dden_sx; + const auto* __restrict__ dens_y = task->dden_sy; + const auto* __restrict__ dens_z = task->dden_sz; + + const auto* __restrict__ denz_x = task->dden_zx; + const auto* __restrict__ denz_y = task->dden_zy; + const auto* __restrict__ denz_z = task->dden_zz; + + for( uint32_t ipt = it.get_local_id(2) % sycl::warp_size; + ipt < npts; + ipt += sycl::warp_size ) { + + const double vrhop_i = vrhop[ipt]; + const double vrhom_i = vrhom[ipt]; + const double vrhoS_i = 0.5 * (vrhop_i + vrhom_i); + const double vrhoZ_i = 0.5 * (vrhop_i - vrhom_i); + + const double vtaup_i = 0.5 * vtaup[ipt]; + const double vtaum_i = 0.5 * vtaum[ipt]; + const double vtauS_i = 0.5 * (vtaup_i + vtaum_i); + const double vtauZ_i = 0.5 * (vtaup_i - vtaum_i); + + const double vgammapp_i = vgamma_pp[ipt]; + const double vgammapm_i = vgamma_pm[ipt]; + const double vgammamm_i = vgamma_mm[ipt]; + + const double denSx_i = dens_x[ipt]; + const double denSy_i = dens_y[ipt]; + const double denSz_i = dens_z[ipt]; + const double denZx_i = denz_x[ipt]; + const double denZy_i = denz_y[ipt]; + const double denZz_i = denz_z[ipt]; + + for( uint32_t ibf = 0; ibf < shsz; ++ibf ) { + const double xN = xmatS[ipt + ibf*npts]; + const double xZ = xmatZ[ipt + ibf*npts]; + const double dbfx_mu_i = basis_x_eval[ipt + ibf*npts]; + const double dbfy_mu_i = basis_y_eval[ipt + ibf*npts]; + const double dbfz_mu_i = basis_z_eval[ipt + ibf*npts]; + + g_acc_x_task += vrhoS_i * xN * dbfx_mu_i; + g_acc_y_task += vrhoS_i * xN * dbfy_mu_i; + g_acc_z_task += vrhoS_i * xN * dbfz_mu_i; + g_acc_x_task += vrhoZ_i * xZ * dbfx_mu_i; + g_acc_y_task += vrhoZ_i * xZ * dbfy_mu_i; + g_acc_z_task += vrhoZ_i * xZ * dbfz_mu_i; + + const double xNx = xmatS_x[ipt + ibf*npts]; + const double xNy = xmatS_y[ipt + ibf*npts]; + const double xNz = xmatS_z[ipt + ibf*npts]; + const double xZx = xmatZ_x[ipt + ibf*npts]; + const double xZy = xmatZ_y[ipt + ibf*npts]; + const double xZz = xmatZ_z[ipt + ibf*npts]; + + const double d11nn_xmat_term = denSx_i * xNx + denSy_i * xNy + denSz_i * xNz; + const double d11nz_xmat_term = denSx_i * xZx + denSy_i * xZy + denSz_i * xZz; + const double d11zn_xmat_term = denZx_i * xNx + denZy_i * xNy + denZz_i * xNz; + const double d11zz_xmat_term = denZx_i * xZx + denZy_i * xZy + denZz_i * xZz; + + const double d2bfxx = basis_xx_eval[ipt + ibf*npts]; + const double d2bfxy = basis_xy_eval[ipt + ibf*npts]; + const double d2bfxz = basis_xz_eval[ipt + ibf*npts]; + const double d2bfyy = basis_yy_eval[ipt + ibf*npts]; + const double d2bfyz = basis_yz_eval[ipt + ibf*npts]; + const double d2bfzz = basis_zz_eval[ipt + ibf*npts]; + + { + const double d2n_term_x = d2bfxx*denSx_i + d2bfxy*denSy_i + d2bfxz*denSz_i; + const double d2n_term_y = d2bfxy*denSx_i + d2bfyy*denSy_i + d2bfyz*denSz_i; + const double d2n_term_z = d2bfxz*denSx_i + d2bfyz*denSy_i + d2bfzz*denSz_i; + const double d2z_term_x = d2bfxx*denZx_i + d2bfxy*denZy_i + d2bfxz*denZz_i; + const double d2z_term_y = d2bfxy*denZx_i + d2bfyy*denZy_i + d2bfyz*denZz_i; + const double d2z_term_z = d2bfxz*denZx_i + d2bfyz*denZy_i + d2bfzz*denZz_i; + + g_acc_x_task += 0.5 * (vgammapp_i + vgammapm_i + vgammamm_i) * (d2n_term_x * xN + d11nn_xmat_term * dbfx_mu_i); + g_acc_x_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2z_term_x * xN + d11zn_xmat_term * dbfx_mu_i); + g_acc_x_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2n_term_x * xZ + d11nz_xmat_term * dbfx_mu_i); + g_acc_x_task += 0.5 * (vgammapp_i - vgammapm_i + vgammamm_i) * (d2z_term_x * xZ + d11zz_xmat_term * dbfx_mu_i); + + g_acc_y_task += 0.5 * (vgammapp_i + vgammapm_i + vgammamm_i) * (d2n_term_y * xN + d11nn_xmat_term * dbfy_mu_i); + g_acc_y_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2z_term_y * xN + d11zn_xmat_term * dbfy_mu_i); + g_acc_y_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2n_term_y * xZ + d11nz_xmat_term * dbfy_mu_i); + g_acc_y_task += 0.5 * (vgammapp_i - vgammapm_i + vgammamm_i) * (d2z_term_y * xZ + d11zz_xmat_term * dbfy_mu_i); + + g_acc_z_task += 0.5 * (vgammapp_i + vgammapm_i + vgammamm_i) * (d2n_term_z * xN + d11nn_xmat_term * dbfz_mu_i); + g_acc_z_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2z_term_z * xN + d11zn_xmat_term * dbfz_mu_i); + g_acc_z_task += 0.5 * (vgammapp_i - vgammamm_i) * (d2n_term_z * xZ + d11nz_xmat_term * dbfz_mu_i); + g_acc_z_task += 0.5 * (vgammapp_i - vgammapm_i + vgammamm_i) * (d2z_term_z * xZ + d11zz_xmat_term * dbfz_mu_i); + } + + { + const double d2n_term_x = d2bfxx*xNx + d2bfxy*xNy + d2bfxz*xNz; + const double d2n_term_y = d2bfxy*xNx + d2bfyy*xNy + d2bfyz*xNz; + const double d2n_term_z = d2bfxz*xNx + d2bfyz*xNy + d2bfzz*xNz; + const double d2z_term_x = d2bfxx*xZx + d2bfxy*xZy + d2bfxz*xZz; + const double d2z_term_y = d2bfxy*xZx + d2bfyy*xZy + d2bfyz*xZz; + const double d2z_term_z = d2bfxz*xZx + d2bfyz*xZy + d2bfzz*xZz; + g_acc_x_task += vtauS_i * d2n_term_x; + g_acc_y_task += vtauS_i * d2n_term_y; + g_acc_z_task += vtauS_i * d2n_term_z; + + g_acc_x_task += vtauZ_i * d2z_term_x; + g_acc_y_task += vtauZ_i * d2z_term_y; + g_acc_z_task += vtauZ_i * d2z_term_z; + } + }// Loop over bfns within a shell + + } // Loop over points + + g_acc_x += g_acc_x_task; + g_acc_y += g_acc_y_task; + g_acc_z += g_acc_z_task; + + //write to Parent atom with translational invariance + if constexpr( with_weight_derivatives ) { + atomic_add_device( EXC_GRAD + 3*iParent + 0, 2.0 * g_acc_x_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 1, 2.0 * g_acc_y_task ); + atomic_add_device( EXC_GRAD + 3*iParent + 2, 2.0 * g_acc_z_task ); + } + + } // Loop over tasks assigned to shell + + constexpr auto warp_size = sycl::warp_size; + g_acc_x = -2. * sycl::warp_reduce_sum( g_acc_x ); + g_acc_y = -2. * sycl::warp_reduce_sum( g_acc_y ); + g_acc_z = -2. * sycl::warp_reduce_sum( g_acc_z ); + + if( (it.get_local_id(2) % sycl::warp_size) == 0 ) { + atomic_add_device( EXC_GRAD + 3*iCen + 0, g_acc_x ); + atomic_add_device( EXC_GRAD + 3*iCen + 1, g_acc_y ); + atomic_add_device( EXC_GRAD + 3*iCen + 2, g_acc_z ); + } + + } // Loop over shells + +} + +void increment_exc_grad_mgga( integrator_ks_scheme ks_scheme, size_t nshell, bool need_lapl, + ShellToTaskDevice* shell_to_task, XCDeviceTask* device_tasks, + double* EXC_GRAD, bool with_weight_derivatives, device_queue queue ) { + + if(need_lapl) GAUXC_GENERIC_EXCEPTION("SYCL + MGGA/LAPL EXC GRAD NYI"); + + ::sycl::queue& stream = queue.queue_as(); + sycl::launch_dim threads(512), blocks(1,1,nshell); + + switch(ks_scheme) { + case RKS: + if (with_weight_derivatives) { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_mgga_rks_kernel", + [=](){ + increment_exc_grad_mgga_rks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } else { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_mgga_rks_kernel", + [=](){ + increment_exc_grad_mgga_rks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } + break; + case UKS: + if (with_weight_derivatives) { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_mgga_uks_kernel", + [=](){ + increment_exc_grad_mgga_uks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } else { + sycl::launch_kernel( stream, blocks, threads, "increment_exc_grad_mgga_uks_kernel", + [=](){ + increment_exc_grad_mgga_uks_kernel( + nshell, shell_to_task, device_tasks, EXC_GRAD ); + }); + } + break; + default: GAUXC_GENERIC_EXCEPTION("GGA EXC GRAD + GKS NYI"); + } +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/onemkl_extensions.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/onemkl_extensions.cxx new file mode 100644 index 000000000..ed8a20fce --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/onemkl_extensions.cxx @@ -0,0 +1,203 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include +#include "exceptions/onemkl_exception.hpp" + +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/onemkl_util.hpp" +#include "device/common/device_blas.hpp" + +// Only the BLAS domain is needed. The umbrella also pulls in +// the C LAPACK prototypes, whose parameter named "csl" collides with the csl +// macro from buffer_adaptor.hpp +#include + +namespace GauXC { + +oneapi::mkl::transpose device_op_to_onemkl( DeviceBlasOp op ) { + switch( op ) { + case DeviceBlasOp::NoTrans: return oneapi::mkl::transpose::nontrans; + case DeviceBlasOp::Trans: return oneapi::mkl::transpose::trans; + default: + GAUXC_GENERIC_EXCEPTION("Unsupported DeviceBlasOp"); + return oneapi::mkl::transpose::nontrans; + } +} + +oneapi::mkl::uplo device_uplo_to_onemkl( DeviceBlasUplo uplo ) { + switch(uplo) { + case DeviceBlasUplo::Upper: return oneapi::mkl::uplo::upper; + case DeviceBlasUplo::Lower: return oneapi::mkl::uplo::lower; + default: + GAUXC_GENERIC_EXCEPTION("Unsupported DeviceBlasUplo"); + return oneapi::mkl::uplo::lower; + } +} + +template +void increment( const T* X, T* Y, ::sycl::queue& stream ) { + GAUXC_SYCL_ERROR( "Increment Launch Failed", + stream.single_task( [=](){ (*Y) += (*X); } ) + ); +} + +template +void increment( device_blas_handle generic_handle, const T* X, T* Y, int N) { + const int threads = sycl::warp_size * sycl::max_warps_per_thread_block; + const int blocks = util::div_ceil( N, threads ); + auto& handle = generic_handle.blas_handle_as(); + auto& stream = util::get_queue(handle); + + GAUXC_SYCL_ERROR( "Increment Vec Launch Failed", + stream.parallel_for( ::sycl::nd_range<1>( blocks*threads, threads ), + [=](::sycl::nd_item<1> it) { + const auto tid = it.get_global_id(0); + if( tid < (size_t)N ) Y[tid] += X[tid]; + }) + ); +} + +template + void increment( device_blas_handle generic_handle, const double* X, double* Y, int N ); + +template <> +void dot( device_blas_handle generic_handle, + int N, + const double* X, + int INCX, + const double* Y, + int INCY, + double* RES ) { + + auto& handle = generic_handle.blas_handle_as(); + auto& stream = util::get_queue(handle); + + GAUXC_ONEMKL_ERROR("ONEMKL DDOT FAILED", + oneapi::mkl::blas::column_major::dot( stream, N, X, INCX, Y, INCY, RES ) ); + +} + +template +void gdot( device_blas_handle generic_handle, + int N, + const T* X, + int INCX, + const T* Y, + int INCY, + T* SCR, + T* RES ) { + + + dot( generic_handle, N, X, INCX, Y, INCY, SCR ); + auto& handle = generic_handle.blas_handle_as(); + auto& stream = util::get_queue(handle); + increment( SCR, RES, stream ); + +} + +template +void gdot( device_blas_handle generic_handle, + int N, + const double* X, + int INCX, + const double* Y, + int INCY, + double* SCR, + double* RES ); + + + + + + + + + + +template +void hadamard_product( device_blas_handle generic_handle, + int M, + int N, + const T* A, + int LDA, + T* B, + int LDB ) { + + + auto& handle = generic_handle.blas_handle_as(); + auto& stream = util::get_queue(handle); + + ::sycl::range<2> local( sycl::max_warps_per_thread_block, sycl::warp_size ); + ::sycl::range<2> global( util::div_ceil( N, local[0] ) * local[0], + util::div_ceil( M, local[1] ) * local[1] ); + + GAUXC_SYCL_ERROR( "Hadamard Product Launch Failed", + stream.parallel_for( ::sycl::nd_range<2>(global, local), + [=](::sycl::nd_item<2> it) { + const size_t tid_x = it.get_global_id(1); + const size_t tid_y = it.get_global_id(0); + if( tid_x < (size_t)M and tid_y < (size_t)N ) { + B[ tid_x + tid_y*LDB ] *= A[ tid_x + tid_y*LDA ]; + } + }) + ); + +} + +template +void hadamard_product( device_blas_handle generic_handle, + int M, + int N, + const double* A, + int LDA, + double* B, + int LDB ); + + + + +template <> +void gemm( device_blas_handle generic_handle, + DeviceBlasOp TA, DeviceBlasOp TB, + int M, int N, int K, double ALPHA, + const double* A, int LDA, const double* B, int LDB, + double BETA, double* C, int LDC ) { + + + auto& handle = generic_handle.blas_handle_as(); + auto& stream = util::get_queue(handle); + + GAUXC_ONEMKL_ERROR("ONEMKL DGEMM FAILED", + oneapi::mkl::blas::column_major::gemm( stream, device_op_to_onemkl(TA), + device_op_to_onemkl(TB), M, N, K, ALPHA, A, LDA, B, LDB, BETA, C, LDC ) ); + +} + + +template <> +void syr2k( device_blas_handle generic_handle, + DeviceBlasUplo UPLO, DeviceBlasOp Trans, + int M, int K, double ALPHA, + const double* A, int LDA, const double* B, int LDB, + double BETA, double* C, int LDC ) { + + auto& handle = generic_handle.blas_handle_as(); + auto& stream = util::get_queue(handle); + + GAUXC_ONEMKL_ERROR("ONEMKL DSYR2K FAILED", + oneapi::mkl::blas::column_major::syr2k( stream, device_uplo_to_onemkl(UPLO), + device_op_to_onemkl(Trans), M, K, ALPHA, A, LDA, B, LDB, BETA, C, LDC ) ); + +} + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/pack_submat.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/pack_submat.cxx new file mode 100644 index 000000000..61b5d253c --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/pack_submat.cxx @@ -0,0 +1,233 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/pack_submat.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include +#include "device_specific/sycl_util.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "sycl_launch.hpp" + +namespace GauXC { + +#define WARP_X 16 +#define WARP_Y 1 +#define UNROLL_FACTOR 4 +#define EFF_UNROLL 4 +#define CUT_X 8 +#define CUT_Y 8 + +template +void sym_submat_set_combined_kernel( size_t ntasks, + XCDeviceTask* device_tasks, + T* A, + size_t LDA, + const int block_y, + const int block_x) { + + auto it = GauXC::sycl::this_item(); + const int batch_id = it.get_group(0); + auto& task = device_tasks[ batch_id ]; + + if constexpr (skip_single_cut ) { + if( task.bfn_screening.ncut == 1 ) return; + } + + const auto* submat_cut_device = task.bfn_screening.submat_cut; + const auto* submat_block_device = task.bfn_screening.submat_block; + const auto LDAS = task.bfn_screening.nbe; + auto* ASmall_device = task.nbe_scr; + + //if( LDAS == LDAB ) return; + + const int tid_xx = it.get_local_id(2) % WARP_X; + const int tid_xy = it.get_local_id(2) / WARP_X; + + const int tid_yx = it.get_local_id(1) % CUT_X; + const int tid_yy = it.get_local_id(1) / CUT_X; + + const int start_cut_y = submat_block_device[block_y]; + const int end_cut_y = submat_block_device[block_y+1]; + const int start_cut_x = submat_block_device[block_x]; + const int end_cut_x = submat_block_device[block_x+1]; + + for( int i_cut = tid_yy + start_cut_y; i_cut < end_cut_y; i_cut += CUT_Y ) { + const int3 i_data = *((int3*)(submat_cut_device + 3*i_cut)); + const int i_cut_first = i_data.x; + const int delta_i = i_data.y; + const int i_cut_small = i_data.z; + + for( int j_cut = tid_yx + start_cut_x; j_cut < end_cut_x; j_cut += CUT_X ) { + const int3 j_data = *((int3*)(submat_cut_device + 3*j_cut)); + const int j_cut_first = j_data.x; + const int delta_j = j_data.y; + const int j_cut_small = j_data.z; + + auto* ASmall_begin = ASmall_device + i_cut_small + j_cut_small*LDAS; + auto* ABig_begin = A + i_cut_first + j_cut_first*LDA; + + int J; + for( J = tid_xy; J < (delta_j / EFF_UNROLL) * EFF_UNROLL; J += EFF_UNROLL ) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + + double val[UNROLL_FACTOR]; + double* address[UNROLL_FACTOR]; +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + val[k] = ABig_begin[I + (J + k*WARP_Y)*LDA]; + address[k] = ASmall_begin + I + (J + k*WARP_Y) * LDAS; + } +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + *address[k] = val[k]; + } + } + } + + for ( ; J < delta_j; J += WARP_Y) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + ASmall_begin[I + J*LDAS] = ABig_begin[I + J*LDA]; + } + } + } + } +} + + + + +void sym_pack_submat( size_t ntasks, XCDeviceTask* device_tasks, const double* A, + int32_t LDA, int32_t submat_block_size, device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + sycl::launch_dim threads( sycl::warp_size/2, sycl::max_warps_per_thread_block * 2, 1 ); + sycl::launch_dim blocks( 1,1, ntasks ); + + auto n_launch = util::div_ceil( LDA, submat_block_size ); + for (int i = 0; i < n_launch; i++) + for (int j = 0; j < n_launch; j++) { + sycl::launch_kernel( stream, blocks, threads, "sym_submat_set_combined_kernel", + [=](){ + sym_submat_set_combined_kernel( ntasks, device_tasks, A, LDA, i, j ); + }); + } +} + + + + + + + + +template +void asym_submat_set_combined_kernel( size_t ntasks, + XCDeviceTask* device_tasks, + T* A, + size_t LDA, + const int block_y, + const int block_x) { + + auto it = GauXC::sycl::this_item(); + const int batch_id = it.get_group(0); + auto& task = device_tasks[ batch_id ]; + + if constexpr (skip_single_cut ) { + if( task.bfn_screening.ncut == 1 ) return; + } + + const auto* row_submat_cut_device = task.bfn_screening.submat_cut; + const auto* row_submat_block_device = task.bfn_screening.submat_block; + const auto* col_submat_cut_device = task.cou_screening.submat_cut; + const auto* col_submat_block_device = task.cou_screening.submat_block; + + const auto LDAS = task.bfn_screening.nbe; + auto* ASmall_device = task.nbe_scr; + + //if( LDAS == LDAB ) return; + + const int tid_xx = it.get_local_id(2) % WARP_X; + const int tid_xy = it.get_local_id(2) / WARP_X; + + const int tid_yx = it.get_local_id(1) % CUT_X; + const int tid_yy = it.get_local_id(1) / CUT_X; + + const int start_cut_y = row_submat_block_device[block_y]; + const int end_cut_y = row_submat_block_device[block_y+1]; + const int start_cut_x = col_submat_block_device[block_x]; + const int end_cut_x = col_submat_block_device[block_x+1]; + + for( int i_cut = tid_yy + start_cut_y; i_cut < end_cut_y; i_cut += CUT_Y ) { + const int3 i_data = *((int3*)(row_submat_cut_device + 3*i_cut)); + const int i_cut_first = i_data.x; + const int delta_i = i_data.y; + const int i_cut_small = i_data.z; + + for( int j_cut = tid_yx + start_cut_x; j_cut < end_cut_x; j_cut += CUT_X ) { + const int3 j_data = *((int3*)(col_submat_cut_device + 3*j_cut)); + const int j_cut_first = j_data.x; + const int delta_j = j_data.y; + const int j_cut_small = j_data.z; + + auto* ASmall_begin = ASmall_device + i_cut_small + j_cut_small*LDAS; + auto* ABig_begin = A + i_cut_first + j_cut_first*LDA; + + int J; + for( J = tid_xy; J < (delta_j / EFF_UNROLL) * EFF_UNROLL; J += EFF_UNROLL ) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + + double val[UNROLL_FACTOR]; + double* address[UNROLL_FACTOR]; +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + val[k] = ABig_begin[I + (J + k*WARP_Y)*LDA]; + address[k] = ASmall_begin + I + (J + k*WARP_Y) * LDAS; + } +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + *address[k] = val[k]; + } + } + } + + for ( ; J < delta_j; J += WARP_Y) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + ASmall_begin[I + J*LDAS] = ABig_begin[I + J*LDA]; + } + } + } + } +} + + + + +void asym_pack_submat( size_t ntasks, XCDeviceTask* device_tasks, const double* A, + int32_t LDA, int32_t submat_block_size, device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + sycl::launch_dim threads( sycl::warp_size/2, sycl::max_warps_per_thread_block * 2, 1 ); + sycl::launch_dim blocks( 1,1, ntasks ); + + auto n_launch = util::div_ceil( LDA, submat_block_size ); + for (int i = 0; i < n_launch; i++) + for (int j = 0; j < n_launch; j++) { + sycl::launch_kernel( stream, blocks, threads, "asym_submat_set_combined_kernel", + [=](){ + asym_submat_set_combined_kernel( ntasks, device_tasks, A, LDA, i, j ); + }); + } +} + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_atomics.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_atomics.hpp new file mode 100644 index 000000000..6ad59dbb8 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_atomics.hpp @@ -0,0 +1,44 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include + +namespace GauXC { + +/// Device-scope relaxed atomic add on global memory. +template +inline void atomic_add_device( T* addr, T val ) { + ::sycl::atomic_ref< T, ::sycl::memory_order::relaxed, + ::sycl::memory_scope::device, + ::sycl::access::address_space::global_space > ref( *addr ); + ref.fetch_add( val ); +} + +/// Work-group-scope relaxed atomic add on work-group local memory +template +inline void atomic_add_local( T* addr, T val ) { + ::sycl::atomic_ref< T, ::sycl::memory_order::relaxed, + ::sycl::memory_scope::work_group, + ::sycl::access::address_space::local_space > ref( *addr ); + ref.fetch_add( val ); +} + +/// Work-group-scope relaxed atomic post-increment on work-group local memory +template +inline T atomic_fetch_add_local( T* addr, T val ) { + ::sycl::atomic_ref< T, ::sycl::memory_order::relaxed, + ::sycl::memory_scope::work_group, + ::sycl::access::address_space::local_space > ref( *addr ); + return ref.fetch_add( val ); +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_extensions.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_extensions.hpp new file mode 100644 index 000000000..0272f92bd --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_extensions.hpp @@ -0,0 +1,50 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include "sycl_launch.hpp" + +namespace GauXC { +namespace sycl { + +template +inline T warp_reduce_sum( T val ) { + + auto sg = this_subgroup(); + for(int i=(warp_sz/2); i>=1; i/=2) + val += ::sycl::permute_group_by_xor( sg, val, i ); + + return val; +} + +template +inline T warp_reduce_prod( T val ) { + + auto sg = this_subgroup(); + for(int i=(warp_sz/2); i>=1; i/=2) + val *= ::sycl::permute_group_by_xor( sg, val, i ); + + return val; +} + +template +inline T warp_reduce_max( T val ) { + + auto sg = this_subgroup(); + for(int i=(warp_sz/2); i>=1; i/=2) + val = ::sycl::fmax( val, ::sycl::permute_group_by_xor( sg, val, i ) ); + + return val; +} + +} +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_inc_potential.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_inc_potential.cxx new file mode 100644 index 000000000..45299516f --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_inc_potential.cxx @@ -0,0 +1,232 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/inc_potential.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include +#include "device_specific/sycl_util.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "sycl_launch.hpp" +#include "sycl_atomics.hpp" + + +namespace GauXC { + + +#define WARP_X 16 +#define WARP_Y 1 +#define UNROLL_FACTOR 4 +#define EFF_UNROLL 4 +#define CUT_X 8 +#define CUT_Y 8 + + +void sym_inc_by_submat_combined_kernel( size_t ntasks, + XCDeviceTask* device_tasks, + double* A, + size_t LDA, + const int block_y, + const int block_x ) { + + auto it = GauXC::sycl::this_item(); + const int batch_id = it.get_group(0); + auto& task = device_tasks[ batch_id ]; + + const auto* submat_cut_device = task.bfn_screening.submat_cut; + const auto* submat_block_device = task.bfn_screening.submat_block; + const auto LDAS = task.bfn_screening.nbe; + auto* ASmall_device = task.nbe_scr; + + //if( LDAS == LDAB ) return; + const int tid_xx = it.get_local_id(2) % WARP_X; + const int tid_xy = it.get_local_id(2) / WARP_X; + + const int tid_yx = it.get_local_id(1) % CUT_X; + const int tid_yy = it.get_local_id(1) / CUT_X; + + const int start_cut_y = submat_block_device[block_y]; + const int end_cut_y = submat_block_device[block_y+1]; + const int start_cut_x = submat_block_device[block_x]; + const int end_cut_x = submat_block_device[block_x+1]; + + for( int i_cut = tid_yy + start_cut_y; i_cut < end_cut_y; i_cut += CUT_Y ) { + const int3 i_data = *((int3*)(submat_cut_device + 3*i_cut)); + const int i_cut_first = i_data.x; + const int delta_i = i_data.y; + const int i_cut_small = i_data.z; + + for( int j_cut = tid_yx + start_cut_x; j_cut < end_cut_x; j_cut += CUT_X ) { + const int3 j_data = *((int3*)(submat_cut_device + 3*j_cut)); + const int j_cut_first = j_data.x; + const int delta_j = j_data.y; + const int j_cut_small = j_data.z; + + auto* ASmall_begin = ASmall_device + i_cut_small + j_cut_small*LDAS; + auto* ABig_begin = A + i_cut_first + j_cut_first*LDA; + + int J; + for( J = tid_xy; J < (delta_j / EFF_UNROLL) * EFF_UNROLL; J += EFF_UNROLL ) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + + double val[UNROLL_FACTOR]; + double* address[UNROLL_FACTOR]; +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + val[k] = ASmall_begin[I + (J+k*WARP_Y)*LDAS]; + address[k] = ABig_begin + I + (J+k*WARP_Y)*LDA; + } +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + atomic_add_device(address[k], val[k] ); + } + } + } + + for ( ; J < delta_j; J += WARP_Y) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + atomic_add_device(ABig_begin + I + J*LDA, ASmall_begin[I + J*LDAS] ); + } + } + + } + } +} + + +void sym_task_inc_potential( size_t ntasks, + XCDeviceTask* device_tasks, + double* V_device, + size_t LDV, + size_t submat_block_size, + device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + sycl::launch_dim threads( sycl::warp_size/2, sycl::max_warps_per_thread_block * 2, 1 ); + sycl::launch_dim blocks( 1,1, ntasks ); + + auto n_launch = util::div_ceil( LDV, submat_block_size ); + for (size_t i = 0; i < n_launch; i++) + for (size_t j = 0; j < n_launch; j++) { + sycl::launch_kernel( stream, blocks, threads, "sym_inc_by_submat_combined_kernel", + [=](){ + sym_inc_by_submat_combined_kernel( ntasks, device_tasks, V_device, LDV, + (int)i, (int)j ); + }); + } + +} + + + + + +void asym_inc_by_submat_combined_kernel( size_t ntasks, + XCDeviceTask* device_tasks, + double* A, + size_t LDA, + const int block_y, + const int block_x ) { + + auto it = GauXC::sycl::this_item(); + const int batch_id = it.get_group(0); + auto& task = device_tasks[ batch_id ]; + + const auto* row_submat_cut_device = task.bfn_screening.submat_cut; + const auto* row_submat_block_device = task.bfn_screening.submat_block; + const auto* col_submat_cut_device = task.cou_screening.submat_cut; + const auto* col_submat_block_device = task.cou_screening.submat_block; + + const auto LDAS = task.bfn_screening.nbe; + auto* ASmall_device = task.nbe_scr; + + //if( LDAS == LDAB ) return; + const int tid_xx = it.get_local_id(2) % WARP_X; + const int tid_xy = it.get_local_id(2) / WARP_X; + + const int tid_yx = it.get_local_id(1) % CUT_X; + const int tid_yy = it.get_local_id(1) / CUT_X; + + const int start_cut_y = row_submat_block_device[block_y]; + const int end_cut_y = row_submat_block_device[block_y+1]; + const int start_cut_x = col_submat_block_device[block_x]; + const int end_cut_x = col_submat_block_device[block_x+1]; + + for( int i_cut = tid_yy + start_cut_y; i_cut < end_cut_y; i_cut += CUT_Y ) { + const int3 i_data = *((int3*)(row_submat_cut_device + 3*i_cut)); + const int i_cut_first = i_data.x; + const int delta_i = i_data.y; + const int i_cut_small = i_data.z; + + for( int j_cut = tid_yx + start_cut_x; j_cut < end_cut_x; j_cut += CUT_X ) { + const int3 j_data = *((int3*)(col_submat_cut_device + 3*j_cut)); + const int j_cut_first = j_data.x; + const int delta_j = j_data.y; + const int j_cut_small = j_data.z; + + auto* ASmall_begin = ASmall_device + i_cut_small + j_cut_small*LDAS; + auto* ABig_begin = A + i_cut_first + j_cut_first*LDA; + + int J; + for( J = tid_xy; J < (delta_j / EFF_UNROLL) * EFF_UNROLL; J += EFF_UNROLL ) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + + double val[UNROLL_FACTOR]; + double* address[UNROLL_FACTOR]; +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + val[k] = ASmall_begin[I + (J+k*WARP_Y)*LDAS]; + address[k] = ABig_begin + I + (J+k*WARP_Y)*LDA; + } +#pragma unroll + for (int k = 0; k < UNROLL_FACTOR; k++) { + atomic_add_device(address[k], val[k] ); + } + } + } + + for ( ; J < delta_j; J += WARP_Y) { + for( int I = tid_xx; I < delta_i; I += WARP_X ) { + atomic_add_device(ABig_begin + I + J*LDA, ASmall_begin[I + J*LDAS] ); + } + } + + } + } +} + + +void asym_task_inc_potential( size_t ntasks, + XCDeviceTask* device_tasks, + double* V_device, + size_t LDV, + size_t submat_block_size, + device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + sycl::launch_dim threads( sycl::warp_size/2, sycl::max_warps_per_thread_block * 2, 1 ); + sycl::launch_dim blocks( 1,1, ntasks ); + + auto n_launch = util::div_ceil( LDV, submat_block_size ); + for (size_t i = 0; i < n_launch; i++) + for (size_t j = 0; j < n_launch; j++) { + sycl::launch_kernel( stream, blocks, threads, "asym_inc_by_submat_combined_kernel", + [=](){ + asym_inc_by_submat_combined_kernel( ntasks, device_tasks, V_device, LDV, + (int)i, (int)j ); + }); + } + +} + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_launch.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_launch.hpp new file mode 100644 index 000000000..fcd2617f5 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_launch.hpp @@ -0,0 +1,72 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include "device_specific/sycl_device_constants.hpp" +#include "exceptions/sycl_exception.hpp" + +namespace GauXC { +namespace sycl { + +struct launch_dim { + size_t x, y, z; + launch_dim( size_t x_ = 1, size_t y_ = 1, size_t z_ = 1 ) : + x(x_), y(y_), z(z_) { } +}; + +inline ::sycl::range<3> local_range( launch_dim threads ) { + return ::sycl::range<3>( threads.z, threads.y, threads.x ); +} + +inline ::sycl::range<3> global_range( launch_dim blocks, launch_dim threads ) { + return ::sycl::range<3>( blocks.z * threads.z, blocks.y * threads.y, + blocks.x * threads.x ); +} + +inline ::sycl::nd_range<3> nd_range( launch_dim blocks, launch_dim threads ) { + return ::sycl::nd_range<3>( global_range(blocks,threads), local_range(threads) ); +} + +inline ::sycl::nd_item<3> this_item() { + return ::sycl::ext::oneapi::this_work_item::get_nd_item<3>(); +} + +inline ::sycl::sub_group this_subgroup() { + return ::sycl::ext::oneapi::this_work_item::get_sub_group(); +} + +inline ::sycl::group<3> this_group() { + return ::sycl::ext::oneapi::this_work_item::get_work_group<3>(); +} + +template +inline T& local_mem() { + return *::sycl::ext::oneapi::group_local_memory_for_overwrite( this_group() ); +} + +template +void launch_kernel( ::sycl::queue& stream, launch_dim blocks, + launch_dim threads, const char* name, KernelOp op ) { + + if( !blocks.x or !blocks.y or !blocks.z ) return; + + GAUXC_SYCL_ERROR( std::string(name) + " Launch Failed", + stream.parallel_for( nd_range(blocks, threads), + [=](::sycl::nd_item<3>) [[sycl::reqd_sub_group_size(warp_size)]] { + op(); + }) + ); + +} + +} +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_1d.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_1d.cxx new file mode 100644 index 000000000..6ca06fcf8 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_1d.cxx @@ -0,0 +1,388 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "sycl_ssf_1d.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_util.hpp" +#include "common/integrator_constants.hpp" +#include +#include +#include +#include "sycl_atomics.hpp" + +static constexpr auto eps_d = std::numeric_limits::epsilon(); + +namespace GauXC { + + +// SIMT over points: 1D kernel +inline void modify_weights_ssf_kernel_1d( + size_t npts, + size_t natoms, + const double* RAB, + int32_t ldRAB, + const double* coords, + const double* dist_scratch, + int32_t lddist, + const int32_t* iparent_device, + const double* dist_nearest_device, + double* weights_device +) { + + auto it = ::sycl::ext::oneapi::this_work_item::get_nd_item<1>(); + + // Frisch partition functions + auto gFrisch = [](double x) { + + const double s_x = x * 1.5625; // / integrator::magic_ssf_factor<>; + const double s_x2 = s_x * s_x; + const double s_x3 = s_x * s_x2; + const double s_x5 = s_x3 * s_x2; + const double s_x7 = s_x5 * s_x2; + + //return (35.*(s_x - s_x3) + 21.*s_x5 - 5.*s_x7) / 16.; + return ((35.)*(s_x - s_x3) + (21.)*s_x5 - (5.)*s_x7); + }; + + auto sFrisch = [&] (double x) { + if( ::sycl::fabs(x) < integrator::magic_ssf_factor<> ) return (0.5 - (0.5/16.) * gFrisch(x)); + else if( x >= integrator::magic_ssf_factor<> ) return 0.; + else return 1.; + }; + + constexpr double weight_tol = integrator::ssf_weight_tol; + + const size_t tid_x = it.get_global_id(0); + const size_t nt_x = it.get_global_range(0); + + for( size_t ipt = tid_x; ipt < npts; ipt += nt_x ) { + + const auto iParent = iparent_device[ipt]; + + double sum = 0.; + double parent_weight = 0.; + + const double* const local_dist_scratch = dist_scratch + ipt * lddist; + const double dist_cutoff = 0.5 * (1 - integrator::magic_ssf_factor<> ) * + dist_nearest_device[ipt]; + if( local_dist_scratch[iParent] < dist_cutoff ) continue; + + + // Do iParent First + { + + const double ri = local_dist_scratch[ iParent ]; + const double* const local_rab = RAB + iParent * ldRAB; + + parent_weight = 1.; + for( int jCenter = 0; jCenter < natoms; jCenter++ ) + if( parent_weight > weight_tol ) { + if( iParent != jCenter ) { + + const double rj = local_dist_scratch[ jCenter ]; + + const double mu = (ri - rj) * local_rab[ jCenter ]; // XXX: RAB is symmetric + parent_weight *= sFrisch( mu ); + + } + } else break; + + //__syncwarp(); + sum += parent_weight; + + } + + if( parent_weight < eps_d ) { + weights_device[ipt] = 0.; + continue; + } + + for( int iCenter = 0; iCenter < natoms; iCenter++ ) + if( iParent != iCenter ) { + + const double ri = local_dist_scratch[ iCenter ]; + + const double* const local_rab = RAB + iCenter * ldRAB; + + double ps = 1.; + for( int jCenter = 0; jCenter < natoms; jCenter++ ) + if( ps > weight_tol ) { + if( iCenter != jCenter ) { + + const double rj = local_dist_scratch[ jCenter ]; + + const double mu = (ri - rj) * local_rab[ jCenter ]; // XXX: RAB is symmetric + ps *= sFrisch( mu ); + + } + } else break; + + //__syncwarp(); + sum += ps; + + } + weights_device[ipt] *= parent_weight / sum; + } + + +} + +void partition_weights_ssf_1d( int32_t npts, int32_t natoms, const double* RAB, + int32_t ldRAB, const double* coords, const double* dist, int32_t lddist, + const int32_t* iparent, const double* dist_nearest, double* weights, + ::sycl::queue& stream ) { + + const size_t threads = sycl::max_threads_per_thread_block; + const size_t blocks = util::div_ceil( npts, threads ); + + GAUXC_SYCL_ERROR( "SSF 1D Weights Launch Failed", + stream.parallel_for( ::sycl::nd_range<1>( blocks*threads, threads ), + [=](::sycl::nd_item<1>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + modify_weights_ssf_kernel_1d( npts, natoms, RAB, ldRAB, coords, + dist, lddist, iparent, dist_nearest, weights ); + }) + ); + +} + +inline void eval_weight_1st_deriv_contracted_ssf_kernel_1d( + size_t npts, + size_t natoms, + const double* RAB, + int32_t ldRAB, + const double* coords, + const double* points_x, + const double* points_y, + const double* points_z, + const double* dist_scratch, + int32_t lddist, + const int32_t* iparent_device, + const double* dist_nearest_device, + const double* __restrict__ w_times_f_device, + double* __restrict__ exc_grad_w_device +) { + + auto it = ::sycl::ext::oneapi::this_work_item::get_nd_item<1>(); + + // Frisch partition functions + auto gFrisch = [](double x) { + + const double s_x = x * 1.5625; // / integrator::magic_ssf_factor<>; + const double s_x2 = s_x * s_x; + const double s_x3 = s_x * s_x2; + const double s_x5 = s_x3 * s_x2; + const double s_x7 = s_x5 * s_x2; + + return ((35.)*(s_x - s_x3) + (21.)*s_x5 - (5.)*s_x7); + }; + + auto sFrisch = [&] (double x) { + if( ::sycl::fabs(x) < integrator::magic_ssf_factor<> ) return (0.5 - (0.5/16.) * gFrisch(x)); + else if( x >= integrator::magic_ssf_factor<> ) return 0.; + else return 1.; + }; + + auto tFrisch = [&](double x) { + const double s_x = x * 1.5625; // / integrator::magic_ssf_factor<>; + const double s_x2 = s_x * s_x; + const double s_x3 = s_x * s_x2; + const double numerator = (35.) * (s_x3 + (3.) * s_x2 + (3.) * s_x + (1.)); + const double denominator = (x - integrator::magic_ssf_factor<>) * ((5.)*s_x3 + (20.)*s_x2 + (29.)*s_x + (16.)); + return numerator / denominator ; + }; + + constexpr double safe_magic_ssf_bound = integrator::magic_ssf_factor<> - 1e-4; + constexpr double weight_tol = integrator::ssf_weight_tol; + constexpr double w_times_f_thresh = 1.e-12; + + const size_t tid_x = it.get_global_id(0); + const size_t nt_x = it.get_global_range(0); + + for( size_t ipt = tid_x; ipt < npts; ipt += nt_x ) { + + const auto w_times_f_i = w_times_f_device[ipt]; + if (::sycl::fabs(w_times_f_i) < w_times_f_thresh) continue; // weight derivative = 0 when p_A = 0 + const auto iParent = iparent_device[ipt]; + + double sum = 0.; + double parent_weight = 0.; + + const double* const local_dist_scratch = dist_scratch + ipt * lddist; + const double dist_cutoff = 0.18 * dist_nearest_device[ipt]; // 0.5 * (1-integrator::magic_ssf_factor<>) * task.dist_nearest + if( local_dist_scratch[iParent] < dist_cutoff ) continue; //weight derivative = 0 when p_A = 1 + + // Do iParent First + { + const double ri = local_dist_scratch[ iParent ]; + const double* const local_rab = RAB + iParent * ldRAB; + + parent_weight = 1.; + for( int jCenter = 0; jCenter < natoms; jCenter++ ) + if( parent_weight > weight_tol ) { + if( iParent != jCenter ) { + + const double rj = local_dist_scratch[ jCenter ]; + + const double mu = (ri - rj) * local_rab[ jCenter ]; // XXX: RAB is symmetric + parent_weight *= sFrisch( mu ); + + } + } else break; + + sum += parent_weight; + } + + // caculate sum + for( int iCenter = 0; iCenter < natoms; iCenter++ ) + if ( iParent != iCenter ) { + const double ri = local_dist_scratch[ iCenter ]; + const double* const local_rab = RAB + iCenter * ldRAB; + double ps = 1.; + for( int jCenter = 0; jCenter < natoms; jCenter++ ) + if( ps > weight_tol ) { + if( iCenter != jCenter ) { + + const double rj = local_dist_scratch[ jCenter ]; + const double mu = (ri - rj) * local_rab[ jCenter ]; // XXX: RAB is symmetric + ps *= sFrisch( mu ); + } + } else break; + + sum += ps; + + } + + double sum_inv = 1. / sum; + + const double point_x = points_x[ipt]; + const double point_y = points_y[ipt]; + const double point_z = points_z[ipt]; + + // Now do derivative + for( int iB = 0; iB < natoms; iB++ ) + if( iParent != iB ) + { + double exc_grad_w_iBx = 0.0, exc_grad_w_iBy = 0.0, exc_grad_w_iBz = 0.0; + + const double* const local_Rinv_B = RAB + iB * ldRAB; + const double rB = local_dist_scratch[ iB ]; + const double coords_B_x = coords[3*iB + 0]; + const double coords_B_y = coords[3*iB + 1]; + const double coords_B_z = coords[3*iB + 2]; + + // first term + const double rA = local_dist_scratch[ iParent ]; + const double rAB_inv = local_Rinv_B[ iParent ]; + const double mu_AB = (rA - rB) * rAB_inv; + if( ::sycl::fabs(mu_AB) < safe_magic_ssf_bound) { + // first term is tFrisch(mu_AB) * (PA-Z)/Z * w_times_f_i * nabla_B mu_BA + double coef1 = tFrisch(mu_AB) * rAB_inv * (parent_weight - sum) * sum_inv * w_times_f_i / rB; + exc_grad_w_iBx = coef1 * (coords_B_x - point_x + mu_AB * ( coords_B_x - coords[3*iParent + 0]) * rAB_inv * rB); + exc_grad_w_iBy = coef1 * (coords_B_y - point_y + mu_AB * ( coords_B_y - coords[3*iParent + 1]) * rAB_inv * rB); + exc_grad_w_iBz = coef1 * (coords_B_z - point_z + mu_AB * ( coords_B_z - coords[3*iParent + 2]) * rAB_inv * rB); + } + + // second term and third term + // first need to calculate PB + double PB = 1.; + for( int jCenter = 0; jCenter < natoms; jCenter++ ) + if( PB > weight_tol ) { + if( iB != jCenter ) { + const double rj = local_dist_scratch[ jCenter ]; + const double mu = (rB - rj) * local_Rinv_B[ jCenter ]; + PB *= sFrisch( mu ); + } + } else break; + + if( PB > weight_tol ) + for( int iC = 0; iC < natoms; iC++ ) { + if (iB == iC) continue; + const double rBC_inv = local_Rinv_B[iC]; + const double rC = local_dist_scratch[iC]; + const double mu_BC = (rB - rC) * rBC_inv; + + if(::sycl::fabs(mu_BC) < safe_magic_ssf_bound){ + const double t_BC = tFrisch(mu_BC); + const double coef = PB * t_BC * rBC_inv * sum_inv * w_times_f_i; + + const double coords_C_x = coords[3*iC + 0]; + const double coords_C_y = coords[3*iC + 1]; + const double coords_C_z = coords[3*iC + 2]; + + // second term + { + const double rB_inv = 1. / rB; + exc_grad_w_iBx -= coef * ((coords_B_x - point_x) * rB_inv - mu_BC * (coords_B_x - coords_C_x) * rBC_inv); + exc_grad_w_iBy -= coef * ((coords_B_y - point_y) * rB_inv - mu_BC * (coords_B_y - coords_C_y) * rBC_inv); + exc_grad_w_iBz -= coef * ((coords_B_z - point_z) * rB_inv - mu_BC * (coords_B_z - coords_C_z) * rBC_inv); + } + + if(iC != iParent) { + // third term + const double rC_inv = 1. / rC; + const double C_x = coef * ((coords_C_x - point_x) * rC_inv + mu_BC * (coords_C_x - coords_B_x) * rBC_inv); + const double C_y = coef * ((coords_C_y - point_y) * rC_inv + mu_BC * (coords_C_y - coords_B_y) * rBC_inv); + const double C_z = coef * ((coords_C_z - point_z) * rC_inv + mu_BC * (coords_C_z - coords_B_z) * rBC_inv); + + atomic_add_device(exc_grad_w_device + 3*iC + 0, C_x); + atomic_add_device(exc_grad_w_device + 3*iC + 1, C_y); + atomic_add_device(exc_grad_w_device + 3*iC + 2, C_z); + + // Update parent atom + atomic_add_device(exc_grad_w_device + 3*iParent + 0, -C_x); + atomic_add_device(exc_grad_w_device + 3*iParent + 1, -C_y); + atomic_add_device(exc_grad_w_device + 3*iParent + 2, -C_z); + } + } + } + + atomic_add_device(exc_grad_w_device + 3*iB + 0, exc_grad_w_iBx); + atomic_add_device(exc_grad_w_device + 3*iB + 1, exc_grad_w_iBy); + atomic_add_device(exc_grad_w_device + 3*iB + 2, exc_grad_w_iBz); + + // Update parent atom + atomic_add_device(exc_grad_w_device + 3*iParent + 0, -exc_grad_w_iBx); + atomic_add_device(exc_grad_w_device + 3*iParent + 1, -exc_grad_w_iBy); + atomic_add_device(exc_grad_w_device + 3*iParent + 2, -exc_grad_w_iBz); + + } + + } + +} + + + +void eval_weight_1st_deriv_contracted_ssf_1d( int32_t npts, int32_t natoms, const double* RAB, + int32_t ldRAB, const double* coords, + const double* points_x, const double* points_y, const double* points_z, + const double* dist, int32_t lddist, + const int32_t* iparent, const double* dist_nearest, const double* w_times_f, + double* exc_grad_w, ::sycl::queue& stream){ + + const size_t threads = sycl::max_threads_per_thread_block/4; + const size_t blocks = util::div_ceil( npts, threads ); + + GAUXC_SYCL_ERROR( "SSF 1D Weight Derivative Launch Failed", + stream.parallel_for( ::sycl::nd_range<1>( blocks*threads, threads ), + [=](::sycl::nd_item<1>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + eval_weight_1st_deriv_contracted_ssf_kernel_1d( npts, natoms, RAB, + ldRAB, coords, points_x, points_y, points_z, dist, lddist, iparent, + dist_nearest, w_times_f, exc_grad_w ); + }) + ); + +} + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_1d.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_1d.hpp new file mode 100644 index 000000000..e3af02453 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_1d.hpp @@ -0,0 +1,29 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include + +namespace GauXC { + +void partition_weights_ssf_1d( int32_t npts, int32_t natoms, const double* RAB, + int32_t ldRAB, const double* coords, const double* dist, int32_t lddist, + const int32_t* iparent, const double* dist_nearest, double* weights, + ::sycl::queue& stream); + +void eval_weight_1st_deriv_contracted_ssf_1d( int32_t npts, int32_t natoms, const double* RAB, + int32_t ldRAB, const double* coords, + const double* points_x, const double* points_y, const double* points_z, + const double* dist, int32_t lddist, + const int32_t* iparent, const double* dist_nearest, const double* w_times_f, + double* exc_grad_w, ::sycl::queue& stream); + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_2d.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_2d.cxx new file mode 100644 index 000000000..609534e02 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_2d.cxx @@ -0,0 +1,53 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "sycl_ssf_2d.hpp" +#include "device/sycl/sycl_aos_scheme1.hpp" +#include "device_specific/sycl_util.hpp" + +namespace GauXC { + +void partition_weights_ssf_2d( int32_t npts, int32_t natoms, const double* RAB, + int32_t ldRAB, const double* coords, const double* dist, int32_t lddist, + const int32_t* iparent, const double* dist_nearest, double* weights, + ::sycl::queue& stream ) { + + constexpr auto weight_unroll = + alg_constants::SyclAoSScheme1::weight_unroll; + constexpr auto weight_thread_block = + alg_constants::SyclAoSScheme1::weight_thread_block; + constexpr auto weight_thread_block_per_sm = + alg_constants::SyclAoSScheme1::weight_thread_block_per_sm; + + const auto num_sm = + stream.get_device().get_info< ::sycl::info::device::max_compute_units >(); + + constexpr auto warps_per_block = weight_thread_block / sycl::warp_size; + + ::sycl::range<2> local ( warps_per_block, sycl::warp_size ); + ::sycl::range<2> global( num_sm * weight_thread_block_per_sm * warps_per_block, + sycl::warp_size ); + + GAUXC_SYCL_ERROR( "SSF 2D Weights Launch Failed", + stream.submit([&](::sycl::handler& cgh) { + cgh.parallel_for( ::sycl::nd_range<2>(global, local), + [=](::sycl::nd_item<2>) + [[sycl::reqd_sub_group_size(GauXC::sycl::warp_size)]] { + modify_weights_ssf_kernel_2d< weight_unroll, weight_thread_block, + weight_thread_block_per_sm >( npts, natoms, RAB, + ldRAB, coords, dist, lddist, iparent, dist_nearest, weights ); + }); + }) + ); + +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_2d.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_2d.hpp new file mode 100644 index 000000000..958f3ae02 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/sycl_ssf_2d.hpp @@ -0,0 +1,210 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include +#include + +#include "sycl_extensions.hpp" +#include "sycl_atomics.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_vector_types.hpp" +#include "common/integrator_constants.hpp" + +inline constexpr static auto eps_d = std::numeric_limits::epsilon(); + +namespace GauXC { + +inline double gFrisch(double x) { + // Frisch partition functions + const double s_x = x * 1.5625; + const double s_x2 = s_x * s_x; + const double s_x3 = s_x * s_x2; + const double s_x5 = s_x3 * s_x2; + const double s_x7 = s_x5 * s_x2; + + return ((35.) *(s_x - s_x3) + (21.) *s_x5 - (5.) *s_x7); +} + + +inline double sFrisch(double x) { + + if( ::sycl::fabs(x) < integrator::magic_ssf_factor<> ) return (0.5 - (0.5/ 16.0) * gFrisch(x)); + else if( x >= integrator::magic_ssf_factor<> ) return 0.; + else return 1.; +} + +template +void modify_weights_ssf_kernel_2d( + int32_t npts, int32_t natoms, + const double* RAB, int32_t ldRAB, const double* coords, const double* dist, + size_t lddist, const int32_t* iparent, const double* dist_nearest, + double* weights ) { + + static_assert( weight_unroll == 4 ); + + auto it = ::sycl::ext::oneapi::this_work_item::get_nd_item<2>(); + constexpr uint32_t warps_per_block = weight_thread_block / sycl::warp_size; + auto& jCounter_sm = *::sycl::ext::oneapi::group_local_memory_for_overwrite( + ::sycl::ext::oneapi::this_work_item::get_work_group<2>() ); + + constexpr double weight_tol = integrator::ssf_weight_tol; + + const auto sg = it.get_sub_group(); + + const auto local_x = it.get_local_id(1); + const auto local_y = it.get_local_id(0); + const auto range_x = it.get_local_range(1); + + int natom_block = ((natoms + range_x - 1) / range_x) * range_x; + + const size_t tid_x = it.get_global_id(0); + const size_t nt_x = it.get_global_range(0); + + int* jCounter = &jCounter_sm[local_y]; + + // Each sub-group will work together on a point + for( size_t ipt = tid_x; ipt < (size_t)npts; ipt += nt_x ) { + + const auto iParent = iparent[ipt]; + + double sum = 0.; + double parent_weight = 0.; + + const double* const local_dist_scratch = dist + ipt * lddist; + const double dist_cutoff = 0.5 * (1 - integrator::magic_ssf_factor<> ) * + dist_nearest[ipt]; + if( local_dist_scratch[iParent] < dist_cutoff ) continue; + + // Do iParent First + { + + const double ri = local_dist_scratch[ iParent ]; + const double* const local_rab = RAB + iParent * ldRAB; + + parent_weight = 1.; + for( int jCenter = local_x; jCenter < natom_block; jCenter += range_x ) { + double contribution = 1.0; + if (jCenter < natoms && iParent != jCenter) { + const double rj = local_dist_scratch[ jCenter ]; + const double mu = (ri - rj) * local_rab[ jCenter ]; // XXX: RAB is symmetric + contribution = sFrisch( mu ); + } + contribution = sycl::warp_reduce_prod(contribution); + contribution = ::sycl::group_broadcast(sg, contribution, 0); + + parent_weight *= contribution; + + if (parent_weight < weight_tol) break; + } + } + + if( parent_weight < eps_d ) { + if (local_x == 0) + weights[ipt] = 0.; + ::sycl::group_barrier(sg); + continue; + } + + // Initialize each counter to 0 + if (local_x == 0) { + jCounter[0] = 0; + } + ::sycl::group_barrier(sg); + + // Each work-item will process an iCenter. Atomic operations are used to + // assign an iCenter value to each work-item. + int iCenter = atomic_fetch_add_local(jCounter, 1); + if (iCenter >= iParent) iCenter++; // iCenter == iParent is skipped + + // The entire sub-group processes the same jCenter value at the same time + int jCenter = 0; + + const double* local_rab = RAB + iCenter * ldRAB; + double ri = local_dist_scratch[ iCenter ]; + double ps = 1.; + int iCount = 0; + bool cont = (iCenter < natoms); + + // We will continue iterating until all of the work-items have cont == false + while (::sycl::any_of_group(sg, cont)) { + if (cont) { + double2 rj[weight_unroll/2]; + double2 rab_val[weight_unroll/2]; + double mu[weight_unroll]; + iCount += weight_unroll; + + #pragma unroll + for (int k = 0; k < weight_unroll/2; k++) { + rj[k] = *((const double2*)(local_dist_scratch + jCenter) + k); + rab_val[k] = *((const double2*)(local_rab + jCenter) + k); + } + + #pragma unroll + for (int k = 0; k < weight_unroll/2; k++) { + mu[2*k+0] = (ri - rj[k].x) * rab_val[k].x; // XXX: RAB is symmetric + mu[2*k+1] = (ri - rj[k].y) * rab_val[k].y; + } + + #pragma unroll + for (int k = 0; k < weight_unroll; k++) { + if((iCenter != jCenter + k) && (jCenter + k < natoms)) { + mu[k] = sFrisch( mu[k] ); + ps *= mu[k]; + } + } + + // A work-item is done with an iCenter based on 2 conditions: weight + // tolerance, or if it has seen all of the jCenters + if( !(ps > weight_tol && iCount < (int)lddist )) { + // In that case the work-item begins processing another iCenter + sum += ps; + iCenter = atomic_fetch_add_local(jCounter, 1); + if (iCenter >= iParent) iCenter++; + + // If there are no more iCenters left to process, signal ready to exit + cont = (iCenter < natoms); + ri = local_dist_scratch[ iCenter ]; + local_rab = RAB + iCenter * ldRAB; + ps = 1.; + iCount = 0; + } + } + // Wraps jCenter around. This was faster than modulo + jCenter += weight_unroll; + jCenter = (jCenter < ldRAB) ? jCenter : 0; + } + + // All of the work-items then sum their contributions. Only lane 0 needs to + // add the parent contribution. + ::sycl::group_barrier(sg); + sum = sycl::warp_reduce_sum(sum); + if (local_x == 0) { + sum += parent_weight; + weights[ipt] *= parent_weight / sum; + } + + ::sycl::group_barrier(sg); + } + +} + + +void partition_weights_ssf_2d( int32_t npts, int32_t natoms, const double* RAB, + int32_t ldRAB, const double* coords, const double* dist, int32_t lddist, + const int32_t* iparent, const double* dist_nearest, double* weights, + ::sycl::queue& stream ); + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/symmetrize_mat.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/symmetrize_mat.cxx new file mode 100644 index 000000000..6e6b10d7d --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/symmetrize_mat.cxx @@ -0,0 +1,134 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/symmetrize_mat.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_util.hpp" +#include "sycl_launch.hpp" +#include +#include + +namespace GauXC { + +void symmetrize_matrix_device( size_t N, double* A, size_t LDA ) { + + auto it = GauXC::sycl::this_item(); + + constexpr uint32_t block_size = sycl::warp_size; + + auto& buffer = GauXC::sycl::local_mem(); + + const size_t num_blocks = ((N + block_size - 1) / block_size); + + for (int i = it.get_group(2); i < num_blocks; i += it.get_group_range(2)) { + // TODO This could be load balanced if need be + const int i_coord = i * block_size; + for (int j = i; j < num_blocks; j++) { + const int j_coord = j * block_size; + + // Read in block to buffer + // TODO These could be vector reads/writes if this becomes significant + if (i_coord + it.get_local_id(1) < N && j_coord + it.get_local_id(2) < N) { + buffer[it.get_local_id(1)][it.get_local_id(2)] = + A[(i_coord + it.get_local_id(1)) * LDA + j_coord + it.get_local_id(2)]; + } + ::sycl::group_barrier( it.get_group() ); + + // Write buffer + if (j_coord + it.get_local_id(1) < N && i_coord + it.get_local_id(2) < N) { + if ((j_coord != i_coord || it.get_local_id(2) < it.get_local_id(1))) { // handles the diagonal block + A[(j_coord + it.get_local_id(1)) * LDA + i_coord + it.get_local_id(2)] = + buffer[it.get_local_id(2)][it.get_local_id(1)]; + } + } + ::sycl::group_barrier( it.get_group() ); + } + } +} + +void symmetrize_matrix_inc_device( size_t N, double* A, size_t LDA ) { + + auto it = GauXC::sycl::this_item(); + + constexpr uint32_t block_size = sycl::warp_size; + + auto& buffer_0 = GauXC::sycl::local_mem(); + auto& buffer_1 = GauXC::sycl::local_mem(); + + const size_t num_blocks = ((N + block_size - 1) / block_size); + + for (int i = it.get_group(2); i < num_blocks; i += it.get_group_range(2)) { + // TODO This could be load balanced if need be + const int i_coord = i * block_size; + for (int j = i; j < num_blocks; j++) { + const int j_coord = j * block_size; + + // Read in block to buffer + // TODO These could be vector reads/writes if this becomes significant + if (i_coord + it.get_local_id(1) < N && j_coord + it.get_local_id(2) < N) { + buffer_0[it.get_local_id(1)][it.get_local_id(2)] = + A[(i_coord + it.get_local_id(1)) * LDA + j_coord + it.get_local_id(2)]; + } + if (j_coord + it.get_local_id(1) < N && i_coord + it.get_local_id(2) < N) { + buffer_1[it.get_local_id(1)][it.get_local_id(2)] = + A[(j_coord + it.get_local_id(1)) * LDA + i_coord + it.get_local_id(2)]; + } + ::sycl::group_barrier( it.get_group() ); + + buffer_0[it.get_local_id(1)][it.get_local_id(2)] += + buffer_1[it.get_local_id(2)][it.get_local_id(1)]; + buffer_0[it.get_local_id(1)][it.get_local_id(2)] *= 0.5; + ::sycl::group_barrier( it.get_group() ); + + // Write buffer + if (j_coord + it.get_local_id(1) < N && i_coord + it.get_local_id(2) < N) { + //if ((j_coord != i_coord || it.get_local_id(2) < it.get_local_id(1))) { // handles the diagonal block + A[(j_coord + it.get_local_id(1)) * LDA + i_coord + it.get_local_id(2)] = + buffer_0[it.get_local_id(2)][it.get_local_id(1)]; + //} + } + if (i_coord + it.get_local_id(1) < N && j_coord + it.get_local_id(2) < N) { + //if ((j_coord != i_coord || it.get_local_id(2) > it.get_local_id(1))) { // handles the diagonal block + A[(i_coord + it.get_local_id(1)) * LDA + j_coord + it.get_local_id(2)] = + buffer_0[it.get_local_id(1)][it.get_local_id(2)]; + //} + } + ::sycl::group_barrier( it.get_group() ); + } + } +} + + + +void symmetrize_matrix( int32_t N, double* A, size_t LDA, device_queue queue ) { + ::sycl::queue& stream = queue.queue_as(); + const size_t num_blocks = ((N + sycl::warp_size - 1) / sycl::warp_size); + // Warp size must equal max_warps_per_thread_block must equal 32 + sycl::launch_dim threads(sycl::warp_size, sycl::max_warps_per_thread_block, 1), blocks(num_blocks); + + sycl::launch_kernel( stream, blocks, threads, "symmetrize_matrix_device", + [=](){ + symmetrize_matrix_device( N, A, LDA ); + }); +} + +void symmetrize_matrix_inc( int32_t N, double* A, size_t LDA, device_queue queue ) { + ::sycl::queue& stream = queue.queue_as(); + const size_t num_blocks = ((N + sycl::warp_size - 1) / sycl::warp_size); + // Warp size must equal max_warps_per_thread_block must equal 32 + sycl::launch_dim threads(sycl::warp_size, sycl::max_warps_per_thread_block, 1), blocks(num_blocks); + + sycl::launch_kernel( stream, blocks, threads, "symmetrize_matrix_inc_device", + [=](){ + symmetrize_matrix_inc_device( N, A, LDA ); + }); +} +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars.cxx new file mode 100644 index 000000000..61b0b6120 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars.cxx @@ -0,0 +1,267 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/uvvars.hpp" +#include "device_specific/sycl_util.hpp" +#include "sycl_launch.hpp" +#include + +#include "uvvars_lda.hpp" +#include "uvvars_gga.hpp" +#include "uvvars_mgga.hpp" + +namespace GauXC { + +#define EVAL_UVARS_KERNEL(xc_approx) \ + ::sycl::queue& stream = queue.queue_as(); \ + switch ( ks_scheme ) { \ + case RKS: \ + sycl::launch_kernel( stream, blocks, threads, "eval_uvars_" #xc_approx "_rks_kernel", \ + [=](){ eval_uvars_##xc_approx##_rks_kernel( ntasks, device_tasks ); } ); \ + break; \ + case UKS: \ + sycl::launch_kernel( stream, blocks, threads, "eval_uvars_" #xc_approx "_uks_kernel", \ + [=](){ eval_uvars_##xc_approx##_uks_kernel( ntasks, device_tasks ); } ); \ + break; \ + case GKS: \ + sycl::launch_kernel( stream, blocks, threads, "eval_uvars_" #xc_approx "_gks_kernel", \ + [=](){ eval_uvars_##xc_approx##_gks_kernel( ntasks, device_tasks ); } ); \ + break; \ + default: \ + GAUXC_GENERIC_EXCEPTION( "Unexpected KS scheme when attempting to evaluate U vars" ); \ + } + + +#define EVAL_TMAT_KERNEL(xc_approx) \ + ::sycl::queue& stream = queue.queue_as(); \ + switch ( ks_scheme ) { \ + case RKS: \ + sycl::launch_kernel( stream, blocks, threads, "eval_tmat_" #xc_approx "_rks_kernel", \ + [=](){ eval_tmat_##xc_approx##_rks_kernel( ntasks, device_tasks ); } ); \ + break; \ + case UKS: \ + sycl::launch_kernel( stream, blocks, threads, "eval_tmat_" #xc_approx "_uks_kernel", \ + [=](){ eval_tmat_##xc_approx##_uks_kernel( ntasks, device_tasks ); } ); \ + break; \ + case GKS: \ + GAUXC_GENERIC_EXCEPTION( "GKS + evaluate trial U vars NYI" ); \ + break; \ + default: \ + GAUXC_GENERIC_EXCEPTION( "Unexpected KS scheme when attempting to evaluate U vars" ); \ + } + + +#define EVAL_VVARS_KERNEL(xc_approx) \ + ::sycl::queue& stream = queue.queue_as(); \ + switch ( den_select ) { \ + case DEN_S: \ + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_" #xc_approx "_kern", \ + [=](){ eval_vvar_##xc_approx##_kern( ntasks, device_tasks ); } ); \ + break; \ + case DEN_Z: \ + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_" #xc_approx "_kern", \ + [=](){ eval_vvar_##xc_approx##_kern( ntasks, device_tasks ); } ); \ + break; \ + case DEN_Y: \ + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_" #xc_approx "_kern", \ + [=](){ eval_vvar_##xc_approx##_kern( ntasks, device_tasks ); } ); \ + break; \ + case DEN_X: \ + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_" #xc_approx "_kern", \ + [=](){ eval_vvar_##xc_approx##_kern( ntasks, device_tasks ); } ); \ + break; \ + default: \ + GAUXC_GENERIC_EXCEPTION( "Unexpected KS scheme when attempting to evaluate V vars" ); \ + } + +// Internal implementation with trial parameter +void eval_tmat_lda( size_t ntasks, int32_t npts_max, integrator_ks_scheme ks_scheme, + XCDeviceTask* device_tasks, device_queue queue ) { + sycl::launch_dim threads( sycl::max_warps_per_thread_block * sycl::warp_size, 1, 1 ); + sycl::launch_dim blocks( util::div_ceil( npts_max, threads.x ), 1, ntasks ); + EVAL_TMAT_KERNEL(lda); +} + +void eval_uvars_lda( size_t ntasks, int32_t npts_max, integrator_ks_scheme ks_scheme, + XCDeviceTask* device_tasks, device_queue queue ) { + sycl::launch_dim threads( sycl::max_warps_per_thread_block * sycl::warp_size, 1, 1 ); + sycl::launch_dim blocks( util::div_ceil( npts_max, threads.x ), 1, ntasks ); + EVAL_UVARS_KERNEL(lda); +} + +// Internal implementation with trial as template parameter +template +void eval_vvars_lda_impl( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + XCDeviceTask* device_tasks, device_queue queue ) { + sycl::launch_dim threads( sycl::warp_size, sycl::max_warps_per_thread_block, 1 ); + sycl::launch_dim blocks( util::div_ceil( nbf_max, threads.x ), + util::div_ceil( npts_max, threads.y ), + ntasks ); + EVAL_VVARS_KERNEL(lda); +} +void eval_vvars_lda( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + XCDeviceTask* device_tasks, device_queue queue ) { + eval_vvars_lda_impl(ntasks, nbf_max, npts_max, den_select, device_tasks, queue); +} +void eval_vvars_lda_trial( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + XCDeviceTask* device_tasks, device_queue queue ) { + eval_vvars_lda_impl(ntasks, nbf_max, npts_max, den_select, device_tasks, queue); +} + +// Internal implementation with trial parameter +void eval_tmat_gga( size_t ntasks, int32_t npts_max, integrator_ks_scheme ks_scheme, + XCDeviceTask* device_tasks, device_queue queue ) { + sycl::launch_dim threads( GGA_KERNEL_SM_WARPS * sycl::warp_size, 1, 1 ); + sycl::launch_dim blocks( util::div_ceil( npts_max, threads.x ), 1, ntasks ); + EVAL_TMAT_KERNEL(gga); +} +void eval_uvars_gga( size_t ntasks, int32_t npts_max, integrator_ks_scheme ks_scheme, + XCDeviceTask* device_tasks, device_queue queue ) { + sycl::launch_dim threads( GGA_KERNEL_SM_WARPS * sycl::warp_size, 1, 1 ); + sycl::launch_dim blocks( util::div_ceil( npts_max, threads.x ), 1, ntasks ); + EVAL_UVARS_KERNEL(gga); +} + + +// Internal implementation with trial as template parameter +template +void eval_vvars_gga_impl( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + XCDeviceTask* device_tasks, device_queue queue ) { + ::sycl::queue& stream = queue.queue_as(); + sycl::launch_dim threads( sycl::warp_size, sycl::max_warps_per_thread_block, 1 ); + sycl::launch_dim blocks( util::div_ceil( nbf_max, threads.x ), + util::div_ceil( npts_max, threads.y ), + ntasks ); + + switch ( den_select ) { + case DEN_S: + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_gga_kern", + [=](){ eval_vvar_gga_kern( ntasks, device_tasks ); } ); + break; + case DEN_Z: + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_gga_kern", + [=](){ eval_vvar_gga_kern( ntasks, device_tasks ); } ); + break; + case DEN_Y: + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_gga_kern", + [=](){ eval_vvar_gga_kern( ntasks, device_tasks ); } ); + break; + case DEN_X: + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_gga_kern", + [=](){ eval_vvar_gga_kern( ntasks, device_tasks ); } ); + break; + default: + GAUXC_GENERIC_EXCEPTION( "Unexpected KS scheme when attempting to evaluate V vars" ); + } +} +void eval_vvars_gga( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + XCDeviceTask* device_tasks, device_queue queue ) { + eval_vvars_gga_impl(ntasks, nbf_max, npts_max, den_select, device_tasks, queue); +} +void eval_vvars_gga_trial( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + XCDeviceTask* device_tasks, device_queue queue ) { + eval_vvars_gga_impl(ntasks, nbf_max, npts_max, den_select, device_tasks, queue); +} + +// Internal implementation with trial parameter +void eval_tmat_mgga( size_t ntasks, int32_t npts_max, integrator_ks_scheme ks_scheme, + bool need_lapl, XCDeviceTask* device_tasks, device_queue queue ) { + ::sycl::queue& stream = queue.queue_as(); + + sycl::launch_dim threads( GGA_KERNEL_SM_WARPS * sycl::warp_size, 1, 1 ); + sycl::launch_dim blocks( util::div_ceil( npts_max, threads.x ), 1, ntasks ); + + if(need_lapl) { + GAUXC_GENERIC_EXCEPTION("MGGA + LAPL + eval tmat NYI"); + } + if(ks_scheme == RKS) { + sycl::launch_kernel( stream, blocks, threads, "eval_tmat_mgga_rks_kernel", + [=](){ eval_tmat_mgga_rks_kernel( ntasks, device_tasks ); } ); + } else if(ks_scheme == UKS) { + sycl::launch_kernel( stream, blocks, threads, "eval_tmat_mgga_uks_kernel", + [=](){ eval_tmat_mgga_uks_kernel( ntasks, device_tasks ); } ); + } else { + GAUXC_GENERIC_EXCEPTION("GKS + MGGA + DEVICE NYI"); + } +} + +void eval_uvars_mgga( size_t ntasks, int32_t npts_max, integrator_ks_scheme ks_scheme, + bool need_lapl, XCDeviceTask* device_tasks, device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + // Evaluate GAMMA + eval_uvars_gga(ntasks, npts_max, ks_scheme, device_tasks, queue); + + if(ks_scheme == RKS) { + return; // Nothing left to do + } else if(ks_scheme == UKS) { + sycl::launch_dim threads( sycl::max_warps_per_thread_block * sycl::warp_size, 1, 1 ); + sycl::launch_dim blocks( util::div_ceil( npts_max, threads.x ), 1, ntasks ); + if(need_lapl) { + sycl::launch_kernel( stream, blocks, threads, "eval_uvars_mgga_uks_kernel", + [=](){ eval_uvars_mgga_uks_kernel( ntasks, device_tasks ); } ); + } else { + sycl::launch_kernel( stream, blocks, threads, "eval_uvars_mgga_uks_kernel", + [=](){ eval_uvars_mgga_uks_kernel( ntasks, device_tasks ); } ); + } + } else { + GAUXC_GENERIC_EXCEPTION("GKS + MGGA + DEVICE NYI"); + } + +} + +// Internal implementation with trial as template parameter +template +void eval_vvars_mgga_impl( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + bool need_lapl, XCDeviceTask* device_tasks, device_queue queue ) { + // First evaluate GGA variables + eval_vvars_gga_impl(ntasks, nbf_max, npts_max, den_select, device_tasks, queue); + + ::sycl::queue& stream = queue.queue_as(); + sycl::launch_dim threads( sycl::warp_size, sycl::max_warps_per_thread_block, 1 ); + sycl::launch_dim blocks( util::div_ceil( nbf_max, threads.x ), + util::div_ceil( npts_max, threads.y ), + ntasks ); + + switch ( den_select ) { + case DEN_S: + if (need_lapl) { + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_mgga_kern", + [=](){ eval_vvar_mgga_kern( ntasks, device_tasks ); } ); + } else { + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_mgga_kern", + [=](){ eval_vvar_mgga_kern( ntasks, device_tasks ); } ); + } + break; + case DEN_Z: + if (need_lapl) { + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_mgga_kern", + [=](){ eval_vvar_mgga_kern( ntasks, device_tasks ); } ); + } else { + sycl::launch_kernel( stream, blocks, threads, "eval_vvar_mgga_kern", + [=](){ eval_vvar_mgga_kern( ntasks, device_tasks ); } ); + } + break; + default: + GAUXC_GENERIC_EXCEPTION( "Unexpected KS scheme when attempting to evaluate V vars" ); + } +} +void eval_vvars_mgga( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + bool need_lapl, XCDeviceTask* device_tasks, device_queue queue ) { + eval_vvars_mgga_impl(ntasks, nbf_max, npts_max, den_select, need_lapl, device_tasks, queue); +} +void eval_vvars_mgga_trial( size_t ntasks, int32_t nbf_max, int32_t npts_max, density_id den_select, + bool need_lapl, XCDeviceTask* device_tasks, device_queue queue ) { + eval_vvars_mgga_impl(ntasks, nbf_max, npts_max, den_select, need_lapl, device_tasks, queue); +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_gga.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_gga.hpp new file mode 100644 index 000000000..859b3a465 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_gga.hpp @@ -0,0 +1,570 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_util.hpp" +#include "device/xc_device_data.hpp" +#include "sycl_extensions.hpp" +#include "sycl_atomics.hpp" +#include "sycl_launch.hpp" + +#define VVAR_KERNEL_SM_BLOCK 32 +#define GGA_KERNEL_SM_WARPS 16 + +namespace GauXC { + +template +void eval_vvar_gga_kern( size_t ntasks, + XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + constexpr auto warp_size = sycl::warp_size; + auto& den_shared = GauXC::sycl::local_mem(); + + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + + double* den_eval_device = nullptr; + double* den_x_eval_device = nullptr; + double* den_y_eval_device = nullptr; + double* den_z_eval_device = nullptr; + + if constexpr (trial){ + if constexpr (den_select == DEN_S) { + den_eval_device = task.tden_s; + den_x_eval_device = task.tdden_sx; + den_y_eval_device = task.tdden_sy; + den_z_eval_device = task.tdden_sz; + } + if constexpr (den_select == DEN_Z) { + den_eval_device = task.tden_z; + den_x_eval_device = task.tdden_zx; + den_y_eval_device = task.tdden_zy; + den_z_eval_device = task.tdden_zz; + } + if constexpr (den_select == DEN_Y) { + den_eval_device = task.tden_y; + den_x_eval_device = task.tdden_yx; + den_y_eval_device = task.tdden_yy; + den_z_eval_device = task.tdden_yz; + } + if constexpr (den_select == DEN_X) { + den_eval_device = task.tden_x; + den_x_eval_device = task.tdden_xx; + den_y_eval_device = task.tdden_xy; + den_z_eval_device = task.tdden_xz; + } + }else{ + if constexpr (den_select == DEN_S) { + den_eval_device = task.den_s; + den_x_eval_device = task.dden_sx; + den_y_eval_device = task.dden_sy; + den_z_eval_device = task.dden_sz; + } + if constexpr (den_select == DEN_Z) { + den_eval_device = task.den_z; + den_x_eval_device = task.dden_zx; + den_y_eval_device = task.dden_zy; + den_z_eval_device = task.dden_zz; + } + if constexpr (den_select == DEN_Y) { + den_eval_device = task.den_y; + den_x_eval_device = task.dden_yx; + den_y_eval_device = task.dden_yy; + den_z_eval_device = task.dden_yz; + } + if constexpr (den_select == DEN_X) { + den_eval_device = task.den_x; + den_x_eval_device = task.dden_xx; + den_y_eval_device = task.dden_xy; + den_z_eval_device = task.dden_xz; + } + } + + const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + + const auto* den_basis_prod_device = task.zmat; + + const int threadIdx_x = it.get_local_id(2); + const int threadIdx_y = it.get_local_id(1); + const int blockDim_x = it.get_local_range(2); + const int blockDim_y = it.get_local_range(1); + + for ( int bid_x = it.get_group(2) * blockDim_x; + bid_x < nbf; + bid_x += blockDim_x * it.get_group_range(2) ) { + + for ( int bid_y = it.get_group(1) * VVAR_KERNEL_SM_BLOCK; + bid_y < npts; + bid_y += VVAR_KERNEL_SM_BLOCK * it.get_group_range(1) ) { + + for (int sm_y = threadIdx_y; sm_y < VVAR_KERNEL_SM_BLOCK; sm_y += blockDim_y) { + den_shared[0][threadIdx_x][sm_y] = 0.; + den_shared[1][threadIdx_x][sm_y] = 0.; + den_shared[2][threadIdx_x][sm_y] = 0.; + den_shared[3][threadIdx_x][sm_y] = 0.; + + if (bid_y + threadIdx_x < npts and bid_x + sm_y < nbf) { + const double* db_col = den_basis_prod_device + (bid_x + sm_y)*npts; + const double* bf_col = basis_eval_device + (bid_x + sm_y)*npts; + const double* bf_x_col = dbasis_x_eval_device + (bid_x + sm_y)*npts; + const double* bf_y_col = dbasis_y_eval_device + (bid_x + sm_y)*npts; + const double* bf_z_col = dbasis_z_eval_device + (bid_x + sm_y)*npts; + + den_shared[0][threadIdx_x][sm_y] = bf_col [ bid_y + threadIdx_x ] * db_col[ bid_y + threadIdx_x ]; + den_shared[1][threadIdx_x][sm_y] = bf_x_col[ bid_y + threadIdx_x ] * db_col[ bid_y + threadIdx_x ]; + den_shared[2][threadIdx_x][sm_y] = bf_y_col[ bid_y + threadIdx_x ] * db_col[ bid_y + threadIdx_x ]; + den_shared[3][threadIdx_x][sm_y] = bf_z_col[ bid_y + threadIdx_x ] * db_col[ bid_y + threadIdx_x ]; + } + } + ::sycl::group_barrier( it.get_group() ); + + + for (int sm_y = threadIdx_y; sm_y < VVAR_KERNEL_SM_BLOCK; sm_y += blockDim_y) { + const int tid_y = bid_y + sm_y; + double den_reg = den_shared[0][sm_y][threadIdx_x]; + double dx_reg = den_shared[1][sm_y][threadIdx_x]; + double dy_reg = den_shared[2][sm_y][threadIdx_x]; + double dz_reg = den_shared[3][sm_y][threadIdx_x]; + + // Warp blocks are stored col major + den_reg = sycl::warp_reduce_sum( den_reg ); + dx_reg = 2. * sycl::warp_reduce_sum( dx_reg ); + dy_reg = 2. * sycl::warp_reduce_sum( dy_reg ); + dz_reg = 2. * sycl::warp_reduce_sum( dz_reg ); + + + if( threadIdx_x == 0 and tid_y < npts ) { + atomic_add_device( den_eval_device + tid_y, den_reg ); + atomic_add_device( den_x_eval_device + tid_y, dx_reg ); + atomic_add_device( den_y_eval_device + tid_y, dy_reg ); + atomic_add_device( den_z_eval_device + tid_y, dz_reg ); + } + } + ::sycl::group_barrier( it.get_group() ); + } + } + +} + +inline void eval_uvars_gga_rks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + const auto* dden_sx_eval_device = task.dden_sx; + const auto* dden_sy_eval_device = task.dden_sy; + const auto* dden_sz_eval_device = task.dden_sz; + auto* gamma_eval_device = task.gamma; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const double dx = dden_sx_eval_device[ tid ]; + const double dy = dden_sy_eval_device[ tid ]; + const double dz = dden_sz_eval_device[ tid ]; + + gamma_eval_device[ tid ] = dx*dx + dy*dy + dz*dz; + } + +} + +inline void eval_tmat_gga_rks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + const auto* dden_sx_eval_device = task.dden_sx; + const auto* dden_sy_eval_device = task.dden_sy; + const auto* dden_sz_eval_device = task.dden_sz; + const auto* tdden_sx_eval_device = task.tdden_sx; + const auto* tdden_sy_eval_device = task.tdden_sy; + const auto* tdden_sz_eval_device = task.tdden_sz; + + const auto* weight_device = task.weights; + const auto* vgamma_device = task.vgamma; + const auto* v2rho2_device = task.v2rho2; + const auto* v2rhogamma_device = task.v2rhogamma; + const auto* v2gamma2_device = task.v2gamma2; + const auto* trho_device = task.tden_s; + + auto* FXC_A_device = task.FXC_A_s; + auto* FXC_Bx_device = task.FXC_Bx_s; + auto* FXC_By_device = task.FXC_By_s; + auto* FXC_Bz_device = task.FXC_Bz_s; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const auto dx = dden_sx_eval_device[ tid ]; + const auto dy = dden_sy_eval_device[ tid ]; + const auto dz = dden_sz_eval_device[ tid ]; + const auto tdx = tdden_sx_eval_device[ tid ]; + const auto tdy = tdden_sy_eval_device[ tid ]; + const auto tdz = tdden_sz_eval_device[ tid ]; + const auto tgamma = tdx*dx + tdy*dy + tdz*dz; + + const auto FXC_A = v2rho2_device[ tid ] * trho_device[ tid ] + 2.0 * v2rhogamma_device[tid] * tgamma; + const auto B_coef = v2rhogamma_device[tid] * trho_device[tid] + 2.0 * v2gamma2_device[tid] * tgamma; + FXC_A_device[ tid ] = weight_device[ tid ] * FXC_A ; + FXC_Bx_device[ tid ] = 2.0 * weight_device[ tid ] * ( B_coef * dx + vgamma_device[ tid ] * tdx ); + FXC_By_device[ tid ] = 2.0 * weight_device[ tid ] * ( B_coef * dy + vgamma_device[ tid ] * tdy ); + FXC_Bz_device[ tid ] = 2.0 * weight_device[ tid ] * ( B_coef * dz + vgamma_device[ tid ] * tdz ); + } + +} + +inline void eval_uvars_gga_uks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + auto* den_pos_eval_device = task.den_s; + const auto* den_pos_x_eval_device = task.dden_sx; + const auto* den_pos_y_eval_device = task.dden_sy; + const auto* den_pos_z_eval_device = task.dden_sz; + + auto* den_neg_eval_device = task.den_z; + const auto* den_neg_x_eval_device = task.dden_zx; + const auto* den_neg_y_eval_device = task.dden_zy; + const auto* den_neg_z_eval_device = task.dden_zz; + + auto* gamma_pp_eval_device = task.gamma_pp; + auto* gamma_pm_eval_device = task.gamma_pm; + auto* gamma_mm_eval_device = task.gamma_mm; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const double ps = den_pos_eval_device[ tid ]; + const double pz = den_neg_eval_device[ tid ]; + const double dndx = den_pos_x_eval_device[ tid ]; + const double dndy = den_pos_y_eval_device[ tid ]; + const double dndz = den_pos_z_eval_device[ tid ]; + const double dMzdx = den_neg_x_eval_device[ tid ]; + const double dMzdy = den_neg_y_eval_device[ tid ]; + const double dMzdz = den_neg_z_eval_device[ tid ]; + + // (del n).(del n) + const auto dn_sq = dndx*dndx + dndy*dndy + dndz*dndz; + // (del Mz).(del Mz) + const auto dMz_sq = dMzdx*dMzdx + dMzdy*dMzdy + dMzdz*dMzdz; + // (del n).(del Mz) + const auto dn_dMz = dndx*dMzdx + dndy*dMzdy + dndz*dMzdz; + + gamma_pp_eval_device[ tid ] = 0.25*(dn_sq + dMz_sq) + 0.5*dn_dMz; + gamma_pm_eval_device[ tid ] = 0.25*(dn_sq - dMz_sq); + gamma_mm_eval_device[ tid ] = 0.25*(dn_sq + dMz_sq) - 0.5*dn_dMz; + + den_pos_eval_device[ tid ] = 0.5*(ps + pz); + den_neg_eval_device[ tid ] = 0.5*(ps - pz); + } + +} + +inline void eval_tmat_gga_uks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + const auto* tden_s_device = task.tden_s; + const auto* tden_z_device = task.tden_z; + const auto* weight_device = task.weights; + + const auto* tden_pos_x_eval_device = task.tdden_sx; + const auto* tden_pos_y_eval_device = task.tdden_sy; + const auto* tden_pos_z_eval_device = task.tdden_sz; + const auto* den_pos_x_eval_device = task.dden_sx; + const auto* den_pos_y_eval_device = task.dden_sy; + const auto* den_pos_z_eval_device = task.dden_sz; + + const auto* tden_neg_x_eval_device = task.tdden_zx; + const auto* tden_neg_y_eval_device = task.tdden_zy; + const auto* tden_neg_z_eval_device = task.tdden_zz; + const auto* den_neg_x_eval_device = task.dden_zx; + const auto* den_neg_y_eval_device = task.dden_zy; + const auto* den_neg_z_eval_device = task.dden_zz; + + const auto* vgamma_aa_device = task.vgamma_pp; + const auto* vgamma_ab_device = task.vgamma_pm; + const auto* vgamma_bb_device = task.vgamma_mm; + const auto* v2rho2_a_a_device = task.v2rho2_a_a; + const auto* v2rho2_a_b_device = task.v2rho2_a_b; + const auto* v2rho2_b_b_device = task.v2rho2_b_b; + const auto* v2rhogamma_a_aa_device = task.v2rhogamma_a_aa; + const auto* v2rhogamma_a_ab_device = task.v2rhogamma_a_ab; + const auto* v2rhogamma_a_bb_device = task.v2rhogamma_a_bb; + const auto* v2rhogamma_b_aa_device = task.v2rhogamma_b_aa; + const auto* v2rhogamma_b_ab_device = task.v2rhogamma_b_ab; + const auto* v2rhogamma_b_bb_device = task.v2rhogamma_b_bb; + const auto* v2gamma2_aa_aa_device = task.v2gamma2_aa_aa; + const auto* v2gamma2_aa_ab_device = task.v2gamma2_aa_ab; + const auto* v2gamma2_aa_bb_device = task.v2gamma2_aa_bb; + const auto* v2gamma2_ab_ab_device = task.v2gamma2_ab_ab; + const auto* v2gamma2_ab_bb_device = task.v2gamma2_ab_bb; + const auto* v2gamma2_bb_bb_device = task.v2gamma2_bb_bb; + + auto* FXC_A_s_device = task.FXC_A_s; + auto* FXC_A_z_device = task.FXC_A_z; + auto* FXC_Bx_s_device = task.FXC_Bx_s; + auto* FXC_Bx_z_device = task.FXC_Bx_z; + auto* FXC_By_s_device = task.FXC_By_s; + auto* FXC_By_z_device = task.FXC_By_z; + auto* FXC_Bz_s_device = task.FXC_Bz_s; + auto* FXC_Bz_z_device = task.FXC_Bz_z; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const auto ps = tden_s_device[ tid ]; + const auto pz = tden_z_device[ tid ]; + const auto trho_a_device = 0.5*(ps + pz); + const auto trho_b_device = 0.5*(ps - pz); + + const auto tdndx = tden_pos_x_eval_device[ tid ]; + const auto tdndy = tden_pos_y_eval_device[ tid ]; + const auto tdndz = tden_pos_z_eval_device[ tid ]; + const auto tdMzdx = tden_neg_x_eval_device[ tid ]; + const auto tdMzdy = tden_neg_y_eval_device[ tid ]; + const auto tdMzdz = tden_neg_z_eval_device[ tid ]; + const auto tdden_a_x = 0.5*(tdndx + tdMzdx); + const auto tdden_a_y = 0.5*(tdndy + tdMzdy); + const auto tdden_a_z = 0.5*(tdndz + tdMzdz); + const auto tdden_b_x = 0.5*(tdndx - tdMzdx); + const auto tdden_b_y = 0.5*(tdndy - tdMzdy); + const auto tdden_b_z = 0.5*(tdndz - tdMzdz); + + const auto dndx = den_pos_x_eval_device[ tid ]; + const auto dndy = den_pos_y_eval_device[ tid ]; + const auto dndz = den_pos_z_eval_device[ tid ]; + const auto dMzdx = den_neg_x_eval_device[ tid ]; + const auto dMzdy = den_neg_y_eval_device[ tid ]; + const auto dMzdz = den_neg_z_eval_device[ tid ]; + const auto dden_a_x = 0.5*(dndx + dMzdx); + const auto dden_a_y = 0.5*(dndy + dMzdy); + const auto dden_a_z = 0.5*(dndz + dMzdz); + const auto dden_b_x = 0.5*(dndx - dMzdx); + const auto dden_b_y = 0.5*(dndy - dMzdy); + const auto dden_b_z = 0.5*(dndz - dMzdz); + + const auto tgamma_pp = tdden_a_x * dden_a_x + tdden_a_y * dden_a_y + tdden_a_z * dden_a_z; + const auto tgamma_pm = tdden_a_x * dden_b_x + tdden_a_y * dden_b_y + tdden_a_z * dden_b_z + + tdden_b_x * dden_a_x + tdden_b_y * dden_a_y + tdden_b_z * dden_a_z; + const auto tgamma_mm = tdden_b_x * dden_b_x + tdden_b_y * dden_b_y + tdden_b_z * dden_b_z; + + + const auto A_a = v2rho2_a_a_device[tid] * trho_a_device + 2.0 * v2rhogamma_a_aa_device[tid] * tgamma_pp + + v2rhogamma_a_ab_device[tid] * tgamma_pm + 2.0 * v2rhogamma_a_bb_device[tid] * tgamma_mm + + v2rho2_a_b_device[tid] * trho_b_device; + const auto A_b = v2rho2_b_b_device[tid] * trho_b_device + 2.0 * v2rhogamma_b_bb_device[tid] * tgamma_mm + + v2rhogamma_b_ab_device[tid] * tgamma_pm + 2.0 * v2rhogamma_b_aa_device[tid] * tgamma_pp + + v2rho2_a_b_device[tid] * trho_a_device; + FXC_A_s_device[ tid ] = 0.5 * weight_device[ tid ] * (A_a + A_b); + FXC_A_z_device[ tid ] = 0.5 * weight_device[ tid ] * (A_a - A_b); + // Calculate B coefficients for alpha spin + const double B_coef1_a = v2rhogamma_a_aa_device[tid] * trho_a_device + 2.0 * v2gamma2_aa_aa_device[tid] * tgamma_pp + + v2gamma2_aa_ab_device[tid] * tgamma_pm + 2.0 * v2gamma2_aa_bb_device[tid] * tgamma_mm + + v2rhogamma_b_aa_device[tid] * trho_b_device; + + const double B_coef2_a = v2rhogamma_a_ab_device[tid] * trho_a_device + 2.0 * v2gamma2_aa_ab_device[tid] * tgamma_pp + + v2gamma2_ab_ab_device[tid] * tgamma_pm + 2.0 * v2gamma2_ab_bb_device[tid] * tgamma_mm + + v2rhogamma_b_ab_device[tid] * trho_b_device; + + // Calculate gradient components for alpha spin + const double Bx_a = 2.0 * B_coef1_a * dden_a_x + B_coef2_a * dden_b_x + + 2.0 * vgamma_aa_device[tid] * tdden_a_x + vgamma_ab_device[tid] * tdden_b_x; + + const double By_a = 2.0 * B_coef1_a * dden_a_y + B_coef2_a * dden_b_y + + 2.0 * vgamma_aa_device[tid] * tdden_a_y + vgamma_ab_device[tid] * tdden_b_y; + + const double Bz_a = 2.0 * B_coef1_a * dden_a_z + B_coef2_a * dden_b_z + + 2.0 * vgamma_aa_device[tid] * tdden_a_z + vgamma_ab_device[tid] * tdden_b_z; + + // Calculate B coefficients for beta spin + const double B_coef1_b = v2rhogamma_b_bb_device[tid] * trho_b_device + 2.0 * v2gamma2_bb_bb_device[tid] * tgamma_mm + + v2gamma2_ab_bb_device[tid] * tgamma_pm + 2.0 * v2gamma2_aa_bb_device[tid] * tgamma_pp + + v2rhogamma_a_bb_device[tid] * trho_a_device; + + const double B_coef2_b = v2rhogamma_b_ab_device[tid] * trho_b_device + 2.0 * v2gamma2_ab_bb_device[tid] * tgamma_mm + + v2gamma2_ab_ab_device[tid] * tgamma_pm + 2.0 * v2gamma2_aa_ab_device[tid] * tgamma_pp + + v2rhogamma_a_ab_device[tid] * trho_a_device; + + const double Bx_b = 2.0 * B_coef1_b * dden_b_x + B_coef2_b * dden_a_x + + 2.0 * vgamma_bb_device[tid] * tdden_b_x + vgamma_ab_device[tid] * tdden_a_x; + + const double By_b = 2.0 * B_coef1_b * dden_b_y + B_coef2_b * dden_a_y + + 2.0 * vgamma_bb_device[tid] * tdden_b_y + vgamma_ab_device[tid] * tdden_a_y; + + const double Bz_b = 2.0 * B_coef1_b * dden_b_z + B_coef2_b * dden_a_z + + 2.0 * vgamma_bb_device[tid] * tdden_b_z + vgamma_ab_device[tid] * tdden_a_z; + + FXC_Bx_s_device[tid] = 0.5 * weight_device[tid] * (Bx_a + Bx_b); + FXC_By_s_device[tid] = 0.5 * weight_device[tid] * (By_a + By_b); + FXC_Bz_s_device[tid] = 0.5 * weight_device[tid] * (Bz_a + Bz_b); + FXC_Bx_z_device[tid] = 0.5 * weight_device[tid] * (Bx_a - Bx_b); + FXC_By_z_device[tid] = 0.5 * weight_device[tid] * (By_a - By_b); + FXC_Bz_z_device[tid] = 0.5 * weight_device[tid] * (Bz_a - Bz_b); + + + } + +} + +inline void eval_uvars_gga_gks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + auto* den_s_eval_device = task.den_s; + const auto* dden_sx_eval_device = task.dden_sx; + const auto* dden_sy_eval_device = task.dden_sy; + const auto* dden_sz_eval_device = task.dden_sz; + + auto* den_z_eval_device = task.den_z; + const auto* dden_zx_eval_device = task.dden_zx; + const auto* dden_zy_eval_device = task.dden_zy; + const auto* dden_zz_eval_device = task.dden_zz; + + const auto* den_y_eval_device = task.den_y; + const auto* dden_yx_eval_device = task.dden_yx; + const auto* dden_yy_eval_device = task.dden_yy; + const auto* dden_yz_eval_device = task.dden_yz; + + const auto* den_x_eval_device = task.den_x; + const auto* dden_xx_eval_device = task.dden_xx; + const auto* dden_xy_eval_device = task.dden_xy; + const auto* dden_xz_eval_device = task.dden_xz; + + auto* gamma_pp_eval_device = task.gamma_pp; + auto* gamma_pm_eval_device = task.gamma_pm; + auto* gamma_mm_eval_device = task.gamma_mm; + + auto* H_z_eval_device = task.H_z; + auto* H_y_eval_device = task.H_y; + auto* H_x_eval_device = task.H_x; + auto* K_z_eval_device = task.K_z; + auto* K_y_eval_device = task.K_y; + auto* K_x_eval_device = task.K_x; + + const double dtolsq = 1e-24; // TODO: make variable + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const double dndz = dden_sz_eval_device[ tid ]; + const double dndy = dden_sy_eval_device[ tid ]; + const double dndx = dden_sx_eval_device[ tid ]; + + const double dMzdz = dden_zz_eval_device[ tid ]; + const double dMzdy = dden_zy_eval_device[ tid ]; + const double dMzdx = dden_zx_eval_device[ tid ]; + + const double dMydz = dden_yz_eval_device[ tid ]; + const double dMydy = dden_yy_eval_device[ tid ]; + const double dMydx = dden_yx_eval_device[ tid ]; + + const double dMxdz = dden_xz_eval_device[ tid ]; + const double dMxdy = dden_xy_eval_device[ tid ]; + const double dMxdx = dden_xx_eval_device[ tid ]; + + const auto ps = den_s_eval_device[ tid ]; + const auto pz = den_z_eval_device[ tid ]; + const auto py = den_y_eval_device[ tid ]; + const auto px = den_x_eval_device[ tid ]; + + const auto mtemp = pz*pz + px*px + py*py; + double mnorm = 0.; + + const auto dels_dot_dels = dndx * dndx + dndy * dndy + dndz * dndz; + const auto delz_dot_delz = dMzdx * dMzdx + dMzdy * dMzdy + dMzdz * dMzdz; + const auto delx_dot_delx = dMxdx * dMxdx + dMxdy * dMxdy + dMxdz * dMxdz; + const auto dely_dot_dely = dMydx * dMydx + dMydy * dMydy + dMydz * dMydz; + + const auto dels_dot_delz = dndx * dMzdx + dndy * dMzdy + dndz * dMzdz; + const auto dels_dot_delx = dndx * dMxdx + dndy * dMxdy + dndz * dMxdz; + const auto dels_dot_dely = dndx * dMydx + dndy * dMydy + dndz * dMydz; + + const auto sum = delz_dot_delz + delx_dot_delx + dely_dot_dely; + const auto s_sum = + dels_dot_delz * pz + dels_dot_delx * px + dels_dot_dely * py; + + const auto inv_sqsum2 = + ::sycl::rsqrt(dels_dot_delz * dels_dot_delz + dels_dot_delx * dels_dot_delx + + dels_dot_dely * dels_dot_dely); + const auto sqsum2 = 1./inv_sqsum2; + + double sign = 1.; + if( ::sycl::signbit(s_sum)) + sign = -1.; + + + if (mtemp > dtolsq) { + const double inv_mnorm = ::sycl::rsqrt(mtemp); + mnorm = 1./inv_mnorm; + K_z_eval_device[ tid ] = pz * inv_mnorm; + K_y_eval_device[ tid ] = py * inv_mnorm; + K_x_eval_device[ tid ] = px * inv_mnorm; + H_z_eval_device[ tid ] = sign * dels_dot_delz * inv_sqsum2; + H_y_eval_device[ tid ] = sign * dels_dot_dely * inv_sqsum2; + H_x_eval_device[ tid ] = sign * dels_dot_delx * inv_sqsum2; + } + else { + mnorm = (1. / 3.) * (px + py + pz); + K_z_eval_device[ tid ] = 1. / 3.; + K_y_eval_device[ tid ] = 1. / 3.; + K_x_eval_device[ tid ] = 1. / 3.; + + H_z_eval_device[ tid ] = sign / 3.; + H_y_eval_device[ tid ] = sign / 3.; + H_x_eval_device[ tid ] = sign / 3.; + } + + gamma_pp_eval_device[ tid ] = 0.25*(dels_dot_dels + sum) + 0.5*sign*sqsum2; + gamma_pm_eval_device[ tid ] = 0.25*(dels_dot_dels - sum); + gamma_mm_eval_device[ tid ] = 0.25*(dels_dot_dels + sum) - 0.5*sign*sqsum2; + + den_s_eval_device[ tid ] = 0.5*(ps + mnorm); + den_z_eval_device[ tid ] = 0.5*(ps - mnorm); + + } + +} + +} // namespace GauXC diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_lda.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_lda.hpp new file mode 100644 index 000000000..69401b018 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_lda.hpp @@ -0,0 +1,215 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_util.hpp" +#include "device/xc_device_data.hpp" +#include "sycl_extensions.hpp" +#include "sycl_atomics.hpp" +#include "sycl_launch.hpp" + +namespace GauXC { + +template +void eval_vvar_lda_kern( size_t ntasks, + XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + + double* den_eval_device = nullptr; + // use the "U" variable (+/- for UKS) even though at this point the density (S/Z) is stored + if constexpr (trial){ + if constexpr (den_select == DEN_S) den_eval_device = task.tden_s; + if constexpr (den_select == DEN_Z) den_eval_device = task.tden_z; + if constexpr (den_select == DEN_Y) den_eval_device = task.tden_y; + if constexpr (den_select == DEN_X) den_eval_device = task.tden_x; + }else{ + if constexpr (den_select == DEN_S) den_eval_device = task.den_s; + if constexpr (den_select == DEN_Z) den_eval_device = task.den_z; + if constexpr (den_select == DEN_Y) den_eval_device = task.den_y; + if constexpr (den_select == DEN_X) den_eval_device = task.den_x; + } + + const auto* basis_eval_device = task.bf; + + const auto* den_basis_prod_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + double den_reg = 0.; + + if( tid_x < nbf and tid_y < npts ) { + + const double* bf_col = basis_eval_device + tid_x*npts; + const double* db_col = den_basis_prod_device + tid_x*npts; + + den_reg = bf_col[ tid_y ] * db_col[ tid_y ]; + + } + + // Warp blocks are stored col major + constexpr auto warp_size = sycl::warp_size; + den_reg = sycl::warp_reduce_sum( den_reg ); + + + if( it.get_local_id(2) == 0 and tid_y < npts ) { + atomic_add_device( den_eval_device + tid_y, den_reg ); + } + +} + +inline void eval_uvars_lda_rks_kernel(size_t ntasks, XCDeviceTask* tasks_device) { + // eval_vvars populated uvar storage already in the case of LDA+RKS + return; +} +inline void eval_tmat_lda_rks_kernel(size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + const auto* v2rho2_device = task.v2rho2; + const auto* weight_device = task.weights; + auto* tden_s_eval_device = task.tden_s; + auto* FXC_A_device = task.FXC_A_s; + + const int tid = it.get_global_id(2); + if( tid < npts ) { + FXC_A_device[ tid ] = v2rho2_device[ tid ] * tden_s_eval_device[ tid ] * weight_device[ tid ]; + } + + return; +} + + +inline void eval_uvars_lda_uks_kernel(size_t ntasks, + XCDeviceTask* tasks_device ) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + + const auto npts = task.npts; + + auto* den_pos_eval_device = task.den_s; + auto* den_neg_eval_device = task.den_z; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const auto ps = den_pos_eval_device[ tid ]; + const auto pz = den_neg_eval_device[ tid ]; + den_pos_eval_device[ tid ] = 0.5*(ps + pz); + den_neg_eval_device[ tid ] = 0.5*(ps - pz); + } +} + +inline void eval_tmat_lda_uks_kernel(size_t ntasks, + XCDeviceTask* tasks_device ) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + + const auto npts = task.npts; + + auto* tden_s_device = task.tden_s; + auto* tden_z_device = task.tden_z; + auto* FXC_A_s_device = task.FXC_A_s; + auto* FXC_A_z_device = task.FXC_A_z; + const auto* weight_device = task.weights; + + const auto* v2rho2_a_a_device = task.v2rho2_a_a; + const auto* v2rho2_a_b_device = task.v2rho2_a_b; + const auto* v2rho2_b_b_device = task.v2rho2_b_b; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const auto ps = tden_s_device[ tid ]; + const auto pz = tden_z_device[ tid ]; + const auto trho_a_device = 0.5*(ps + pz); + const auto trho_b_device = 0.5*(ps - pz); + const auto A_a = v2rho2_a_a_device[tid] * trho_a_device + v2rho2_a_b_device[tid] * trho_b_device; + const auto A_b = v2rho2_b_b_device[tid] * trho_b_device + v2rho2_a_b_device[tid] * trho_a_device; + FXC_A_s_device[ tid ] = 0.5 * weight_device[ tid ] * (A_a + A_b); + FXC_A_z_device[ tid ] = 0.5 * weight_device[ tid ] * (A_a - A_b); + } +} + +inline void eval_uvars_lda_gks_kernel(size_t ntasks, + XCDeviceTask* tasks_device ) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + + const auto npts = task.npts; + + auto* den_z_eval_device = task.den_s; + auto* den_s_eval_device = task.den_z; + auto* den_y_eval_device = task.den_y; + auto* den_x_eval_device = task.den_x; + auto* K_z_eval_device = task.K_z; + auto* K_y_eval_device = task.K_y; + auto* K_x_eval_device = task.K_x; + const double dtolsq = 1e-24; // TODO: make variable + + const int tid = it.get_global_id(2); + + + if( tid < npts ) { + const auto ps = den_s_eval_device[ tid ]; + const auto pz = den_z_eval_device[ tid ]; + const auto py = den_y_eval_device[ tid ]; + const auto px = den_x_eval_device[ tid ]; + const auto mtemp = pz*pz + px*px + py*py; + double mnorm = 0.; + + if (mtemp > dtolsq) { + const double inv_mnorm = ::sycl::rsqrt(mtemp); + mnorm = 1./inv_mnorm; + K_z_eval_device[ tid ] = pz * inv_mnorm; + K_y_eval_device[ tid ] = py * inv_mnorm; + K_x_eval_device[ tid ] = px * inv_mnorm; + } + else { + mnorm = (1. / 3.) * (px + py + pz); + K_z_eval_device[ tid ] = 1. / 3.; + K_y_eval_device[ tid ] = 1. / 3.; + K_x_eval_device[ tid ] = 1. / 3.; + } + + den_s_eval_device[ tid ] = 0.5*(ps + mnorm); + den_z_eval_device[ tid ] = 0.5*(ps - mnorm); + + } +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_mgga.hpp b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_mgga.hpp new file mode 100644 index 000000000..895547fd4 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/uvvars_mgga.hpp @@ -0,0 +1,465 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "device_specific/sycl_device_constants.hpp" +#include "device_specific/sycl_util.hpp" +#include "device/xc_device_data.hpp" +#include "sycl_extensions.hpp" +#include "sycl_atomics.hpp" +#include "sycl_launch.hpp" + +#define MGGA_KERNEL_SM_BLOCK 32 + +namespace GauXC { + + + +template +void eval_vvar_mgga_kern( size_t ntasks, + XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + constexpr auto warp_size = sycl::warp_size; + auto& den_shared = GauXC::sycl::local_mem(); + + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + double* tau_eval_device = nullptr; + double* lapl_eval_device = nullptr; + + if constexpr (trial){ + if constexpr (den_select == DEN_S) { + tau_eval_device = task.ttau_s; + if constexpr (need_lapl) { + lapl_eval_device = task.tlapl_s; + } + } + if constexpr (den_select == DEN_Z) { + tau_eval_device = task.ttau_z; + if constexpr (need_lapl) { + lapl_eval_device = task.tlapl_z; + } + } + } else{ + if constexpr (den_select == DEN_S) { + tau_eval_device = task.tau_s; + if constexpr (need_lapl) { + lapl_eval_device = task.lapl_s; + } + } + if constexpr (den_select == DEN_Z) { + tau_eval_device = task.tau_z; + if constexpr (need_lapl) { + lapl_eval_device = task.lapl_z; + } + } + } + + //const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + decltype(dbasis_x_eval_device) basis_lapl_eval_device = nullptr; + if constexpr (need_lapl) { + basis_lapl_eval_device = task.d2bflapl; + } + + //const auto* den_basis_prod_device = task.zmat; + const auto* den_basis_dx_prod_device = task.xmat_x; + const auto* den_basis_dy_prod_device = task.xmat_y; + const auto* den_basis_dz_prod_device = task.xmat_z; + decltype(den_basis_dx_prod_device) den_basis_prod_device = nullptr; + if constexpr (need_lapl) { + den_basis_prod_device = task.zmat; + } + + const int threadIdx_x = it.get_local_id(2); + const int threadIdx_y = it.get_local_id(1); + const int blockDim_x = it.get_local_range(2); + const int blockDim_y = it.get_local_range(1); + + for ( int bid_x = it.get_group(2) * blockDim_x; + bid_x < nbf; + bid_x += blockDim_x * it.get_group_range(2) ) { + + for ( int bid_y = it.get_group(1) * MGGA_KERNEL_SM_BLOCK; + bid_y < npts; + bid_y += MGGA_KERNEL_SM_BLOCK * it.get_group_range(1) ) { + + for (int sm_y = threadIdx_y; sm_y < MGGA_KERNEL_SM_BLOCK; sm_y += blockDim_y) { + den_shared[0][threadIdx_x][sm_y] = 0.; + den_shared[1][threadIdx_x][sm_y] = 0.; + den_shared[2][threadIdx_x][sm_y] = 0.; + if constexpr (need_lapl) + den_shared[3][threadIdx_x][sm_y] = 0.; + + if (bid_y + threadIdx_x < npts and bid_x + sm_y < nbf) { + const double* db_x_col = den_basis_dx_prod_device + (bid_x + sm_y)*npts; + const double* db_y_col = den_basis_dy_prod_device + (bid_x + sm_y)*npts; + const double* db_z_col = den_basis_dz_prod_device + (bid_x + sm_y)*npts; + + const double* bf_x_col = dbasis_x_eval_device + (bid_x + sm_y)*npts; + const double* bf_y_col = dbasis_y_eval_device + (bid_x + sm_y)*npts; + const double* bf_z_col = dbasis_z_eval_device + (bid_x + sm_y)*npts; + + + den_shared[0][threadIdx_x][sm_y] = bf_x_col[ bid_y + threadIdx_x ] * db_x_col[ bid_y + threadIdx_x ]; + den_shared[1][threadIdx_x][sm_y] = bf_y_col[ bid_y + threadIdx_x ] * db_y_col[ bid_y + threadIdx_x ]; + den_shared[2][threadIdx_x][sm_y] = bf_z_col[ bid_y + threadIdx_x ] * db_z_col[ bid_y + threadIdx_x ]; + + + if constexpr (need_lapl) { + const double* db_col = den_basis_prod_device + (bid_x + sm_y)*npts; + const double* bf_l_col = basis_lapl_eval_device + (bid_x + sm_y)*npts; + den_shared[3][threadIdx_x][sm_y] = bf_l_col[ bid_y + threadIdx_x ] * db_col[ bid_y + threadIdx_x ]; + } + } + } + ::sycl::group_barrier( it.get_group() ); + + + for (int sm_y = threadIdx_y; sm_y < MGGA_KERNEL_SM_BLOCK; sm_y += blockDim_y) { + const int tid_y = bid_y + sm_y; + + double tx_reg = den_shared[0][sm_y][threadIdx_x]; + double ty_reg = den_shared[1][sm_y][threadIdx_x]; + double tz_reg = den_shared[2][sm_y][threadIdx_x]; + // Warp blocks are stored col major + double tau_reg = 0.0; + tau_reg = 0.5 * sycl::warp_reduce_sum( tx_reg ); + tau_reg += 0.5 * sycl::warp_reduce_sum( ty_reg ); + tau_reg += 0.5 * sycl::warp_reduce_sum( tz_reg ); + + double lapl_reg = 0.0; + if constexpr (need_lapl) { + lapl_reg = den_shared[3][sm_y][threadIdx_x]; + lapl_reg = sycl::warp_reduce_sum( lapl_reg ); + lapl_reg = 2. * lapl_reg + 4. * tau_reg; + } + + if( threadIdx_x == 0 and tid_y < npts ) { + atomic_add_device( tau_eval_device + tid_y, tau_reg ); + if constexpr (need_lapl) { + atomic_add_device( lapl_eval_device + tid_y, lapl_reg ); + } + } + } + ::sycl::group_barrier( it.get_group() ); + } + } +} + + + + +template +void eval_uvars_mgga_uks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + auto* tau_pos_eval_device = task.tau_s; + auto* tau_neg_eval_device = task.tau_z; + + double* lapl_pos_eval_device = nullptr; + double* lapl_neg_eval_device = nullptr; + if constexpr (need_lapl) { + lapl_pos_eval_device = task.lapl_s; + lapl_neg_eval_device = task.lapl_z; + } + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const double ts = tau_pos_eval_device[ tid ]; + const double tz = tau_neg_eval_device[ tid ]; + tau_pos_eval_device[ tid ] = 0.5*(ts + tz); + tau_neg_eval_device[ tid ] = 0.5*(ts - tz); + + if constexpr (need_lapl) { + const double ls = lapl_pos_eval_device[ tid ]; + const double lz = lapl_neg_eval_device[ tid ]; + lapl_pos_eval_device[ tid ] = 0.5*(ls + lz); + lapl_neg_eval_device[ tid ] = 0.5*(ls - lz); + } + } + +} + + +inline void eval_tmat_mgga_rks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + const auto* dden_sx_eval_device = task.dden_sx; + const auto* dden_sy_eval_device = task.dden_sy; + const auto* dden_sz_eval_device = task.dden_sz; + const auto* tdden_sx_eval_device = task.tdden_sx; + const auto* tdden_sy_eval_device = task.tdden_sy; + const auto* tdden_sz_eval_device = task.tdden_sz; + + const auto* weight_device = task.weights; + const auto* vgamma_device = task.vgamma; + const auto* v2rho2_device = task.v2rho2; + const auto* v2rhogamma_device = task.v2rhogamma; + const auto* v2gamma2_device = task.v2gamma2; + const auto* v2rhotau_device = task.v2rhotau; + const auto* v2tau2_device = task.v2tau2; + const auto* v2gammatau_device = task.v2gammatau; + const auto* trho_device = task.tden_s; + const auto* ttau_device = task.ttau_s; + + auto* FXC_A_device = task.FXC_A_s; + auto* FXC_Bx_device = task.FXC_Bx_s; + auto* FXC_By_device = task.FXC_By_s; + auto* FXC_Bz_device = task.FXC_Bz_s; + auto* FXC_C_device = task.FXC_C_s; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const auto dx = dden_sx_eval_device[ tid ]; + const auto dy = dden_sy_eval_device[ tid ]; + const auto dz = dden_sz_eval_device[ tid ]; + const auto tdx = tdden_sx_eval_device[ tid ]; + const auto tdy = tdden_sy_eval_device[ tid ]; + const auto tdz = tdden_sz_eval_device[ tid ]; + const auto tgamma = tdx*dx + tdy*dy + tdz*dz; + + const auto FXC_A = v2rho2_device[ tid ] * trho_device[ tid ] + 2.0 * v2rhogamma_device[tid] * tgamma + + v2rhotau_device[ tid ] * ttau_device[ tid ]; + FXC_A_device[ tid ] = weight_device[ tid ] * FXC_A; + + const auto FXC_C = v2rhotau_device[ tid ] * trho_device[ tid ] + 2.0 * v2gammatau_device[ tid ] * tgamma + + v2tau2_device[ tid ] * ttau_device[ tid ]; + FXC_C_device[ tid ] = weight_device[ tid ] * FXC_C; + + const auto B_coef = v2rhogamma_device[tid] * trho_device[tid] + 2.0 * v2gamma2_device[tid] * tgamma + + v2gammatau_device[ tid ] * ttau_device[ tid ]; + FXC_Bx_device[ tid ] = 2.0 * weight_device[ tid ] * ( B_coef * dx + vgamma_device[ tid ] * tdx ); + FXC_By_device[ tid ] = 2.0 * weight_device[ tid ] * ( B_coef * dy + vgamma_device[ tid ] * tdy ); + FXC_Bz_device[ tid ] = 2.0 * weight_device[ tid ] * ( B_coef * dz + vgamma_device[ tid ] * tdz ); + } + +} + + + +inline void eval_tmat_mgga_uks_kernel( size_t ntasks, XCDeviceTask* tasks_device) { + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + const auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + + const auto* tden_s_device = task.tden_s; + const auto* tden_z_device = task.tden_z; + const auto* ttau_s_device = task.ttau_s; + const auto* ttau_z_device = task.ttau_z; + const auto* weight_device = task.weights; + + const auto* tden_pos_x_eval_device = task.tdden_sx; + const auto* tden_pos_y_eval_device = task.tdden_sy; + const auto* tden_pos_z_eval_device = task.tdden_sz; + const auto* den_pos_x_eval_device = task.dden_sx; + const auto* den_pos_y_eval_device = task.dden_sy; + const auto* den_pos_z_eval_device = task.dden_sz; + + const auto* tden_neg_x_eval_device = task.tdden_zx; + const auto* tden_neg_y_eval_device = task.tdden_zy; + const auto* tden_neg_z_eval_device = task.tdden_zz; + const auto* den_neg_x_eval_device = task.dden_zx; + const auto* den_neg_y_eval_device = task.dden_zy; + const auto* den_neg_z_eval_device = task.dden_zz; + + const double* vgamma_aa_device = task.vgamma_pp; + const double* vgamma_ab_device = task.vgamma_pm; + const double* vgamma_bb_device = task.vgamma_mm; + const double* v2rho2_a_a_device = task.v2rho2_a_a; + const double* v2rho2_a_b_device = task.v2rho2_a_b; + const double* v2rho2_b_b_device = task.v2rho2_b_b; + const double* v2rhogamma_a_aa_device = task.v2rhogamma_a_aa; + const double* v2rhogamma_a_ab_device = task.v2rhogamma_a_ab; + const double* v2rhogamma_a_bb_device = task.v2rhogamma_a_bb; + const double* v2rhogamma_b_aa_device = task.v2rhogamma_b_aa; + const double* v2rhogamma_b_ab_device = task.v2rhogamma_b_ab; + const double* v2rhogamma_b_bb_device = task.v2rhogamma_b_bb; + const double* v2gamma2_aa_aa_device = task.v2gamma2_aa_aa; + const double* v2gamma2_aa_ab_device = task.v2gamma2_aa_ab; + const double* v2gamma2_aa_bb_device = task.v2gamma2_aa_bb; + const double* v2gamma2_ab_ab_device = task.v2gamma2_ab_ab; + const double* v2gamma2_ab_bb_device = task.v2gamma2_ab_bb; + const double* v2gamma2_bb_bb_device = task.v2gamma2_bb_bb; + const double* v2rhotau_a_a_device = task.v2rhotau_a_a; + const double* v2rhotau_a_b_device = task.v2rhotau_a_b; + const double* v2rhotau_b_a_device = task.v2rhotau_b_a; + const double* v2rhotau_b_b_device = task.v2rhotau_b_b; + const double* v2gammatau_aa_a_device= task.v2gammatau_aa_a; + const double* v2gammatau_aa_b_device= task.v2gammatau_aa_b; + const double* v2gammatau_ab_a_device= task.v2gammatau_ab_a; + const double* v2gammatau_ab_b_device= task.v2gammatau_ab_b; + const double* v2gammatau_bb_a_device= task.v2gammatau_bb_a; + const double* v2gammatau_bb_b_device= task.v2gammatau_bb_b; + const double* v2tau2_a_a_device = task.v2tau2_a_a; + const double* v2tau2_a_b_device = task.v2tau2_a_b; + const double* v2tau2_b_b_device = task.v2tau2_b_b; + + auto* FXC_A_s_device = task.FXC_A_s; + auto* FXC_A_z_device = task.FXC_A_z; + auto* FXC_Bx_s_device = task.FXC_Bx_s; + auto* FXC_Bx_z_device = task.FXC_Bx_z; + auto* FXC_By_s_device = task.FXC_By_s; + auto* FXC_By_z_device = task.FXC_By_z; + auto* FXC_Bz_s_device = task.FXC_Bz_s; + auto* FXC_Bz_z_device = task.FXC_Bz_z; + auto* FXC_C_s_device = task.FXC_C_s; + auto* FXC_C_z_device = task.FXC_C_z; + + const int tid = it.get_global_id(2); + + if( tid < npts ) { + const auto ps = tden_s_device[ tid ]; + const auto pz = tden_z_device[ tid ]; + const auto trho_a_device = 0.5*(ps + pz); + const auto trho_b_device = 0.5*(ps - pz); + const auto ts = ttau_s_device[ tid ]; + const auto tz = ttau_z_device[ tid ]; + const auto tau_a = 0.5*(ts + tz); + const auto tau_b = 0.5*(ts - tz); + + const auto tdndx = tden_pos_x_eval_device[ tid ]; + const auto tdndy = tden_pos_y_eval_device[ tid ]; + const auto tdndz = tden_pos_z_eval_device[ tid ]; + const auto tdMzdx = tden_neg_x_eval_device[ tid ]; + const auto tdMzdy = tden_neg_y_eval_device[ tid ]; + const auto tdMzdz = tden_neg_z_eval_device[ tid ]; + const auto tdden_a_x = 0.5*(tdndx + tdMzdx); + const auto tdden_a_y = 0.5*(tdndy + tdMzdy); + const auto tdden_a_z = 0.5*(tdndz + tdMzdz); + const auto tdden_b_x = 0.5*(tdndx - tdMzdx); + const auto tdden_b_y = 0.5*(tdndy - tdMzdy); + const auto tdden_b_z = 0.5*(tdndz - tdMzdz); + + const auto dndx = den_pos_x_eval_device[ tid ]; + const auto dndy = den_pos_y_eval_device[ tid ]; + const auto dndz = den_pos_z_eval_device[ tid ]; + const auto dMzdx = den_neg_x_eval_device[ tid ]; + const auto dMzdy = den_neg_y_eval_device[ tid ]; + const auto dMzdz = den_neg_z_eval_device[ tid ]; + const auto dden_a_x = 0.5*(dndx + dMzdx); + const auto dden_a_y = 0.5*(dndy + dMzdy); + const auto dden_a_z = 0.5*(dndz + dMzdz); + const auto dden_b_x = 0.5*(dndx - dMzdx); + const auto dden_b_y = 0.5*(dndy - dMzdy); + const auto dden_b_z = 0.5*(dndz - dMzdz); + + const auto tgamma_pp = tdden_a_x * dden_a_x + tdden_a_y * dden_a_y + tdden_a_z * dden_a_z; + const auto tgamma_pm = tdden_a_x * dden_b_x + tdden_a_y * dden_b_y + tdden_a_z * dden_b_z + + tdden_b_x * dden_a_x + tdden_b_y * dden_a_y + tdden_b_z * dden_a_z; + const auto tgamma_mm = tdden_b_x * dden_b_x + tdden_b_y * dden_b_y + tdden_b_z * dden_b_z; + + + const auto A_a = v2rho2_a_a_device[tid] * trho_a_device + 2.0 * v2rhogamma_a_aa_device[tid] * tgamma_pp + + v2rhogamma_a_ab_device[tid] * tgamma_pm + 2.0 * v2rhogamma_a_bb_device[tid] * tgamma_mm + + v2rho2_a_b_device[tid] * trho_b_device + v2rhotau_a_a_device[tid] * tau_a + + v2rhotau_a_b_device[tid] * tau_b; + const auto A_b = v2rho2_b_b_device[tid] * trho_b_device + 2.0 * v2rhogamma_b_bb_device[tid] * tgamma_mm + + v2rhogamma_b_ab_device[tid] * tgamma_pm + 2.0 * v2rhogamma_b_aa_device[tid] * tgamma_pp + + v2rho2_a_b_device[tid] * trho_a_device + v2rhotau_b_b_device[tid] * tau_b + + v2rhotau_b_a_device[tid] * tau_a; + FXC_A_s_device[ tid ] = 0.5 * weight_device[ tid ] * (A_a + A_b); + FXC_A_z_device[ tid ] = 0.5 * weight_device[ tid ] * (A_a - A_b); + + // Compute C coefficients for alpha and beta spin + const auto C_a = v2rhotau_a_a_device[tid] * trho_a_device + v2rhotau_b_a_device[tid] * trho_b_device + + 2.0 * v2gammatau_aa_a_device[tid] * tgamma_pp + v2gammatau_ab_a_device[tid] * tgamma_pm + + 2.0 * v2gammatau_bb_a_device[tid] * tgamma_mm + + v2tau2_a_a_device[tid] * tau_a + v2tau2_a_b_device[tid] * tau_b; + + const auto C_b = v2rhotau_a_b_device[tid] * trho_a_device + v2rhotau_b_b_device[tid] * trho_b_device + + 2.0 * v2gammatau_aa_b_device[tid] * tgamma_pp + v2gammatau_ab_b_device[tid] * tgamma_pm + + 2.0 * v2gammatau_bb_b_device[tid] * tgamma_mm + + v2tau2_a_b_device[tid] * tau_a + v2tau2_b_b_device[tid] * tau_b; + + FXC_C_s_device[tid] = 0.5 * weight_device[tid] * (C_a + C_b); + FXC_C_z_device[tid] = 0.5 * weight_device[tid] * (C_a - C_b); + + // Calculate B coefficients for alpha spin + const double B_coef1_a = v2rhogamma_a_aa_device[tid] * trho_a_device + 2.0 * v2gamma2_aa_aa_device[tid] * tgamma_pp + + v2gamma2_aa_ab_device[tid] * tgamma_pm + 2.0 * v2gamma2_aa_bb_device[tid] * tgamma_mm + + v2rhogamma_b_aa_device[tid] * trho_b_device + v2gammatau_aa_a_device[tid] * tau_a + + v2gammatau_aa_b_device[tid] * tau_b; + + const double B_coef2_a = v2rhogamma_a_ab_device[tid] * trho_a_device + 2.0 * v2gamma2_aa_ab_device[tid] * tgamma_pp + + v2gamma2_ab_ab_device[tid] * tgamma_pm + 2.0 * v2gamma2_ab_bb_device[tid] * tgamma_mm + + v2rhogamma_b_ab_device[tid] * trho_b_device + v2gammatau_ab_a_device[tid] * tau_a + + v2gammatau_ab_b_device[tid] * tau_b; + + // Calculate gradient components for alpha spin + const double Bx_a = 2.0 * B_coef1_a * dden_a_x + B_coef2_a * dden_b_x + + 2.0 * vgamma_aa_device[tid] * tdden_a_x + vgamma_ab_device[tid] * tdden_b_x; + + const double By_a = 2.0 * B_coef1_a * dden_a_y + B_coef2_a * dden_b_y + + 2.0 * vgamma_aa_device[tid] * tdden_a_y + vgamma_ab_device[tid] * tdden_b_y; + + const double Bz_a = 2.0 * B_coef1_a * dden_a_z + B_coef2_a * dden_b_z + + 2.0 * vgamma_aa_device[tid] * tdden_a_z + vgamma_ab_device[tid] * tdden_b_z; + + // Calculate B coefficients for beta spin + const double B_coef1_b = v2rhogamma_b_bb_device[tid] * trho_b_device + 2.0 * v2gamma2_bb_bb_device[tid] * tgamma_mm + + v2gamma2_ab_bb_device[tid] * tgamma_pm + 2.0 * v2gamma2_aa_bb_device[tid] * tgamma_pp + + v2rhogamma_a_bb_device[tid] * trho_a_device + v2gammatau_bb_b_device[tid] * tau_b + + v2gammatau_bb_a_device[tid] * tau_a; + + const double B_coef2_b = v2rhogamma_b_ab_device[tid] * trho_b_device + 2.0 * v2gamma2_ab_bb_device[tid] * tgamma_mm + + v2gamma2_ab_ab_device[tid] * tgamma_pm + 2.0 * v2gamma2_aa_ab_device[tid] * tgamma_pp + + v2rhogamma_a_ab_device[tid] * trho_a_device + v2gammatau_ab_b_device[tid] * tau_b + + v2gammatau_ab_a_device[tid] * tau_a; + + const double Bx_b = 2.0 * B_coef1_b * dden_b_x + B_coef2_b * dden_a_x + + 2.0 * vgamma_bb_device[tid] * tdden_b_x + vgamma_ab_device[tid] * tdden_a_x; + + const double By_b = 2.0 * B_coef1_b * dden_b_y + B_coef2_b * dden_a_y + + 2.0 * vgamma_bb_device[tid] * tdden_b_y + vgamma_ab_device[tid] * tdden_a_y; + + const double Bz_b = 2.0 * B_coef1_b * dden_b_z + B_coef2_b * dden_a_z + + 2.0 * vgamma_bb_device[tid] * tdden_b_z + vgamma_ab_device[tid] * tdden_a_z; + + // Store weighted values in output arrays + FXC_Bx_s_device[tid] = 0.5 * weight_device[tid] * (Bx_a + Bx_b); + FXC_By_s_device[tid] = 0.5 * weight_device[tid] * (By_a + By_b); + FXC_Bz_s_device[tid] = 0.5 * weight_device[tid] * (Bz_a + Bz_b); + FXC_Bx_z_device[tid] = 0.5 * weight_device[tid] * (Bx_a - Bx_b); + FXC_By_z_device[tid] = 0.5 * weight_device[tid] * (By_a - By_b); + FXC_Bz_z_device[tid] = 0.5 * weight_device[tid] * (Bz_a - Bz_b); + + } + +} + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/zmat_fxc.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/zmat_fxc.cxx new file mode 100644 index 000000000..0b644a4ec --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/zmat_fxc.cxx @@ -0,0 +1,250 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/zmat_fxc.hpp" +#include +#include "device_specific/sycl_util.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "sycl_launch.hpp" + +namespace GauXC { + + +template +void zmat_lda_fxc_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const auto* FXC_A_device = task.FXC_A_s; + if constexpr ( den_selector == DEN_Z ) FXC_A_device = task.FXC_A_z; + + const auto* basis_eval_device = task.bf; + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + const double fact = 0.5 * FXC_A_device[tid_x]; + + z_matrix_device[ ibfoff ] = fact * basis_eval_device[ ibfoff ]; + } + +} + + + + + +template +void zmat_gga_fxc_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + + const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + const auto* FXC_A_device = task.FXC_A_s; + const auto* FXC_Bx_device = task.FXC_Bx_s; + const auto* FXC_By_device = task.FXC_By_s; + const auto* FXC_Bz_device = task.FXC_Bz_s; + if constexpr ( den_selector == DEN_Z ) { + FXC_A_device = task.FXC_A_z; + FXC_Bx_device = task.FXC_Bx_z; + FXC_By_device = task.FXC_By_z; + FXC_Bz_device = task.FXC_Bz_z; + } + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + + const double dx = FXC_Bx_device[tid_x] * dbasis_x_eval_device[ ibfoff ]; + const double dy = FXC_By_device[tid_x] * dbasis_y_eval_device[ ibfoff ]; + const double dz = FXC_Bz_device[tid_x] * dbasis_z_eval_device[ ibfoff ]; + + z_matrix_device[ ibfoff ] = + (0.5 * FXC_A_device[tid_x] * basis_eval_device[ ibfoff ] + dx + dy + dz ); + } +} + + + +#define ZMAT_FXC_KERN(xc_approx) \ + ::sycl::queue& stream = queue.queue_as(); \ + sycl::launch_dim threads(sycl::warp_size,sycl::max_warps_per_thread_block,1); \ + sycl::launch_dim blocks( util::div_ceil( max_npts, threads.x ), \ + util::div_ceil( max_nbf, threads.y ), \ + ntasks ); \ + if ( sel == DEN_S ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_fxc_kernel", \ + [=](){ zmat_##xc_approx##_fxc_kernel( ntasks, tasks_device ); } ); \ + else if ( sel == DEN_Z ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_fxc_kernel", \ + [=](){ zmat_##xc_approx##_fxc_kernel( ntasks, tasks_device ); } ); \ + + + +void zmat_lda_fxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + density_id sel, + device_queue queue ) { +ZMAT_FXC_KERN(lda) +} + + + +void zmat_gga_fxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + density_id sel, + device_queue queue ) { +ZMAT_FXC_KERN(gga) +} + + + +void zmat_mgga_fxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + bool do_lapl, + density_id sel, + device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + + sycl::launch_dim threads(sycl::warp_size,sycl::max_warps_per_thread_block,1); + sycl::launch_dim blocks( util::div_ceil( max_npts, threads.x ), + util::div_ceil( max_nbf, threads.y ), + ntasks ); + + if(do_lapl) + GAUXC_GENERIC_EXCEPTION("Fxc contraction + do_lapl NYI"); + + switch(sel) { + case DEN_S: + sycl::launch_kernel( stream, blocks, threads, "zmat_gga_fxc_kernel", + [=](){ zmat_gga_fxc_kernel( ntasks, tasks_device ); } ); + break; + case DEN_Z: + sycl::launch_kernel( stream, blocks, threads, "zmat_gga_fxc_kernel", + [=](){ zmat_gga_fxc_kernel( ntasks, tasks_device ); } ); + break; + } + +} + + + + + + + + + + +template +void mmat_mgga_fxc_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + auto* FXC_C_s_device = task.FXC_C_s; + if constexpr ( id == DEN_Z ) FXC_C_s_device = task.FXC_C_z; + + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + + auto* mmat_x = task.xmat_x; + auto* mmat_y = task.xmat_y; + auto* mmat_z = task.xmat_z; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + + const double fact = 0.25 * FXC_C_s_device[tid_x]; + + mmat_x[ ibfoff ] = fact * dbasis_x_eval_device[ ibfoff ]; + mmat_y[ ibfoff ] = fact * dbasis_y_eval_device[ ibfoff ]; + mmat_z[ ibfoff ] = fact * dbasis_z_eval_device[ ibfoff ]; + } +} + +void mmat_mgga_fxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + bool do_lapl, + density_id sel, + device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + + sycl::launch_dim threads(sycl::warp_size,sycl::max_warps_per_thread_block,1); + sycl::launch_dim blocks( util::div_ceil( max_npts, threads.x ), + util::div_ceil( max_nbf, threads.y ), + ntasks ); + + if(do_lapl) + GAUXC_GENERIC_EXCEPTION("Fxc contraction + do_lapl NYI"); + + switch(sel) { + case DEN_S: + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_fxc_kernel", + [=](){ mmat_mgga_fxc_kernel( ntasks, tasks_device ); } ); + break; + case DEN_Z: + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_fxc_kernel", + [=](){ mmat_mgga_fxc_kernel( ntasks, tasks_device ); } ); + break; + } + +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/kernels/zmat_vxc.cxx b/src/xc_integrator/local_work_driver/device/sycl/kernels/zmat_vxc.cxx new file mode 100644 index 000000000..d7d3d1366 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/kernels/zmat_vxc.cxx @@ -0,0 +1,870 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/zmat_vxc.hpp" +#include +#include "device_specific/sycl_util.hpp" +#include "device_specific/sycl_device_constants.hpp" +#include "sycl_launch.hpp" + +namespace GauXC { + + +void zmat_lda_vxc_rks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const auto* vrho_device = task.vrho; + + const auto* basis_eval_device = task.bf; + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + const double fact = 0.5 * vrho_device[tid_x]; + + z_matrix_device[ ibfoff ] = fact * basis_eval_device[ ibfoff ]; + + } + +} + + + + + + + + + + + + + +template +void zmat_lda_vxc_uks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const double* vrho_pos_device = task.vrho_pos; + const double* vrho_neg_device = task.vrho_neg; + + + const auto* basis_eval_device = task.bf; + + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + const double factp = 0.5 * vrho_pos_device[tid_x]; + const double factm = 0.5 * vrho_neg_device[tid_x]; + double sign = 1.0; + if constexpr ( den_selector == DEN_Z ) sign = -1.0; + + z_matrix_device[ ibfoff ] = 0.5*(factp * basis_eval_device[ ibfoff ] + sign * factm * basis_eval_device[ ibfoff ]); + } + +} + +template +void zmat_lda_vxc_gks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const double* vrho_pos_device = task.vrho_pos; + const double* vrho_neg_device = task.vrho_neg; + + double* K_device; + if constexpr ( den_selector == DEN_Z ) K_device = task.K_z; + if constexpr ( den_selector == DEN_Y ) K_device = task.K_y; + if constexpr ( den_selector == DEN_X ) K_device = task.K_x; + + + + const auto* basis_eval_device = task.bf; + + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + const size_t ibfoff = tid_y * npts + tid_x; + const double factp = 0.5 * vrho_pos_device[tid_x]; + const double factm = 0.5 * vrho_neg_device[tid_x]; + + if constexpr ( den_selector == DEN_S ) { + z_matrix_device[ ibfoff ] = 0.5*(factp * basis_eval_device[ ibfoff ] + factm * basis_eval_device[ ibfoff ]); + } + else { + const double factk = 0.5 * (factp - factm); + z_matrix_device[ ibfoff ] = K_device[ ibfoff ] * factk * basis_eval_device[ ibfoff ]; + } + } + +} + + + + + + + + + +void zmat_gga_vxc_rks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const auto* vrho_device = task.vrho; + const auto* vgamma_device = task.vgamma; + const auto* den_x_eval_device = task.dden_sx; + const auto* den_y_eval_device = task.dden_sy; + const auto* den_z_eval_device = task.dden_sz; + + const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + const double fact_1 = 0.5 * vrho_device[tid_x] ; + const double fact_2 = 2.0 * vgamma_device[tid_x]; + + const double dx = den_x_eval_device[ tid_x ] * dbasis_x_eval_device[ ibfoff ]; + const double dy = den_y_eval_device[ tid_x ] * dbasis_y_eval_device[ ibfoff ]; + const double dz = den_z_eval_device[ tid_x ] * dbasis_z_eval_device[ ibfoff ]; + + z_matrix_device[ ibfoff ] = + fact_1 * basis_eval_device[ ibfoff ] + fact_2 * ( dx + dy + dz ); + + } +} + + + + + + + + +template +void zmat_gga_vxc_uks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + + const double* vrho_pos_device = task.vrho_pos; + const double* vrho_neg_device = task.vrho_neg; + const double* vgamma_pp_device = task.vgamma_pp; + const double* vgamma_pm_device = task.vgamma_pm; + const double* vgamma_mm_device = task.vgamma_mm; + + const auto* den_pos_x_eval_device = task.dden_sx; + const auto* den_pos_y_eval_device = task.dden_sy; + const auto* den_pos_z_eval_device = task.dden_sz; + const auto* den_neg_x_eval_device = task.dden_zx; + const auto* den_neg_y_eval_device = task.dden_zy; + const auto* den_neg_z_eval_device = task.dden_zz; + + + const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + + const double factp = 0.25 * vrho_pos_device[tid_x]; + const double factm = 0.25 * vrho_neg_device[tid_x]; + + const auto gga_fact_pp = vgamma_pp_device[tid_x]; + const auto gga_fact_pm = vgamma_pm_device[tid_x]; + const auto gga_fact_mm = vgamma_mm_device[tid_x]; + + const auto gga_fact_1 = 0.5*(gga_fact_pp + gga_fact_pm + gga_fact_mm); + const auto gga_fact_2 = 0.5*(gga_fact_pp - gga_fact_mm); + const auto gga_fact_3 = 0.5*(gga_fact_pp - gga_fact_pm + gga_fact_mm); + + double sign = 1.0; + + double x_fact, y_fact, z_fact; + + if constexpr ( den_selector == DEN_S ) { + x_fact = gga_fact_1 * den_pos_x_eval_device[ tid_x ] + gga_fact_2 * den_neg_x_eval_device[ tid_x ]; + y_fact = gga_fact_1 * den_pos_y_eval_device[ tid_x ] + gga_fact_2 * den_neg_y_eval_device[ tid_x ]; + z_fact = gga_fact_1 * den_pos_z_eval_device[ tid_x ] + gga_fact_2 * den_neg_z_eval_device[ tid_x ]; + + + } + if constexpr ( den_selector == DEN_Z ) { + sign = -1.0; + x_fact = gga_fact_3 * den_neg_x_eval_device[ tid_x ] + gga_fact_2 * den_pos_x_eval_device[ tid_x ]; + y_fact = gga_fact_3 * den_neg_y_eval_device[ tid_x ] + gga_fact_2 * den_pos_y_eval_device[ tid_x ]; + z_fact = gga_fact_3 * den_neg_z_eval_device[ tid_x ] + gga_fact_2 * den_pos_z_eval_device[ tid_x ]; + + } + + z_matrix_device[ ibfoff ] = x_fact * dbasis_x_eval_device[ ibfoff ] + + y_fact * dbasis_y_eval_device[ ibfoff ] + + z_fact * dbasis_z_eval_device[ ibfoff ] + + (factp + sign * factm) * basis_eval_device[ ibfoff ]; + } +} + + + + +template +void zmat_gga_vxc_gks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + + const double* vrho_pos_device = task.vrho_pos; + const double* vrho_neg_device = task.vrho_neg; + const double* vgamma_pp_device = task.vgamma_pp; + const double* vgamma_pm_device = task.vgamma_pm; + const double* vgamma_mm_device = task.vgamma_mm; + + + // for non-DEN_S + double* K_device; + double* H_device; + if constexpr ( den_selector == DEN_Z ) { K_device = task.K_z; H_device = task.H_z; } + if constexpr ( den_selector == DEN_Y ) { K_device = task.K_y; H_device = task.H_y; } + if constexpr ( den_selector == DEN_X ) { K_device = task.K_x; H_device = task.H_x; } + + const auto* dden_sx_eval_device = task.dden_sx; + const auto* dden_sy_eval_device = task.dden_sy; + const auto* dden_sz_eval_device = task.dden_sz; + const auto* dden_zx_eval_device = task.dden_zx; + const auto* dden_zy_eval_device = task.dden_zy; + const auto* dden_zz_eval_device = task.dden_zz; + const auto* dden_yx_eval_device = task.dden_yx; + const auto* dden_yy_eval_device = task.dden_yy; + const auto* dden_yz_eval_device = task.dden_yz; + const auto* dden_xx_eval_device = task.dden_xx; + const auto* dden_xy_eval_device = task.dden_xy; + const auto* dden_xz_eval_device = task.dden_xz; + + + const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + + const double fact_p = 0.5*vrho_pos_device[tid_x]; + const double fact_m = 0.5*vrho_neg_device[tid_x]; + + const auto gga_fact_pp = vgamma_pp_device[tid_x]; + const auto gga_fact_pm = vgamma_pm_device[tid_x]; + const auto gga_fact_mm = vgamma_mm_device[tid_x]; + + const auto gga_fact_1 = 0.5*(gga_fact_pp + gga_fact_pm + gga_fact_mm); + const auto gga_fact_2 = 0.5*(gga_fact_pp - gga_fact_mm); + const auto gga_fact_3 = 0.5*(gga_fact_pp - gga_fact_pm + gga_fact_mm); + + double s_fact, x_fact, y_fact, z_fact; + + if constexpr ( den_selector == DEN_S ) { + const double* Hz_device = task.H_z; + const double* Hy_device = task.H_y; + const double* Hx_device = task.H_x; + + s_fact = 0.5 * (fact_p + fact_m); + + x_fact = gga_fact_1 * dden_sx_eval_device[ tid_x ] + + gga_fact_2 * (Hz_device[ tid_x ] * dden_zx_eval_device[ tid_x ] + + Hy_device[ tid_x ] * dden_yx_eval_device[ tid_x ] + + Hx_device[ tid_x ] * dden_xx_eval_device[ tid_x ] ); + y_fact = gga_fact_1 * dden_sy_eval_device[ tid_x ] + + gga_fact_2 * (Hz_device[ tid_x ] * dden_zy_eval_device[ tid_x ] + + Hy_device[ tid_x ] * dden_yy_eval_device[ tid_x ] + + Hx_device[ tid_x ] * dden_xy_eval_device[ tid_x ] ); + z_fact = gga_fact_1 * dden_sz_eval_device[ tid_x ] + + gga_fact_2 * (Hz_device[ tid_x ] * dden_zz_eval_device[ tid_x ] + + Hy_device[ tid_x ] * dden_yz_eval_device[ tid_x ] + + Hx_device[ tid_x ] * dden_xz_eval_device[ tid_x ] ); + } + + if constexpr ( den_selector == DEN_Z ) { + s_fact = K_device[ tid_x ] * 0.5 * (fact_p - fact_m); + x_fact = gga_fact_3 * dden_zx_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sx_eval_device[ tid_x ]; + y_fact = gga_fact_3 * dden_zy_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sy_eval_device[ tid_x ]; + z_fact = gga_fact_3 * dden_zz_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sz_eval_device[ tid_x ]; + } + + if constexpr ( den_selector == DEN_Y ) { + s_fact = K_device[ tid_x ] * 0.5 * (fact_p - fact_m); + x_fact = gga_fact_3 * dden_yx_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sx_eval_device[ tid_x ]; + y_fact = gga_fact_3 * dden_yy_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sy_eval_device[ tid_x ]; + z_fact = gga_fact_3 * dden_yz_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sz_eval_device[ tid_x ]; + } + + if constexpr ( den_selector == DEN_X ) { + s_fact = K_device[ tid_x ] * 0.5 * (fact_p - fact_m); + x_fact = gga_fact_3 * dden_xx_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sx_eval_device[ tid_x ]; + y_fact = gga_fact_3 * dden_xy_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sy_eval_device[ tid_x ]; + z_fact = gga_fact_3 * dden_xz_eval_device[ tid_x ] + + gga_fact_2 * H_device[ tid_x ] * dden_sz_eval_device[ tid_x ]; + } + + z_matrix_device[ ibfoff ] = x_fact * dbasis_x_eval_device[ ibfoff ] + + y_fact * dbasis_y_eval_device[ ibfoff ] + + z_fact * dbasis_z_eval_device[ ibfoff ] + + s_fact * basis_eval_device[ ibfoff ]; + + } +} + + + + +template +void zmat_mgga_vxc_rks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const auto* vrho_device = task.vrho; + const auto* vgamma_device = task.vgamma; + const double* vlapl_device = need_lapl ? task.vlapl : nullptr; + const auto* den_x_eval_device = task.dden_sx; + const auto* den_y_eval_device = task.dden_sy; + const auto* den_z_eval_device = task.dden_sz; + + const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + const double* d2basis_lapl_eval_device = + need_lapl ? task.d2bflapl : nullptr; + + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + const double fact_1 = 0.5 * vrho_device[tid_x] ; + const double fact_2 = 2.0 * vgamma_device[tid_x]; + + const double dx = den_x_eval_device[ tid_x ] * dbasis_x_eval_device[ ibfoff ]; + const double dy = den_y_eval_device[ tid_x ] * dbasis_y_eval_device[ ibfoff ]; + const double dz = den_z_eval_device[ tid_x ] * dbasis_z_eval_device[ ibfoff ]; + + double val = + fact_1 * basis_eval_device[ ibfoff ] + fact_2 * ( dx + dy + dz ); + + if constexpr (need_lapl) { + val += vlapl_device[tid_x] * d2basis_lapl_eval_device[ibfoff]; + } + + z_matrix_device[ ibfoff ] = val; + } +} + +template +void zmat_mgga_vxc_uks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + + const double* vrho_pos_device = task.vrho_pos; + const double* vrho_neg_device = task.vrho_neg; + const double* vlapl_pos_device = task.vlapl_pos; + const double* vlapl_neg_device = task.vlapl_neg; + const double* vgamma_pp_device = task.vgamma_pp; + const double* vgamma_pm_device = task.vgamma_pm; + const double* vgamma_mm_device = task.vgamma_mm; + + const auto* den_pos_x_eval_device = task.dden_sx; + const auto* den_pos_y_eval_device = task.dden_sy; + const auto* den_pos_z_eval_device = task.dden_sz; + const auto* den_neg_x_eval_device = task.dden_zx; + const auto* den_neg_y_eval_device = task.dden_zy; + const auto* den_neg_z_eval_device = task.dden_zz; + + + const auto* basis_eval_device = task.bf; + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + const auto* d2basis_lapl_eval_device = task.d2bflapl; + + auto* z_matrix_device = task.zmat; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + + const double factp = 0.25 * vrho_pos_device[tid_x]; + const double factm = 0.25 * vrho_neg_device[tid_x]; + + const auto gga_fact_pp = vgamma_pp_device[tid_x]; + const auto gga_fact_pm = vgamma_pm_device[tid_x]; + const auto gga_fact_mm = vgamma_mm_device[tid_x]; + + const auto gga_fact_1 = 0.5*(gga_fact_pp + gga_fact_pm + gga_fact_mm); + const auto gga_fact_2 = 0.5*(gga_fact_pp - gga_fact_mm); + const auto gga_fact_3 = 0.5*(gga_fact_pp - gga_fact_pm + gga_fact_mm); + + double sign = 1.0; + + double x_fact, y_fact, z_fact; + + if constexpr ( den_selector == DEN_S ) { + x_fact = gga_fact_1 * den_pos_x_eval_device[ tid_x ] + gga_fact_2 * den_neg_x_eval_device[ tid_x ]; + y_fact = gga_fact_1 * den_pos_y_eval_device[ tid_x ] + gga_fact_2 * den_neg_y_eval_device[ tid_x ]; + z_fact = gga_fact_1 * den_pos_z_eval_device[ tid_x ] + gga_fact_2 * den_neg_z_eval_device[ tid_x ]; + } + if constexpr ( den_selector == DEN_Z ) { + sign = -1.0; + x_fact = gga_fact_3 * den_neg_x_eval_device[ tid_x ] + gga_fact_2 * den_pos_x_eval_device[ tid_x ]; + y_fact = gga_fact_3 * den_neg_y_eval_device[ tid_x ] + gga_fact_2 * den_pos_y_eval_device[ tid_x ]; + z_fact = gga_fact_3 * den_neg_z_eval_device[ tid_x ] + gga_fact_2 * den_pos_z_eval_device[ tid_x ]; + } + + auto val = x_fact * dbasis_x_eval_device[ ibfoff ] + + y_fact * dbasis_y_eval_device[ ibfoff ] + + z_fact * dbasis_z_eval_device[ ibfoff ] + + (factp + sign * factm) * basis_eval_device[ ibfoff ]; + + if constexpr (need_lapl) { + const double lfactp = vlapl_pos_device[tid_x]; + const double lfactm = vlapl_neg_device[tid_x]; + + val += 0.5 * (lfactp + sign * lfactm) * d2basis_lapl_eval_device[ ibfoff ]; + } + + z_matrix_device[ ibfoff ] = val; + } +} + + + +#define ZMAT_VXC_KERN(xc_approx) \ + ::sycl::queue& stream = queue.queue_as(); \ + sycl::launch_dim threads(sycl::warp_size,sycl::max_warps_per_thread_block,1); \ + sycl::launch_dim blocks( util::div_ceil( max_npts, threads.x ), \ + util::div_ceil( max_nbf, threads.y ), \ + ntasks ); \ + switch( scheme ) { \ + case RKS: \ + sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_vxc_rks_kernel", \ + [=](){ zmat_##xc_approx##_vxc_rks_kernel( ntasks, tasks_device ); } ); \ + break; \ + case UKS: \ + if ( sel == DEN_S ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_vxc_uks_kernel", \ + [=](){ zmat_##xc_approx##_vxc_uks_kernel( ntasks, tasks_device ); } ); \ + else if ( sel == DEN_Z ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_vxc_uks_kernel", \ + [=](){ zmat_##xc_approx##_vxc_uks_kernel( ntasks, tasks_device ); } ); \ + else GAUXC_GENERIC_EXCEPTION( "zmat_##xc_approx##_vxc invalid density" ); \ + break; \ + case GKS: \ + if ( sel == DEN_S ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_vxc_gks_kernel", \ + [=](){ zmat_##xc_approx##_vxc_gks_kernel( ntasks, tasks_device ); } ); \ + else if ( sel == DEN_Z ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_vxc_gks_kernel", \ + [=](){ zmat_##xc_approx##_vxc_gks_kernel( ntasks, tasks_device ); } ); \ + else if ( sel == DEN_Y ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_vxc_gks_kernel", \ + [=](){ zmat_##xc_approx##_vxc_gks_kernel( ntasks, tasks_device ); } ); \ + else if ( sel == DEN_X ) sycl::launch_kernel( stream, blocks, threads, "zmat_" #xc_approx "_vxc_gks_kernel", \ + [=](){ zmat_##xc_approx##_vxc_gks_kernel( ntasks, tasks_device ); } ); \ + else GAUXC_GENERIC_EXCEPTION( "zmat_##xc_approx##_vxc invalid density" ); \ + break; \ + default: \ + GAUXC_GENERIC_EXCEPTION( "zmat_##xc_approx##_vxc invalid KS scheme" ); \ + } + + + +void zmat_lda_vxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + integrator_ks_scheme scheme, + density_id sel, + device_queue queue ) { +ZMAT_VXC_KERN(lda) +} + + + +void zmat_gga_vxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + integrator_ks_scheme scheme, + density_id sel, + device_queue queue ) { +ZMAT_VXC_KERN(gga) +} + + + +void zmat_mgga_vxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + bool do_lapl, + integrator_ks_scheme scheme, + density_id sel, + device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + + sycl::launch_dim threads(sycl::warp_size,sycl::max_warps_per_thread_block,1); + sycl::launch_dim blocks( util::div_ceil( max_npts, threads.x ), + util::div_ceil( max_nbf, threads.y ), + ntasks ); + + if(scheme == RKS) { + if(do_lapl) + sycl::launch_kernel( stream, blocks, threads, "zmat_mgga_vxc_rks_kernel", + [=](){ zmat_mgga_vxc_rks_kernel( ntasks, tasks_device ); } ); + else + sycl::launch_kernel( stream, blocks, threads, "zmat_mgga_vxc_rks_kernel", + [=](){ zmat_mgga_vxc_rks_kernel( ntasks, tasks_device ); } ); + } else if(scheme == UKS) { + switch(sel) { + case DEN_S: + if(do_lapl) + sycl::launch_kernel( stream, blocks, threads, "zmat_mgga_vxc_uks_kernel", + [=](){ zmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + else + sycl::launch_kernel( stream, blocks, threads, "zmat_mgga_vxc_uks_kernel", + [=](){ zmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + break; + case DEN_Z: + if(do_lapl) + sycl::launch_kernel( stream, blocks, threads, "zmat_mgga_vxc_uks_kernel", + [=](){ zmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + else + sycl::launch_kernel( stream, blocks, threads, "zmat_mgga_vxc_uks_kernel", + [=](){ zmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + break; + } + } else { + GAUXC_GENERIC_EXCEPTION("MGGA + DEVICE + GKS NYI"); + } + +} + + + + + + + + + + + + + + + + +template +void mmat_mgga_vxc_rks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const auto* vtau_device = task.vtau; + const double* vlapl_device = need_lapl ? task.vlapl : nullptr; + + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + + auto* mmat_x = task.xmat_x; + auto* mmat_y = task.xmat_y; + auto* mmat_z = task.xmat_z; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + const size_t ibfoff = tid_y * npts + tid_x; + const double fact_1 = 0.25 * vtau_device[tid_x] + + (need_lapl ? vlapl_device[tid_x] : 0.0); + + mmat_x[ ibfoff ] = fact_1 * dbasis_x_eval_device[ ibfoff ]; + mmat_y[ ibfoff ] = fact_1 * dbasis_y_eval_device[ ibfoff ]; + mmat_z[ ibfoff ] = fact_1 * dbasis_z_eval_device[ ibfoff ]; + } +} + +template +void mmat_mgga_vxc_uks_kernel( size_t ntasks, + XCDeviceTask* tasks_device ) { + + + auto it = GauXC::sycl::this_item(); + const int batch_idx = it.get_group(0); + if( batch_idx >= ntasks ) return; + + auto& task = tasks_device[ batch_idx ]; + const auto npts = task.npts; + const auto nbf = task.bfn_screening.nbe; + const auto* vtau_pos_device = task.vtau_pos; + const auto* vtau_neg_device = task.vtau_neg; + const double* vlapl_pos_device = need_lapl ? task.vlapl_pos : nullptr; + const double* vlapl_neg_device = need_lapl ? task.vlapl_neg : nullptr; + + const auto* dbasis_x_eval_device = task.dbfx; + const auto* dbasis_y_eval_device = task.dbfy; + const auto* dbasis_z_eval_device = task.dbfz; + + auto* mmat_x = task.xmat_x; + auto* mmat_y = task.xmat_y; + auto* mmat_z = task.xmat_z; + + const int tid_x = it.get_global_id(2); + const int tid_y = it.get_global_id(1); + + if( tid_x < npts and tid_y < nbf ) { + + double sign = 1.0; + if(id == DEN_Z) sign = -1; + + const size_t ibfoff = tid_y * npts + tid_x; + const auto tfactp = 0.25 * vtau_pos_device[tid_x]; + const auto tfactm = 0.25 * vtau_neg_device[tid_x]; + const double fact_tau = 0.5 * (tfactp + sign * tfactm); + double fact_lapl = 0.0; + if(need_lapl) { + const auto lfactp = vlapl_pos_device[tid_x]; + const auto lfactm = vlapl_neg_device[tid_x]; + fact_lapl = 0.5 * (lfactp + sign * lfactm); + } + const double fact_1 = fact_tau + fact_lapl; + + mmat_x[ ibfoff ] = fact_1 * dbasis_x_eval_device[ ibfoff ]; + mmat_y[ ibfoff ] = fact_1 * dbasis_y_eval_device[ ibfoff ]; + mmat_z[ ibfoff ] = fact_1 * dbasis_z_eval_device[ ibfoff ]; + } +} + +//void print_zmat_stats( size_t ntasks, +// XCDeviceTask* tasks_device) { +// +// for(size_t iT = 0; iT < ntasks; ++iT) { +// auto& task = tasks_device[iT]; +// const auto npts = task.npts; +// const auto nbf = task.bfn_screening.nbe; +// +// const auto* zmat = task.zmat; +// const auto* bmat = task.bf; +// const auto* blmat = task.d2bflapl; +// +// double znrm = 0.0, bnrm = 0.0, blnrm = 0.0; +// for(auto j = 0; j < npts*nbf; ++j) { +// znrm += zmat[j] * zmat[j]; +// bnrm += bmat[j] * bmat[j]; +// blnrm += blmat[j] * blmat[j]; +// } +// +// const auto* eps = task.eps; +// const auto* vgamma = task.vgamma; +// const auto* vtau = task.vtau; +// const auto* vlapl = task.vlapl; +// const auto* vrho = task.vrho; +// const auto* gamma = task.gamma; +// const auto* tau = task.tau; +// const auto* lapl = task.lapl; +// const auto* rho = task.den; +// double enrm = 0.0, gnrm = 0.0, tnrm = 0.0, rnrm = 0.0, lnrm = 0.0; +// double vgnrm = 0.0, vtnrm = 0.0, vrnrm = 0.0, vlnrm = 0.0; +// for(auto j = 0; j < npts; ++j) { +// enrm += eps[j] * eps[j]; +// vrnrm += vrho[j] * vrho[j]; +// vgnrm += vgamma[j] * vgamma[j]; +// vtnrm += vtau[j] * vtau[j]; +// vlnrm += vlapl[j] * vlapl[j]; +// +// rnrm += rho[j] * rho[j]; +// gnrm += gamma[j] * gamma[j]; +// tnrm += tau[j] * tau[j]; +// lnrm += lapl[j] * lapl[j]; +// } +// +// printf("ITASK = %lu B = %.6e BL = %.6e R = %.6e G = %.6e T = %.6e L = %.6e E = %.6e VR = %.6e VG = %6e VT = %.6e VL = %.6e Z = %.6e \n", +// iT, bnrm, blnrm, rnrm, gnrm, tnrm, lnrm, enrm, vrnrm, vgnrm, vtnrm, vlnrm, znrm); +// } +// +//} + +void mmat_mgga_vxc( size_t ntasks, + int32_t max_nbf, + int32_t max_npts, + XCDeviceTask* tasks_device, + bool do_lapl, + integrator_ks_scheme scheme, + density_id sel, + device_queue queue ) { + + ::sycl::queue& stream = queue.queue_as(); + + + sycl::launch_dim threads(sycl::warp_size,sycl::max_warps_per_thread_block,1); + sycl::launch_dim blocks( util::div_ceil( max_npts, threads.x ), + util::div_ceil( max_nbf, threads.y ), + ntasks ); + + if(scheme == RKS) { + if(do_lapl) + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_vxc_rks_kernel", + [=](){ mmat_mgga_vxc_rks_kernel( ntasks, tasks_device ); } ); + else + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_vxc_rks_kernel", + [=](){ mmat_mgga_vxc_rks_kernel( ntasks, tasks_device ); } ); + } else if(scheme == UKS) { + switch(sel) { + case DEN_S: + if(do_lapl) + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_vxc_uks_kernel", + [=](){ mmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + else + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_vxc_uks_kernel", + [=](){ mmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + break; + case DEN_Z: + if(do_lapl) + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_vxc_uks_kernel", + [=](){ mmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + else + sycl::launch_kernel( stream, blocks, threads, "mmat_mgga_vxc_uks_kernel", + [=](){ mmat_mgga_vxc_uks_kernel( ntasks, tasks_device ); } ); + break; + } + } else { + GAUXC_GENERIC_EXCEPTION("MGGA + DEVICE + GKS NYI"); + } + + + //print_zmat_stats(ntasks,tasks_device); +} + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1.cxx b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1.cxx new file mode 100644 index 000000000..15d25ee9c --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1.cxx @@ -0,0 +1,88 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "sycl_aos_scheme1.hpp" +#include "device/sycl/sycl_backend.hpp" +#include "sycl_aos_scheme1_weights.hpp" +#include "device/common/device_blas.hpp" + +namespace GauXC { + +template +std::unique_ptr SyclAoSScheme1::create_device_data(const DeviceRuntimeEnvironment& rt) { + return std::make_unique(rt); +} + +template +void SyclAoSScheme1::partition_weights( XCDeviceData* _data ) { + auto* data = dynamic_cast(_data); + if( !data ) GAUXC_BAD_LWD_DATA_CAST(); + + auto device_backend = dynamic_cast(data->device_backend_); + if( !device_backend ) GAUXC_BAD_BACKEND_CAST(); + + + // Compute distances from grid to atomic centers + const auto ldatoms = data->get_ldatoms(); + auto base_stack = data->base_stack; + auto static_stack = data->static_stack; + auto scheme1_stack = data->scheme1_stack; + sycl_aos_scheme1_weights_wrapper( data->total_npts_task_batch, data->global_dims.natoms, + base_stack.points_x_device, base_stack.points_y_device, base_stack.points_z_device, + static_stack.rab_device, ldatoms, static_stack.coords_device, + scheme1_stack.dist_scratch_device, ldatoms, scheme1_stack.iparent_device, + scheme1_stack.dist_nearest_device, base_stack.weights_device, + device_backend->master_stream->queue ); +} + +template +void SyclAoSScheme1::eval_weight_1st_deriv_contracted( XCDeviceData* _data, XCWeightAlg alg ) { + if( alg != XCWeightAlg::SSF ) { + GAUXC_GENERIC_EXCEPTION("Weight Algorithm NYI for SYCL AoS Scheme1"); + } + auto* data = dynamic_cast(_data); + if( !data ) GAUXC_BAD_LWD_DATA_CAST(); + + auto device_backend = dynamic_cast(data->device_backend_); + if( !device_backend ) GAUXC_BAD_BACKEND_CAST(); + + // make w times f vector + const bool is_UKS = data->allocated_terms.ks_scheme == UKS; + const bool is_GKS = data->allocated_terms.ks_scheme == GKS; + const bool is_pol = is_UKS or is_GKS; + auto base_stack = data->base_stack; + if( is_pol ) + increment( data->device_backend_->master_blas_handle(), base_stack.den_z_eval_device, + base_stack.den_s_eval_device, data->total_npts_task_batch ); + + hadamard_product(data->device_backend_->master_blas_handle(), data->total_npts_task_batch, 1, base_stack.den_s_eval_device, 1, + base_stack.eps_eval_device, 1); + + + // Compute distances from grid to atomic centers + const auto ldatoms = data->get_ldatoms(); + auto static_stack = data->static_stack; + auto scheme1_stack = data->scheme1_stack; + sycl_aos_scheme1_weight_1st_deriv_wrapper( data->total_npts_task_batch, data->global_dims.natoms, + base_stack.points_x_device, base_stack.points_y_device, base_stack.points_z_device, + static_stack.rab_device, ldatoms, static_stack.coords_device, + scheme1_stack.dist_scratch_device, ldatoms, scheme1_stack.iparent_device, + scheme1_stack.dist_nearest_device, base_stack.eps_eval_device, + static_stack.exc_grad_device, device_backend->master_stream->queue ); +} + + +template struct SyclAoSScheme1; +template struct SyclAoSScheme1; + + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1.hpp b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1.hpp new file mode 100644 index 000000000..bfe9c905a --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1.hpp @@ -0,0 +1,68 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "device/scheme1_base.hpp" +#include "device/scheme1_onemkl_base.hpp" + +namespace GauXC { + +namespace alg_constants { + +struct SyclAoSScheme1 { + static constexpr uint32_t weight_unroll = 4; + static constexpr uint32_t weight_thread_block = 512; + static constexpr uint32_t weight_thread_block_per_sm = 2; + static constexpr uint32_t max_submat_blocks = 10; + + struct ObaraSaika { + static constexpr int points_per_subtask = 256; + }; + +}; + +} + +template +struct SyclAoSScheme1 : public Base { + + // API Overrides + void partition_weights( XCDeviceData* ) override final; + void eval_weight_1st_deriv_contracted( XCDeviceData*, XCWeightAlg ) override final; + + std::unique_ptr create_device_data(const DeviceRuntimeEnvironment&) override final; + + struct Data; + +}; + +extern template struct SyclAoSScheme1; +extern template struct SyclAoSScheme1; + +template +struct SyclAoSScheme1::Data : public Base::Data { + + virtual ~Data() noexcept; + Data() = delete; + Data(const DeviceRuntimeEnvironment& rt); + + // Final overrides + size_t get_submat_chunk_size(int32_t,int32_t) override final; + size_t get_ldatoms() override final; + size_t get_rab_align() override final; + int get_points_per_subtask() override final; + +}; + +extern template struct SyclAoSScheme1::Data; +extern template struct SyclAoSScheme1::Data; + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_data.cxx b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_data.cxx new file mode 100644 index 000000000..d069eb7d1 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_data.cxx @@ -0,0 +1,74 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "sycl_aos_scheme1.hpp" +#include "buffer_adaptor.hpp" +#include +#include "device/sycl/sycl_backend.hpp" +#include "device_specific/sycl_vector_types.hpp" + +namespace GauXC { + +template +SyclAoSScheme1::Data::~Data() noexcept = default; + +template +SyclAoSScheme1::Data::Data(const DeviceRuntimeEnvironment& rt) : + Base::Data( rt ) { } + +template +size_t SyclAoSScheme1::Data::get_ldatoms() { + constexpr auto weight_unroll = alg_constants::SyclAoSScheme1::weight_unroll; + return util::div_ceil( this->global_dims.natoms, weight_unroll ) * weight_unroll; +} + +template +size_t SyclAoSScheme1::Data::get_rab_align() { + return sizeof(double2); +} + +template +int SyclAoSScheme1::Data::get_points_per_subtask() { + return alg_constants::SyclAoSScheme1::ObaraSaika::points_per_subtask; +} + + + +template +size_t SyclAoSScheme1::Data::get_submat_chunk_size(int32_t LDA, + int32_t dev_id) { + + constexpr auto max_submat_blocks = + alg_constants::SyclAoSScheme1::max_submat_blocks; + + (void)(dev_id); // The chunk is sized against the device this Data is bound to + + auto* backend = dynamic_cast(this->device_backend_); + if( !backend ) GAUXC_BAD_BACKEND_CAST(); + + const size_t l2_cache_size = + backend->device.get_info< ::sycl::info::device::global_mem_cache_size >(); + + int l2_block_size = (int) sqrt(0.75 * ((double) l2_cache_size / 8)); + int min_block_size = LDA / max_submat_blocks; + + int block_size = std::max(l2_block_size, min_block_size); + block_size = std::min(block_size, LDA); + + return block_size; + +} + +template struct SyclAoSScheme1::Data; +template struct SyclAoSScheme1::Data; + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_weights.cxx b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_weights.cxx new file mode 100644 index 000000000..c2d5d8b5b --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_weights.cxx @@ -0,0 +1,60 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "sycl_aos_scheme1.hpp" +#include "device/sycl/sycl_backend.hpp" +#include "kernels/grid_to_center.hpp" +#include "kernels/sycl_ssf_2d.hpp" +#include "kernels/sycl_ssf_1d.hpp" +#include "kernels/sycl_launch.hpp" +#include "sycl_aos_scheme1_weights.hpp" + +namespace GauXC { + +void sycl_aos_scheme1_weights_wrapper( int32_t npts, int32_t natoms, + const double* points_x, const double* points_y, const double* points_z, + const double* RAB, int32_t ldRAB, const double* coords, + double* dist, int32_t lddist, const int32_t* iparent, + const double* dist_nearest, double* weights, ::sycl::queue& stream ) { + + // Compute distances from grid to atomic centers + compute_grid_to_center_dist( npts, natoms, coords, points_x, points_y, points_z, + dist, lddist, stream ); + + partition_weights_ssf_1d( npts, natoms, RAB, ldRAB, coords, dist, lddist, + iparent, dist_nearest, weights, stream ); + +} + + +void sycl_aos_scheme1_weight_1st_deriv_wrapper( + int32_t npts, int32_t natoms, + const double* points_x, const double* points_y, const double* points_z, + const double* RAB, int32_t ldRAB, const double* coords, + double* dist, int32_t lddist, const int32_t* iparent, + const double* dist_nearest, const double* w_times_f, + double* exc_grad_w, ::sycl::queue& stream ){ + + // Compute distances from grid to atomic centers + compute_grid_to_center_dist( npts, natoms, coords, points_x, points_y, points_z, + dist, lddist, stream ); + + eval_weight_1st_deriv_contracted_ssf_1d( npts, natoms, RAB, ldRAB, coords, + points_x, points_y, points_z, dist, lddist, iparent, dist_nearest, + w_times_f, exc_grad_w, stream ); + +} + + + + + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_weights.hpp b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_weights.hpp new file mode 100644 index 000000000..84a13ee12 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/sycl_aos_scheme1_weights.hpp @@ -0,0 +1,31 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include + +namespace GauXC { + +void sycl_aos_scheme1_weights_wrapper( int32_t npts, int32_t natoms, + const double* points_x, const double* points_y, const double* points_z, + const double* RAB, int32_t ldRAB, const double* coords, + double* dist, int32_t lddist, const int32_t* iparent, + const double* dist_nearest, double* weights, ::sycl::queue& stream ); + +void sycl_aos_scheme1_weight_1st_deriv_wrapper( + int32_t npts, int32_t natoms, + const double* points_x, const double* points_y, const double* points_z, + const double* RAB, int32_t ldRAB, const double* coords, + double* dist, int32_t lddist, const int32_t* iparent, + const double* dist_nearest, const double* w_times_f, + double* exc_grad_w, ::sycl::queue& stream ); + +} diff --git a/src/xc_integrator/local_work_driver/device/sycl/xc_functional_eval_wrapper.cxx b/src/xc_integrator/local_work_driver/device/sycl/xc_functional_eval_wrapper.cxx new file mode 100644 index 000000000..dcb4ae750 --- /dev/null +++ b/src/xc_integrator/local_work_driver/device/sycl/xc_functional_eval_wrapper.cxx @@ -0,0 +1,75 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "device/common/xc_functional_eval_wrapper.hpp" +#include "device_specific/sycl_util.hpp" + +namespace GauXC { + +void eval_kern_exc_vxc_lda( const functional_type& func, size_t npts, + const double* rho, double* eps, double* vrho, device_queue queue ) { + + ::sycl::queue* stream = &queue.queue_as().queue; + func.eval_exc_vxc_device( npts, rho, eps, vrho, stream ); + +} + +void eval_kern_exc_vxc_gga( const functional_type& func, size_t npts, + const double* rho, const double* gamma, double* eps, double* vrho, + double* vgamma, device_queue queue ) { + + ::sycl::queue* stream = &queue.queue_as().queue; + func.eval_exc_vxc_device( npts, rho, gamma, eps, vrho, vgamma, stream ); + +} + +void eval_kern_exc_vxc_mgga( const functional_type& func, size_t npts, + const double* rho, const double* gamma, const double* tau, const double* lapl, + double* eps, double* vrho, double* vgamma, double* vtau, double* vlapl, + device_queue queue ) { + + ::sycl::queue* stream = &queue.queue_as().queue; + func.eval_exc_vxc_device( npts, rho, gamma, lapl, tau, eps, vrho, vgamma, vlapl, vtau, stream ); + +} + + + +void eval_kern_vxc_fxc_lda( const functional_type& func, size_t npts, + const double* rho, double* vrho, double* v2rho2, device_queue queue ) { + + ::sycl::queue* stream = &queue.queue_as().queue; + func.eval_vxc_fxc_device( npts, rho, vrho, v2rho2, stream ); +} + +void eval_kern_vxc_fxc_gga( const functional_type& func, size_t npts, + const double* rho, const double* gamma, double* vrho, double* vgamma, + double* v2rho2, double* v2rhogamma, double* v2gamma2, device_queue queue ) { + + ::sycl::queue* stream = &queue.queue_as().queue; + func.eval_vxc_fxc_device( npts, rho, gamma, vrho, vgamma, v2rho2, v2rhogamma, v2gamma2, stream ); +} + +void eval_kern_vxc_fxc_mgga( const functional_type& func, size_t npts, + const double* rho, const double* gamma, const double* lapl, const double* tau, + double* vrho, double* vgamma, double* vlapl, double* vtau, + double* v2rho2, double* v2rhogamma, double* v2rholapl, double* v2rhotau, + double* v2gamma2, double* v2gammalapl, double* v2gammatau, double* v2lapl2, + double* v2lapltau, double* v2tau2, device_queue queue ){ + + ::sycl::queue* stream = &queue.queue_as().queue; + func.eval_vxc_fxc_device( npts, rho, gamma, lapl, tau, vrho, vgamma, vlapl, vtau, + v2rho2, v2rhogamma, v2rholapl, v2rhotau, v2gamma2, v2gammalapl, v2gammatau, + v2lapl2, v2lapltau, v2tau2, stream ); +} + + +} diff --git a/src/xc_integrator/local_work_driver/factory.cxx b/src/xc_integrator/local_work_driver/factory.cxx index 3fbf49ab1..dcea311a2 100644 --- a/src/xc_integrator/local_work_driver/factory.cxx +++ b/src/xc_integrator/local_work_driver/factory.cxx @@ -14,6 +14,7 @@ #ifdef GAUXC_HAS_DEVICE #include "device/cuda/cuda_aos_scheme1.hpp" #include "device/hip/hip_aos_scheme1.hpp" +#include "device/sycl/sycl_aos_scheme1.hpp" #endif namespace GauXC { @@ -54,6 +55,9 @@ LocalWorkDriverFactory::ptr_return_t #ifdef GAUXC_HAS_MAGMA using scheme1_magma = HipAoSScheme1; #endif +#elif defined(GAUXC_HAS_SYCL) + using scheme1_default = SyclAoSScheme1<>; + using scheme1_onemkl = SyclAoSScheme1; #endif #ifdef GAUXC_HAS_DEVICE @@ -66,6 +70,10 @@ LocalWorkDriverFactory::ptr_return_t #ifdef GAUXC_HAS_CUTLASS else if( name == "SCHEME1-CUTLASS" ) return std::make_unique( std::make_unique() ); +#endif +#ifdef GAUXC_HAS_SYCL + else if( name == "SCHEME1-ONEMKL" ) + return std::make_unique( std::make_unique() ); #endif else #endif diff --git a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp index 416a49c5a..84a45e9b7 100644 --- a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp +++ b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp @@ -286,7 +286,7 @@ void IncoreReplicatedXCDeviceIntegrator:: device_data.reset_allocations(); device_data.allocate_static_data_exc_vxc( nbf, nshells, enabled_terms, do_vxc ); - device_data.send_static_data_density_basis( Ps, ldps, Pz, ldpz, Px, ldpx, Py, ldpy, basis ); + device_data.send_static_data_density_basis( Ps, ldps, Pz, ldpz, Py, ldpy, Px, ldpx, basis ); diff --git a/src/xc_integrator/xc_data/device/xc_device_task.hpp b/src/xc_integrator/xc_data/device/xc_device_task.hpp index 58ab323c7..ca7b8b49a 100644 --- a/src/xc_integrator/xc_data/device/xc_device_task.hpp +++ b/src/xc_integrator/xc_data/device/xc_device_task.hpp @@ -10,6 +10,8 @@ * See LICENSE.txt for details */ #pragma once +#include +#include namespace GauXC { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 49e788297..8d44d6026 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -73,6 +73,14 @@ add_executable( gauxc_test c_api_test.cxx ) target_link_libraries( gauxc_test PUBLIC gauxc gauxc_catch2 Eigen3::Eigen ) + +# collocation.cxx and weights.cxx submit SYCL kernels directly, so those two +# translation units carry device code and need the SYCL device flags. The rest +# of the test suite is host-only and is left alone; -fsycl is not free on a +# translation unit that produces no device image. +if( GAUXC_HAS_SYCL ) + gauxc_sycl_device_sources( collocation.cxx weights.cxx ) +endif() if(GAUXC_ENABLE_CUTLASS) include(gauxc-cutlass) target_link_libraries(gauxc_test PUBLIC gauxc_cutlass) diff --git a/tests/collocation.cxx b/tests/collocation.cxx index 741505900..84e29fa8c 100644 --- a/tests/collocation.cxx +++ b/tests/collocation.cxx @@ -13,6 +13,7 @@ #include "collocation_host.hpp" #include "collocation_cuda.hpp" #include "collocation_hip.hpp" +#include "collocation_sycl.hpp" //#define GENERATE_TESTS @@ -100,6 +101,35 @@ TEST_CASE( "Water / cc-pVDZ", "[collocation]" ) { } #endif // GAUXC_HAS_HIP +#ifdef GAUXC_HAS_SYCL + BasisSetMap basis_map( basis, mol ); + SECTION( "SYCL Eval" ) { + test_sycl_collocation( basis, ref_file ); + } + SECTION( "SYCL Shell to Task Eval" ) { + test_sycl_collocation_shell_to_task( basis, basis_map, ref_file ); + } + + SECTION( "SYCL Eval Grad" ) { + test_sycl_collocation_deriv1( basis, ref_file ); + } + SECTION( "SYCL Shell to Task Eval Grad" ) { + test_sycl_collocation_shell_to_task_gradient( basis, basis_map, ref_file ); + } + + SECTION( "SYCL Shell to Task Eval Hessian" ) { + test_sycl_collocation_shell_to_task_hessian( basis, basis_map, ref_file ); + } + + SECTION( "SYCL Shell to Task Eval Laplacian" ) { + test_sycl_collocation_shell_to_task_laplacian( basis, basis_map, ref_file ); + } + + SECTION( "SYCL Shell to Task Eval Laplacian Gradient" ) { + test_sycl_collocation_shell_to_task_lapgrad( basis, basis_map, ref_file ); + } +#endif // GAUXC_HAS_SYCL + diff --git a/tests/collocation_sycl.hpp b/tests/collocation_sycl.hpp new file mode 100644 index 000000000..90be30d48 --- /dev/null +++ b/tests/collocation_sycl.hpp @@ -0,0 +1,444 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#ifdef GAUXC_HAS_SYCL +#include "collocation_common.hpp" +#include "hdf5_test_serialization.hpp" +#include "hdf5_test_serialization_impl.hpp" +#include "device/common/collocation_device.hpp" +#include "device_specific/sycl_util.hpp" +#include + + +auto populate_device_sycl( const BasisSet& basis, + const std::vector& ref_data, + bool pop_grad, bool pop_hess, bool pop_lapl, bool pop_lapl_grad, + ::sycl::queue& q ) { + + std::vector< XCDeviceTask > tasks; + + auto shells_device = util::sycl_malloc>( basis.size(), q ); + std::vector> shells( basis ); + util::sycl_copy( basis.size(), shells_device, shells.data(), q ); + + for( auto& d : ref_data ) { + const auto npts = d.pts.size(); + const auto nbf = d.eval.size() / npts; + + const auto& mask = d.mask; + const auto& pts = d.pts; + + /// XXX: THIS DOES NOT POPULATE A VALID TASK, ONLY WHAT's REQUIRED FOR THIS + // TEST + auto& task = tasks.emplace_back(); + task.npts = npts; + task.bfn_screening.nbe = nbf; + task.bfn_screening.nshells = mask.size(); + task.bfn_screening.shell_offs = util::sycl_malloc( mask.size(), q ); + task.bfn_screening.shell_list = util::sycl_malloc( mask.size(), q ); + + task.points_x = util::sycl_malloc( npts, q ); + task.points_y = util::sycl_malloc( npts, q ); + task.points_z = util::sycl_malloc( npts, q ); + task.bf = util::sycl_malloc( nbf * npts, q ); + if(pop_grad) { + task.dbfx = util::sycl_malloc( nbf * npts, q ); + task.dbfy = util::sycl_malloc( nbf * npts, q ); + task.dbfz = util::sycl_malloc( nbf * npts, q ); + } + + if(pop_hess) { + task.d2bfxx = util::sycl_malloc( nbf * npts, q ); + task.d2bfxy = util::sycl_malloc( nbf * npts, q ); + task.d2bfxz = util::sycl_malloc( nbf * npts, q ); + task.d2bfyy = util::sycl_malloc( nbf * npts, q ); + task.d2bfyz = util::sycl_malloc( nbf * npts, q ); + task.d2bfzz = util::sycl_malloc( nbf * npts, q ); + } + + if(pop_lapl) { + task.d2bflapl = util::sycl_malloc( nbf * npts, q ); + } + + if(pop_lapl_grad) { + task.d3bflapl_x = util::sycl_malloc( nbf * npts, q ); + task.d3bflapl_y = util::sycl_malloc( nbf * npts, q ); + task.d3bflapl_z = util::sycl_malloc( nbf * npts, q ); + } + + auto* pts_x_device = task.points_x; + auto* pts_y_device = task.points_y; + auto* pts_z_device = task.points_z; + auto* offs_device = task.bfn_screening.shell_offs; + auto* mask_device = task.bfn_screening.shell_list; + + + std::vector pts_x, pts_y, pts_z; + for( auto pt : pts ) { + pts_x.emplace_back(pt[0]); + pts_y.emplace_back(pt[1]); + pts_z.emplace_back(pt[2]); + } + util::sycl_copy( npts, pts_x_device, pts_x.data(), q ); + util::sycl_copy( npts, pts_y_device, pts_y.data(), q ); + util::sycl_copy( npts, pts_z_device, pts_z.data(), q ); + + std::vector mask_ul( mask.size() ); + std::copy( mask.begin(), mask.end(), mask_ul.begin() ); + util::sycl_copy( mask.size(), mask_device, mask_ul.data(), q ); + + std::vector offs( mask.size() ); + offs[0] = 0; + for( int i = 1; i < mask.size(); ++i ) + offs[i] = offs[i-1] + basis[mask[i-1]].size(); + util::sycl_copy( offs.size(), offs_device, offs.data(), q ); + + } + + return std::pair(shells_device,tasks); +} + + +void sycl_check_collocation( const std::vector& tasks, + const std::vector& ref_data, + bool check_grad, bool check_hess, bool check_lapl, bool check_lapl_grad, + ::sycl::queue& q) { + + for( int i = 0; i < tasks.size(); i++ ) { + + auto* ref_eval = ref_data[i].eval.data(); + std::vector eval (tasks[i].bfn_screening.nbe * tasks[i].npts); + util::sycl_copy( eval.size(), eval.data(), tasks[i].bf, q ); + + check_collocation_transpose( tasks[i].npts, tasks[i].bfn_screening.nbe, ref_eval, + eval.data(), "IT = " + std::to_string(i) + " BF EVAL" ); + + if( check_grad ) { + auto* ref_deval_x = ref_data[i].deval_x.data(); + auto* ref_deval_y = ref_data[i].deval_y.data(); + auto* ref_deval_z = ref_data[i].deval_z.data(); + + std::vector deval_x (tasks[i].bfn_screening.nbe * tasks[i].npts); + std::vector deval_y (tasks[i].bfn_screening.nbe * tasks[i].npts); + std::vector deval_z (tasks[i].bfn_screening.nbe * tasks[i].npts); + + util::sycl_copy( eval.size(), deval_x.data(), tasks[i].dbfx, q ); + util::sycl_copy( eval.size(), deval_y.data(), tasks[i].dbfy, q ); + util::sycl_copy( eval.size(), deval_z.data(), tasks[i].dbfz, q ); + + auto npts = tasks[i].npts; + auto nbe = tasks[i].bfn_screening.nbe; + check_collocation_transpose( npts, nbe, ref_deval_x, deval_x.data(), "IT = " + std::to_string(i) + " BFX EVAL" ); + check_collocation_transpose( npts, nbe, ref_deval_y, deval_y.data(), "IT = " + std::to_string(i) + " BFY EVAL" ); + check_collocation_transpose( npts, nbe, ref_deval_z, deval_z.data(), "IT = " + std::to_string(i) + " BFZ EVAL" ); + } + + if( check_hess ) { + auto* ref_d2eval_xx = ref_data[i].d2eval_xx.data(); + auto* ref_d2eval_xy = ref_data[i].d2eval_xy.data(); + auto* ref_d2eval_xz = ref_data[i].d2eval_xz.data(); + auto* ref_d2eval_yy = ref_data[i].d2eval_yy.data(); + auto* ref_d2eval_yz = ref_data[i].d2eval_yz.data(); + auto* ref_d2eval_zz = ref_data[i].d2eval_zz.data(); + + std::vector d2eval_xx (tasks[i].bfn_screening.nbe * tasks[i].npts); + std::vector d2eval_xy (tasks[i].bfn_screening.nbe * tasks[i].npts); + std::vector d2eval_xz (tasks[i].bfn_screening.nbe * tasks[i].npts); + std::vector d2eval_yy (tasks[i].bfn_screening.nbe * tasks[i].npts); + std::vector d2eval_yz (tasks[i].bfn_screening.nbe * tasks[i].npts); + std::vector d2eval_zz (tasks[i].bfn_screening.nbe * tasks[i].npts); + + util::sycl_copy( eval.size(), d2eval_xx.data(), tasks[i].d2bfxx, q ); + util::sycl_copy( eval.size(), d2eval_xy.data(), tasks[i].d2bfxy, q ); + util::sycl_copy( eval.size(), d2eval_xz.data(), tasks[i].d2bfxz, q ); + util::sycl_copy( eval.size(), d2eval_yy.data(), tasks[i].d2bfyy, q ); + util::sycl_copy( eval.size(), d2eval_yz.data(), tasks[i].d2bfyz, q ); + util::sycl_copy( eval.size(), d2eval_zz.data(), tasks[i].d2bfzz, q ); + + auto npts = tasks[i].npts; + auto nbe = tasks[i].bfn_screening.nbe; + check_collocation_transpose( npts, nbe, ref_d2eval_xx, d2eval_xx.data(), "IT = " + std::to_string(i) + " BFXX EVAL" ); + check_collocation_transpose( npts, nbe, ref_d2eval_xy, d2eval_xy.data(), "IT = " + std::to_string(i) + " BFXY EVAL" ); + check_collocation_transpose( npts, nbe, ref_d2eval_xz, d2eval_xz.data(), "IT = " + std::to_string(i) + " BFXZ EVAL" ); + check_collocation_transpose( npts, nbe, ref_d2eval_yy, d2eval_yy.data(), "IT = " + std::to_string(i) + " BFYY EVAL" ); + check_collocation_transpose( npts, nbe, ref_d2eval_yz, d2eval_yz.data(), "IT = " + std::to_string(i) + " BFYZ EVAL" ); + check_collocation_transpose( npts, nbe, ref_d2eval_zz, d2eval_zz.data(), "IT = " + std::to_string(i) + " BFZZ EVAL" ); + } + + if( check_lapl ) { + auto npts = tasks[i].npts; + auto nbe = tasks[i].bfn_screening.nbe; + auto* ref_d2eval_lapl = ref_data[i].d2eval_lapl.data(); + std::vector d2eval_lapl(npts * nbe); + util::sycl_copy(eval.size(), d2eval_lapl.data(), tasks[i].d2bflapl, q); + check_collocation_transpose(npts, nbe, ref_d2eval_lapl, d2eval_lapl.data(), "IT = " + std::to_string(i) + "BFLAPL EVAL" ); + } + + if( check_lapl_grad ) { + auto npts = tasks[i].npts; + auto nbe = tasks[i].bfn_screening.nbe; + auto* ref_d3eval_lapl_x = ref_data[i].d3eval_lapl_x.data(); + auto* ref_d3eval_lapl_y = ref_data[i].d3eval_lapl_y.data(); + auto* ref_d3eval_lapl_z = ref_data[i].d3eval_lapl_z.data(); + std::vector d3eval_lapl_x(npts * nbe); + std::vector d3eval_lapl_y(npts * nbe); + std::vector d3eval_lapl_z(npts * nbe); + util::sycl_copy(eval.size(), d3eval_lapl_x.data(), tasks[i].d3bflapl_x, q); + util::sycl_copy(eval.size(), d3eval_lapl_y.data(), tasks[i].d3bflapl_y, q); + util::sycl_copy(eval.size(), d3eval_lapl_z.data(), tasks[i].d3bflapl_z, q); + check_collocation_transpose(npts, nbe, ref_d3eval_lapl_x, d3eval_lapl_x.data(), "IT = " + std::to_string(i) + "BFLAPL_X EVAL" ); + check_collocation_transpose(npts, nbe, ref_d3eval_lapl_y, d3eval_lapl_y.data(), "IT = " + std::to_string(i) + "BFLAPL_Y EVAL" ); + check_collocation_transpose(npts, nbe, ref_d3eval_lapl_z, d3eval_lapl_z.data(), "IT = " + std::to_string(i) + "BFLAPL_Z EVAL" ); + } + + } + +} + + +void test_sycl_collocation_masked_combined( const BasisSet& basis, const std::string& filename, bool grad ) { + + + + std::vector ref_data; + read_collocation_data(ref_data, filename); + + + device_queue stream( std::make_shared() ); + // USM alloc/copy are context-bound, so the raw sycl::queue& (as opposed + // to the type-erased device_queue) is needed for the host-side setup below. + ::sycl::queue& q = stream.queue_as().queue; + auto [shells_device,tasks] = populate_device_sycl( basis, ref_data, grad, false, false, false, q ); + + + const auto nshells_max = std::max_element( tasks.begin(), tasks.end(), + []( const auto& a, const auto& b ) { + return a.bfn_screening.nshells < b.bfn_screening.nshells; + })->bfn_screening.nshells; + + const auto npts_max = std::max_element( tasks.begin(), tasks.end(), + []( const auto& a, const auto& b ) { + return a.npts < b.npts; + })->npts; + + auto* tasks_device = util::sycl_malloc( tasks.size(), q ); + util::sycl_copy( tasks.size(), tasks_device, tasks.data(), q ); + + if(grad) + eval_collocation_masked_combined_deriv1( tasks.size(), npts_max, + nshells_max, shells_device, tasks_device, stream ); + else + eval_collocation_masked_combined( tasks.size(), npts_max, + nshells_max, shells_device, tasks_device, stream ); + + util::sycl_device_sync(q); + + sycl_check_collocation( tasks, ref_data, grad, false, false, false, q ); + + + for( auto& t : tasks ) { + util::sycl_free( q, t.points_x, t.points_y, t.points_z, t.bfn_screening.shell_offs, t.bfn_screening.shell_list, t.bf ); + if(grad) util::sycl_free( q, t.dbfx, t.dbfy, t.dbfz ); + } + util::sycl_free( q, tasks_device, shells_device ); +} + +void test_sycl_collocation( const BasisSet& basis, + const std::string& filename ) { + + test_sycl_collocation_masked_combined( basis, filename, false ); + +} + +void test_sycl_collocation_deriv1( const BasisSet& basis, + const std::string& filename ) { + + test_sycl_collocation_masked_combined( basis, filename, true ); + +} + + +void test_sycl_collocation_shell_to_task( const BasisSet& basis, const BasisSetMap& basis_map, + const std::string& filename, bool grad, bool hess, bool lapl, bool lapl_grad) { + + // Load reference data + std::vector ref_data; + read_collocation_data(ref_data, filename); + + // Populate base task information + device_queue stream( std::make_shared() ); + ::sycl::queue& q = stream.queue_as().queue; + auto [shells_device,tasks] = populate_device_sycl( basis, ref_data, grad, hess, lapl, lapl_grad, q ); + + // Send tasks to device + auto* tasks_device = util::sycl_malloc( tasks.size(), q ); + util::sycl_copy( tasks.size(), tasks_device, tasks.data(), q ); + + + // Form Shell -> Task data structures + std::vector< std::vector > + shell_to_task_idx( basis.size() ), + shell_to_task_off( basis.size() ); + + int itask = 0; + for( auto& d : ref_data ) { + const auto& mask = d.mask; + + // Reform offsets + std::vector offs( mask.size() ); + offs[0] = 0; + for( int i = 1; i < mask.size(); ++i ) + offs[i] = offs[i-1] + basis[mask[i-1]].size(); + + // Form shell -> task + for( auto i = 0; i < mask.size(); ++i ) { + auto ish = mask[i]; + shell_to_task_idx[ish].emplace_back(itask); + shell_to_task_off[ish].emplace_back( offs[i] ); + } + itask++; + + } + + std::vector shell_to_task; + for( auto ish = 0; ish < basis.size(); ++ish ) { + shell_to_task.emplace_back(); + + const auto ntask = shell_to_task_idx[ish].size(); + shell_to_task.back().ntask = ntask; + shell_to_task.back().shell_device = shells_device + ish; + shell_to_task.back().task_idx_device = util::sycl_malloc( ntask, q ); + shell_to_task.back().task_shell_offs_device = + util::sycl_malloc( ntask, q ); + + util::sycl_copy( ntask, shell_to_task.back().task_idx_device, + shell_to_task_idx[ish].data(), q ); + util::sycl_copy( ntask, shell_to_task.back().task_shell_offs_device, + shell_to_task_off[ish].data(), q ); + + } + + + // Sort shells by L + std::vector shell_idx( basis.size() ); + std::iota( shell_idx.begin(), shell_idx.end(), 0 ); + + std::sort( shell_idx.begin(), shell_idx.end(), + [&]( auto i, auto j ){ return basis.at(i).l() < basis.at(j).l(); } ); + + { + std::vector shell_to_task_sorted( basis.size() ); + for( auto i = 0; i < basis.size(); ++i ) + shell_to_task_sorted[i] = shell_to_task[shell_idx[i]]; + shell_to_task = std::move(shell_to_task_sorted); + } + + + // Send Shell -> Task to device + auto* shell_to_task_device = util::sycl_malloc(basis.size(), q); + util::sycl_copy( basis.size(), shell_to_task_device, shell_to_task.data(), q ); + util::sycl_device_sync(q); + + // Form angular momentum batches for collocation eval + auto max_l = std::max_element(basis.begin(),basis.end(), + [](const auto&a, const auto& b){ return a.l() < b.l(); } )->l(); + std::vector l_batched_shell_to_task(max_l+1); + { + auto* p = shell_to_task_device; + auto* h = shell_to_task.data(); + for( auto l = 0; l <= max_l; ++l ) { + auto nsh = basis_map.nshells_with_l(l); + auto pure = basis_map.l_purity(l); + l_batched_shell_to_task[l].nshells_in_batch = nsh; + l_batched_shell_to_task[l].pure = pure; + l_batched_shell_to_task[l].shell_to_task_device = p; + + size_t total_ntask = std::accumulate( h, h + nsh, 0ul, + [](auto& a, auto& b){ return a + b.ntask; } ); + l_batched_shell_to_task[l].ntask_average = total_ntask / nsh; + + p += nsh; + h += nsh; + } + } + + + if( lapl_grad ) + eval_collocation_shell_to_task_lapgrad( max_l, l_batched_shell_to_task.data(), + tasks_device, stream ); + else if( hess ) + eval_collocation_shell_to_task_hessian( max_l, l_batched_shell_to_task.data(), + tasks_device, stream ); + else if( lapl ) + eval_collocation_shell_to_task_laplacian( max_l, l_batched_shell_to_task.data(), + tasks_device, stream ); + else if( grad ) + eval_collocation_shell_to_task_gradient( max_l, l_batched_shell_to_task.data(), + tasks_device, stream ); + else + eval_collocation_shell_to_task( max_l, l_batched_shell_to_task.data(), + tasks_device, stream ); + + + + util::sycl_device_sync(q); + sycl_check_collocation( tasks, ref_data, grad, hess, lapl, lapl_grad, q ); + + + for( auto& t : tasks ) { + util::sycl_free( q, t.points_x, t.points_y, t.points_z, t.bfn_screening.shell_offs, t.bfn_screening.shell_list, t.bf ); + if(grad) util::sycl_free( q, t.dbfx, t.dbfy, t.dbfz ); + if(hess) util::sycl_free( q, t.d2bfxx, t.d2bfxy, t.d2bfxz, t.d2bfyy, t.d2bfyz, t.d2bfzz ); + if(lapl) util::sycl_free( q, t.d2bflapl ); + if(lapl_grad) util::sycl_free( q, t.d3bflapl_x, t.d3bflapl_y, t.d3bflapl_z ); + } + util::sycl_free( q, tasks_device, shells_device, shell_to_task_device ); + for( auto& s : shell_to_task ) { + util::sycl_free( q, s.task_idx_device, s.task_shell_offs_device ); + } +} + + + +void test_sycl_collocation_shell_to_task( const BasisSet& basis, + const BasisSetMap& basis_map, const std::string& filename) { + + test_sycl_collocation_shell_to_task(basis,basis_map,filename,false, false, false, false); + +} +void test_sycl_collocation_shell_to_task_gradient( const BasisSet& basis, + const BasisSetMap& basis_map, const std::string& filename) { + + test_sycl_collocation_shell_to_task(basis,basis_map,filename,true, false, false, false); + +} +void test_sycl_collocation_shell_to_task_hessian( const BasisSet& basis, + const BasisSetMap& basis_map, const std::string& filename) { + + test_sycl_collocation_shell_to_task(basis,basis_map,filename,true, true, false, false); + +} + +void test_sycl_collocation_shell_to_task_laplacian( const BasisSet& basis, + const BasisSetMap& basis_map, const std::string& filename) { + + test_sycl_collocation_shell_to_task(basis,basis_map,filename,true, false, true, false); + +} + +void test_sycl_collocation_shell_to_task_lapgrad( const BasisSet& basis, + const BasisSetMap& basis_map, const std::string& filename) { + + test_sycl_collocation_shell_to_task(basis,basis_map,filename,true, true, true, true); + +} + +#endif // GAUXC_HAS_SYCL diff --git a/tests/ref_data/test_mol_c.hdf5 b/tests/ref_data/test_mol_c.hdf5 new file mode 100644 index 000000000..053775327 Binary files /dev/null and b/tests/ref_data/test_mol_c.hdf5 differ diff --git a/tests/standalone_driver.cxx b/tests/standalone_driver.cxx index f4af33756..a23bf662c 100644 --- a/tests/standalone_driver.cxx +++ b/tests/standalone_driver.cxx @@ -855,4 +855,5 @@ int main(int argc, char** argv) { MPI_Finalize(); #endif + return 0; } diff --git a/tests/ut_main.cxx b/tests/ut_main.cxx index 754205156..35e42e267 100644 --- a/tests/ut_main.cxx +++ b/tests/ut_main.cxx @@ -27,6 +27,10 @@ int main( int argc, char* argv[] ) { MPI_Comm_rank(MPI_COMM_WORLD, &rank); #ifdef GAUXC_HAS_CUDA cudaSetDevice(rank); +#elif defined(GAUXC_HAS_SYCL) + // SYCL has no per-process cudaSetDevice equivalent: device visibility for + // rank-to-GPU binding is controlled externally via ZE_AFFINITY_MASK / + // ONEAPI_DEVICE_SELECTOR, so there is nothing to do here. #endif int result = Catch::Session().run( argc, argv ); MPI_Finalize(); diff --git a/tests/weights.cxx b/tests/weights.cxx index 568822686..ba30fb28d 100644 --- a/tests/weights.cxx +++ b/tests/weights.cxx @@ -21,6 +21,7 @@ #include "weights_host.hpp" #include "weights_cuda.hpp" #include "weights_hip.hpp" +#include "weights_sycl.hpp" //#define GENERATE_TESTS TEST_CASE( "Partition Weights", "[weights]" ) { @@ -71,6 +72,8 @@ TEST_CASE( "Partition Weights", "[weights]" ) { test_cuda_weights( ref_file ); #elif defined(GAUXC_HAS_HIP) test_hip_weights( ref_file ); +#elif defined(GAUXC_HAS_SYCL) + test_sycl_weights( ref_file ); #endif } #endif diff --git a/tests/weights_sycl.hpp b/tests/weights_sycl.hpp new file mode 100644 index 000000000..41136ac8d --- /dev/null +++ b/tests/weights_sycl.hpp @@ -0,0 +1,121 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once +#include "weights_generate.hpp" +#include "hdf5_test_serialization.hpp" +#include "hdf5_test_serialization_impl.hpp" +#include +#include + +#ifdef GAUXC_HAS_SYCL +#include "device_specific/sycl_util.hpp" +#include "device/sycl/sycl_aos_scheme1.hpp" +#include "device/sycl/sycl_aos_scheme1_weights.hpp" + + + + +void test_sycl_weights( const std::string& filename ) { + + ref_weights_data ref_data; + read_weights_data(ref_data, filename); + + //std::vector< std::array > points; + std::vector< double > points_x, points_y, points_z; + std::vector< double > weights, weights_ref; + std::vector< double > dist_nearest; + std::vector< int32_t > iparent; + + for( auto& task : ref_data.tasks_unm ) { + for( auto pt : task.points ) { + points_x.emplace_back(pt[0]); + points_y.emplace_back(pt[1]); + points_z.emplace_back(pt[2]); + } + weights.insert( weights.end(), + task.weights.begin(), + task.weights.end() ); + + size_t npts = task.points.size(); + dist_nearest.insert( dist_nearest.end(), npts, + task.dist_nearest ); + iparent.insert( iparent.end(), npts, task.iParent ); + } + + for( auto& task : ref_data.tasks_mod ) { + weights_ref.insert( weights_ref.end(), + task.weights.begin(), + task.weights.end() ); + } + + size_t npts = points_x.size(); + size_t natoms = ref_data.mol.natoms(); + + constexpr auto weight_unroll = alg_constants::SyclAoSScheme1::weight_unroll; + + size_t LDatoms = util::div_ceil( natoms, weight_unroll ) * weight_unroll; + + std::vector< double > coords( 3 * natoms ); + for( auto iat = 0 ; iat < natoms; ++iat ) { + coords[ 3*iat + 0 ] = ref_data.mol.at(iat).x; + coords[ 3*iat + 1 ] = ref_data.mol.at(iat).y; + coords[ 3*iat + 2 ] = ref_data.mol.at(iat).z; + } + + // SYCL USM allocation is context-bound, so every alloc/copy/free below + // needs the queue that owns the allocations. + util::sycl_queue stream; + ::sycl::queue& q = stream.queue; + + //auto* points_d = util::sycl_malloc( 3*npts, q ); + auto* points_x_d = util::sycl_malloc( npts, q ); + auto* points_y_d = util::sycl_malloc( npts, q ); + auto* points_z_d = util::sycl_malloc( npts, q ); + auto* weights_d = util::sycl_malloc( npts , q ); + auto* iparent_d = util::sycl_malloc( npts , q ); + auto* distnea_d = util::sycl_malloc( npts , q ); + auto* rab_d = util::sycl_malloc( natoms*LDatoms, q ); + auto* coords_d = util::sycl_malloc( 3*natoms, q ); + auto* dist_scr_d= util::sycl_malloc( npts*LDatoms, q ); + + //util::sycl_copy( 3*npts, points_d, points.data()->data(), q ); + util::sycl_copy( npts, points_x_d, points_x.data(), q ); + util::sycl_copy( npts, points_y_d, points_y.data(), q ); + util::sycl_copy( npts, points_z_d, points_z.data(), q ); + util::sycl_copy( npts, weights_d, weights.data(), q ); + util::sycl_copy( npts, iparent_d, iparent.data(), q ); + util::sycl_copy( npts, distnea_d, dist_nearest.data(), q ); + + std::vector rab_inv(natoms*natoms); + for( auto i = 0; i < natoms*natoms; ++i ) + rab_inv[i] = 1./ref_data.meta->rab().data()[i]; + + util::sycl_copy_2d( rab_d, LDatoms * sizeof(double), + rab_inv.data(), natoms * sizeof(double), + natoms * sizeof(double), natoms, q, "RAB H2D"); + + util::sycl_copy( 3*natoms, coords_d, coords.data(), q ); + + sycl_aos_scheme1_weights_wrapper( npts, natoms, points_x_d, points_y_d, points_z_d, rab_d, + LDatoms, coords_d, dist_scr_d, LDatoms, iparent_d, distnea_d, + weights_d, q ); + + util::sycl_device_sync(q); + util::sycl_copy( npts, weights.data(), weights_d, q ); + util::sycl_free( q, points_x_d, points_y_d, points_z_d, weights_d, iparent_d, distnea_d, + rab_d, coords_d, dist_scr_d ); + + for( auto i = 0ul; i < npts; ++i ) + CHECK( weights.at(i) == Approx( weights_ref.at(i) ) ); + +} +#endif diff --git a/tests/xc_integrator.cxx b/tests/xc_integrator.cxx index 08f4ca214..1a684b88b 100644 --- a/tests/xc_integrator.cxx +++ b/tests/xc_integrator.cxx @@ -339,6 +339,12 @@ void test_integrator(std::string reference_file, functional_type& func, PruningS #ifdef GAUXC_HAS_HIP check_grad = false; check_k = false; + #elif defined(GAUXC_HAS_SYCL) + // The Obara-Saika GPU integrals backing EXX/snK are CUDA-only + // (GAUXC_ENABLE_EXX is only defined under GAUXC_HAS_CUDA in + // scheme1_base.cxx), so K cannot be checked on SYCL. Gradients have no + // such restriction and remain enabled. + check_k = false; #endif SECTION( "Incore - MPI Reduction" ) { test_xc_integrator( ExecutionSpace::Device, rt,