diff --git a/.github/workflows/sycl-cpu-ci.yml b/.github/workflows/sycl-cpu-ci.yml new file mode 100644 index 000000000..1bd257c0a --- /dev/null +++ b/.github/workflows/sycl-cpu-ci.yml @@ -0,0 +1,366 @@ +name: SYCL CPU CI + +# Sanity-verifies the SYCL port on a CPU-only GitHub-hosted runner using the +# Intel OpenCL CPU device (no GPU present, none needed). dpctl and dpnp are +# built from source at their latest upstream (IntelPython) commits against the +# same oneAPI toolkit, so the SYCL ABI stays in lockstep with what +# gpu4pyscf/lib expects. +# +# NOTE: gpu4pyscf's local dev setup uses a sparse-solvers dpnp fork +# (abagusetty/dpnp, feature-sparse-linalg-solvers). This CI deliberately +# builds upstream IntelPython/dpnp instead, so a couple of sparse-dependent +# tests may fail here — that is acceptable for this CI. +# +# Production builds AOT-compile for Intel PVC; that cannot run here, so the +# build passes -DGPU4PYSCF_SYCL_AOT=OFF (see gpu4pyscf/lib/CMakeLists.txt) to +# get a generic spir64 JIT build that the OpenCL CPU device can execute. +# +# Triggers: +# - workflow_dispatch (Actions tab) +# - every push and every pull request, on any branch +# - an issue/PR comment containing "@sycl-runner" +# NOTE: GitHub only delivers issue_comment events to the copy of this file +# on the repository's DEFAULT branch. For the comment trigger to work, +# this workflow file must also exist on the default branch; a comment on a +# PR then builds that PR's head ref. +# +# The full test suite will not finish on a 4-core runner and is not the goal; +# the smoke steps below are allowed to partially fail (continue-on-error) so a +# broken kernel build or import regression is still caught early. + +on: + workflow_dispatch: {} + # Run on every push and every pull request, the way lint.yml does, so the + # SYCL build is checked on each commit rather than only on the sycl branch. + # unittest.yml restricts itself to master because its runners are costly; + # this job is a build plus smoke test on a stock GitHub runner. + push: + pull_request: + issue_comment: + types: [created] + +permissions: + # contents: write lets the always-run log step push build/test logs to a + # ci-logs- branch, since API log download needs repo admin auth. + contents: write + issues: write + pull-requests: write + +concurrency: + group: sycl-cpu-ci-${{ github.ref }}-${{ github.event_name }} + # Do NOT cancel in progress. A cold run spends ~45 min building dpctl/dpnp; + # cancelling it on the next push means that work is thrown away and the + # caches it would have populated are never written, so the following run + # starts cold too. Several pushes in quick succession could then leave the + # stack permanently unbuilt. Queue instead. + cancel-in-progress: false + +env: + # Only the OpenCL CPU device may be visible to SYCL. + ONEAPI_DEVICE_SELECTOR: "opencl:cpu" + OMP_NUM_THREADS: "2" + SYCL_CACHE_PERSISTENT: "1" + SYCL_CACHE_DIR: "/tmp/sycl_cache" + CCACHE_DIR: "/home/runner/.ccache" + CCACHE_MAXSIZE: "6G" + VENV: "/home/runner/sycl-venv" + # The SYCL port is used via PYTHONPATH, not pip-installed: `pip install .` + # would run the CUDA setup.py, and the vendored cupy/cupyx shims live under + # /gpu4pyscf/, which must itself be on the path (same layout the + # production Aurora environment uses). + PYTHONPATH: "${{ github.workspace }}:${{ github.workspace }}/gpu4pyscf" + +jobs: + build-and-test: + # For comment events, only react to "@sycl-runner". + if: github.event_name != 'issue_comment' || contains(github.event.comment.body, '@sycl-runner') + runs-on: ubuntu-24.04 + timeout-minutes: 355 + + steps: + - name: Acknowledge "@sycl-runner" comment + if: github.event_name == 'issue_comment' + continue-on-error: true + uses: actions/github-script@v7 + with: + script: | + await github.rest.reactions.createForIssueComment({ + owner: context.repo.owner, + repo: context.repo.repo, + comment_id: context.payload.comment.id, + content: 'rocket', + }); + + - name: Checkout gpu4pyscf + uses: actions/checkout@v4 + with: + # push / pull_request / workflow_dispatch build the ref that + # triggered them. A comment on a PR builds that PR's head ref, so + # "@sycl-runner" tests the branch under discussion rather than a + # fixed one; a comment on a plain issue has no PR to resolve, so it + # falls back to the default branch. + ref: ${{ github.event_name == 'issue_comment' && github.event.issue.pull_request && format('refs/pull/{0}/head', github.event.issue.number) || '' }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install Intel oneAPI toolchain (compilers, MKL, TBB) and build tools + run: | + wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \ + | gpg --dearmor | sudo tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null + echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \ + | sudo tee /etc/apt/sources.list.d/oneAPI.list + sudo apt-get update + sudo apt-get install -y \ + intel-oneapi-compiler-dpcpp-cpp \ + intel-oneapi-compiler-fortran \ + intel-oneapi-mkl-devel \ + intel-oneapi-tbb-devel \ + intel-oneapi-libdpstd-devel \ + intel-oneapi-runtime-opencl \ + intel-oneapi-runtime-dpcpp-cpp \ + ocl-icd-opencl-dev clinfo \ + lld ninja-build cmake ccache + + - name: Ensure an OpenCL CPU device exists + run: | + source /opt/intel/oneapi/setvars.sh --force + # The dpcpp runtime package ships the Intel CPU OpenCL runtime; if + # for some reason no CPU device is registered, fall back to the + # standalone oclcpuexp release the intel/llvm project publishes + # (same runtime dpctl's own CI uses). + if ! sycl-ls | grep -qi "opencl:cpu"; then + OCLCPUEXP_FN="oclcpuexp-2025.21.10.0.10_160000_rel.tar.gz" + DRIVER_PATH="2025-WW45" + wget -q "https://github.com/intel/llvm/releases/download/${DRIVER_PATH}/${OCLCPUEXP_FN}" \ + -O /tmp/oclcpuexp.tar.gz + sudo mkdir -p /opt/intel/oclcpuexp + sudo tar -xzf /tmp/oclcpuexp.tar.gz -C /opt/intel/oclcpuexp + sudo mkdir -p /etc/OpenCL/vendors + echo "/opt/intel/oclcpuexp/x64/libintelocl.so" | sudo tee /etc/OpenCL/vendors/intel_expcpu.icd + fi + icpx --version + sycl-ls + sycl-ls | grep -qi "opencl:cpu" + + - name: Resolve dpctl/dpnp commits and cache keys + id: stack + run: | + # Keying on upstream HEAD meant any commit to IntelPython/dpctl or + # dpnp invalidated the cache and forced a ~45 min rebuild, often on + # a run that changed nothing here. Bucket by ISO week instead: the + # stack is rebuilt at most once a week, and the SHAs resolved inside + # that bucket are recorded so a run is still reproducible. + echo "dpctl_sha=$(git ls-remote https://github.com/IntelPython/dpctl HEAD | cut -f1)" >> "$GITHUB_OUTPUT" + echo "dpnp_sha=$(git ls-remote https://github.com/IntelPython/dpnp HEAD | cut -f1)" >> "$GITHUB_OUTPUT" + echo "oneapi_ver=$(dpkg-query -W -f='${Version}' intel-oneapi-compiler-dpcpp-cpp)" >> "$GITHUB_OUTPUT" + echo "week=$(date -u +%G-W%V)" >> "$GITHUB_OUTPUT" + + - name: Restore dpctl/dpnp stack cache (venv + source builds) + id: stack-cache + # Explicit restore/save split: plain actions/cache only saves when the + # whole job succeeds, so iterating on later steps kept losing the + # ~45 min dpnp build. The save step below runs right after the builds. + uses: actions/cache/restore@v4 + with: + path: | + /home/runner/sycl-venv + /home/runner/dpctl + /home/runner/dpnp + key: sycl-stack-v2-${{ runner.os }}-py311-oneapi${{ steps.stack.outputs.oneapi_ver }}-${{ steps.stack.outputs.week }} + # Fall back to any earlier week on the same oneAPI: a stale but + # working dpctl/dpnp beats a 45 min rebuild, and the weekly key + # above still refreshes it once the bucket rolls over. + restore-keys: | + sycl-stack-v2-${{ runner.os }}-py311-oneapi${{ steps.stack.outputs.oneapi_ver }}- + sycl-stack-v2-${{ runner.os }}-py311- + + - name: Create venv with Python build/runtime dependencies + if: steps.stack-cache.outputs.cache-matched-key == '' + run: | + python -m venv "$VENV" + "$VENV/bin/pip" install --upgrade pip + "$VENV/bin/pip" install numpy cython setuptools wheel scikit-build cmake ninja \ + versioneer pybind11 pytest pytest-timeout \ + pyscf pyscf-dispersion h5py geometric + + # Restored before the dpctl/dpnp builds, not just before the gpu4pyscf + # one: those two are now routed through ccache as well, and they are the + # expensive half of a cold run. + - name: Restore ccache (C++/SYCL objects, incl. dpctl/dpnp) + uses: actions/cache/restore@v4 + with: + path: /home/runner/.ccache + key: sycl-ccache-${{ runner.os }}-${{ github.sha }} + restore-keys: | + sycl-ccache-${{ runner.os }}- + + - name: Build dpctl from source (IntelPython/dpctl master) + if: steps.stack-cache.outputs.cache-matched-key == '' + run: | + source /opt/intel/oneapi/setvars.sh --force + source "$VENV/bin/activate" + set -o pipefail; mkdir -p /tmp/logs + # blob-less partial clone: full history+tags (versioneer needs `git + # describe` to produce a version satisfying dpnp's dpctl>=0.23.0dev0 + # pin; a shallow clone yields 0+untagged, which pip rejects). + git clone --filter=blob:none https://github.com/IntelPython/dpctl /home/runner/dpctl + cd /home/runner/dpctl + # Route the compiler through ccache. The stack cache above is keyed + # by week, so when it does roll over most translation units are + # unchanged from the previous build and come straight from ccache + # rather than being recompiled. + { python scripts/build_locally.py --oneapi \ + --cmake-opts="-DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache" && + pip install --no-build-isolation --no-deps . && + cd /tmp && python -c "import dpctl; print('dpctl OK', dpctl.__version__)" ; } 2>&1 | tee /tmp/logs/20_dpctl_build.log + # NON-editable install on purpose: dpnp's Cython resolves dpctl.pxd + # by scanning sys.path for a real site-packages/dpctl/ directory; a + # PEP-660 editable install hides it behind an import hook. + + - name: Build dpnp from source (IntelPython/dpnp master) + if: steps.stack-cache.outputs.cache-matched-key == '' + run: | + source /opt/intel/oneapi/setvars.sh --force + source "$VENV/bin/activate" + set -o pipefail; mkdir -p /tmp/logs + git clone --filter=blob:none https://github.com/IntelPython/dpnp /home/runner/dpnp + cd /home/runner/dpnp + # ccache as above. dpnp is the expensive half of the stack (291 + # Ninja targets), so this is where it pays off. + { python scripts/build_locally.py --oneapi \ + --cmake-opts="-DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache" && + pip install --no-build-isolation --no-deps . && + cd /tmp && python -c "import dpnp; print('dpnp OK', dpnp.__version__)" ; } 2>&1 | tee /tmp/logs/30_dpnp_build.log + + - name: Save dpctl/dpnp stack cache + # always(): a cancelled or failed run has usually still built part of + # the stack, and saving it means the next run resumes rather than + # restarting. Mirrors the ccache save step below. + if: always() && steps.stack-cache.outputs.cache-matched-key == '' + uses: actions/cache/save@v4 + with: + path: | + /home/runner/sycl-venv + /home/runner/dpctl + /home/runner/dpnp + key: sycl-stack-v2-${{ runner.os }}-py311-oneapi${{ steps.stack.outputs.oneapi_ver }}-${{ steps.stack.outputs.week }} + + - name: Build gpu4pyscf (SYCL backend, JIT spir64, CPU-runnable) + env: + # For sub-builds (ExchCXX ExternalProject etc.) that would otherwise + # fall back to /usr/bin/cc, which cannot compile -fsycl code. + CC: icx + CXX: icpx + FC: ifx + CMAKE_CONFIGURE_ARGS: >- + -DCMAKE_CXX_COMPILER=icpx -DCMAKE_C_COMPILER=icx + -DCMAKE_Fortran_COMPILER=ifx -DUSE_SYCL=ON + -DCMAKE_BUILD_TYPE=Release + -DGPU4PYSCF_SYCL_AOT=OFF + -DCMAKE_C_COMPILER_LAUNCHER=ccache + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + run: | + source /opt/intel/oneapi/setvars.sh --force + source "$VENV/bin/activate" + set -o pipefail; mkdir -p /tmp/logs + python setup.py --sycl build 2>&1 | tee /tmp/logs/40_gpu4pyscf_build.log + ccache -s || true + + - name: Save ccache (even on failure, partial objects speed up retries) + if: always() + uses: actions/cache/save@v4 + with: + path: /home/runner/.ccache + key: sycl-ccache-${{ runner.os }}-${{ github.sha }} + + - name: Smoke test - imports and CPU SYCL device visibility + run: | + source /opt/intel/oneapi/setvars.sh --force + source "$VENV/bin/activate" + set -o pipefail; mkdir -p /tmp/logs + # The gpu4pyscf import is disabled until the dpnp sparse PR lands: + # cupyx/scipy/sparse/linalg.py imports dpnp.scipy.sparse.linalg, + # which does not exist in dpnp yet, so `import gpu4pyscf` fails with + # ModuleNotFoundError before reaching anything this step tests. + # Re-enable the third command once that PR is merged. + { python -c "import dpctl; print('dpctl', dpctl.__version__); [print(d) for d in dpctl.get_devices()]" && + python -c "import dpnp; a = dpnp.arange(10); assert int(a.sum()) == 45; print('dpnp', dpnp.__version__, 'OK')" ; + # python -c "import gpu4pyscf; print('gpu4pyscf import OK')" ; + } 2>&1 | tee /tmp/logs/50_smoke.log + + - name: scf unit tests (smoke, partial completion expected) + # Disabled: the SYCL backend still has known aborts and hangs in + # these suites, so a run says nothing useful about a given change. + # The build and smoke steps above are the meaningful gate for now. + # Flip this back to re-enable once the defects are fixed. + if: false + continue-on-error: true + timeout-minutes: 90 + run: | + source /opt/intel/oneapi/setvars.sh --force + source "$VENV/bin/activate" + set -o pipefail; mkdir -p /tmp/logs + python -m pytest gpu4pyscf/scf/tests/ -v -x --timeout=900 \ + --durations=25 --continue-on-collection-errors 2>&1 | tee /tmp/logs/60_pytest_scf.log + + - name: df unit tests (smoke, partial completion expected) + # Disabled: the SYCL backend still has known aborts and hangs in + # these suites, so a run says nothing useful about a given change. + # The build and smoke steps above are the meaningful gate for now. + # Flip this back to re-enable once the defects are fixed. + if: false + continue-on-error: true + timeout-minutes: 60 + run: | + source /opt/intel/oneapi/setvars.sh --force + source "$VENV/bin/activate" + set -o pipefail; mkdir -p /tmp/logs + python -m pytest gpu4pyscf/df/tests/ -v -x --timeout=900 \ + --durations=25 --continue-on-collection-errors 2>&1 | tee /tmp/logs/70_pytest_df.log + + - name: dft unit tests (smoke, partial completion expected) + # Disabled: the SYCL backend still has known aborts and hangs in + # these suites, so a run says nothing useful about a given change. + # The build and smoke steps above are the meaningful gate for now. + # Flip this back to re-enable once the defects are fixed. + if: false + continue-on-error: true + timeout-minutes: 60 + run: | + source /opt/intel/oneapi/setvars.sh --force + source "$VENV/bin/activate" + set -o pipefail; mkdir -p /tmp/logs + python -m pytest gpu4pyscf/dft/tests/ -v -x --timeout=900 \ + --durations=25 --continue-on-collection-errors 2>&1 | tee /tmp/logs/80_pytest_dft.log + + - name: Upload logs as a run artifact + if: always() + uses: actions/upload-artifact@v4 + with: + name: sycl-cpu-ci-logs-${{ github.run_id }}-${{ github.run_attempt }} + path: /tmp/logs/*.log + if-no-files-found: ignore + retention-days: 14 + + # Only a push to a branch in this repository has a writable + # GITHUB_TOKEN. On a pull request from a fork the token is read-only and + # github.repository names the base repository, so the push below is + # rejected with a 403. The artifact above covers every context; this + # step is the extra convenience of a browsable branch where it can run. + - name: Publish logs to a ci-logs branch (readable without API auth) + if: always() && github.event_name == 'push' + continue-on-error: true + run: | + mkdir -p /tmp/logs /tmp/ci-logs && cd /tmp/ci-logs + for f in /tmp/logs/*.log; do + [ -e "$f" ] && tail -n 1200 "$f" > "$(basename "$f")" + done + ls -la + git init -q -b main . + git config user.email "ci@github-runner" && git config user.name "sycl-cpu-ci logs" + git add -A + git commit -q -m "logs for run ${{ github.run_id }} (attempt ${{ github.run_attempt }}) sha ${{ github.sha }}" --allow-empty + git push -f "https://x-access-token:${{ secrets.GITHUB_TOKEN }}@github.com/${{ github.repository }}" HEAD:refs/heads/ci-logs-${{ github.run_id }} diff --git a/gpu4pyscf/cupy/__init__.py b/gpu4pyscf/cupy/__init__.py new file mode 100644 index 000000000..f8523a09c --- /dev/null +++ b/gpu4pyscf/cupy/__init__.py @@ -0,0 +1,924 @@ +""" +CuPy-compatibility facade over dpnp for gpu4pyscf. + +Exports a fake `cupy` module built from dpnp, patches CuPy-vs-dpnp API +differences in place, and sets up the `cupy.cuda` submodule (master +queue, Stream, Device, Event -- see cupy/cuda.py). + +Import order matters: `cuda.py` monkey-patches dpnp creation APIs to +inject sycl_queue=master. The aliases on `cupy_fake` must be bound +AFTER that patching or they'll reference the unwrapped originals -- +see the rebind section near the bottom of this file. + +Idempotency / single-load guarantee +----------------------------------- +This package can be reached through two dotted paths: + - `cupy` (because we install `cupy_fake` into sys.modules) + - `gpu4pyscf.cupy` (the real package path) +and similarly for the `.cuda` submodule. Without care, Python's import +machinery loads the file TWICE -- once per name -- producing two module +objects with two independent `_master_queues` registries and two +independent sets of wrappers on dpnp. + +Fix: at the end of first-time init we alias BOTH names in sys.modules +to the same cupy_fake / cuda objects. A subsequent `import gpu4pyscf.cupy` +(or `import cupy`, or any variant of `.cuda`) then short-circuits in +the import cache and does not re-execute the file. + +Namespace isolation +------------------- +`cupy_fake` is a *separate* `types.ModuleType` object, NOT the current +module. This matters: `setattr(cupy_fake, 'any', dpnp.any)` must only +populate the fake cupy namespace. If we used `sys.modules[__name__]` +as cupy_fake, the broad dpnp-attribute loop would overwrite Python +builtins (`any`, `max`, `sum`, `abs`, ...) in this module's globals, +breaking every function defined here that calls `any(generator)` -- +e.g. the numpy-einsum dispatcher. Keep cupy_fake separate. +""" +import os +import sys +import types +from abc import ABCMeta + +import numpy as np +import dpnp +from dpnp.dpnp_array import dpnp_array +import dpnp.tensor as dpt + + +# ===================================================================== +# Early short-circuit -- if the facade has already been built under +# another name, just re-alias sys.modules and return it. This handles +# the rare case where Python manages to execute this file a second +# time despite the end-of-file aliasing (e.g. reload, stale finder). +# ===================================================================== +_ALREADY_LOADED = None +for _candidate in ("cupy", "gpu4pyscf.cupy"): + _m = sys.modules.get(_candidate) + if _m is not None and getattr(_m, "__gpu4pyscf_cupy_facade__", False): + _ALREADY_LOADED = _m + break + +if _ALREADY_LOADED is not None: + # Redirect the current import to the pre-existing facade. + sys.modules[__name__] = _ALREADY_LOADED + sys.modules["cupy"] = _ALREADY_LOADED + sys.modules["gpu4pyscf.cupy"] = _ALREADY_LOADED + _existing_cuda = getattr(_ALREADY_LOADED, "cuda", None) + if _existing_cuda is not None: + sys.modules["cupy.cuda"] = _existing_cuda + sys.modules["gpu4pyscf.cupy.cuda"] = _existing_cuda +else: + # ================================================================= + # First-time initialization -- build the fake cupy module. + # ================================================================= + + # ----------------------------------------------------------------- + # cupy.ndarray alias -- callable with CuPy's memptr= kwarg, + # isinstance-compatible with dpnp arrays. + # ----------------------------------------------------------------- + def _resolve_dpnp_impl(): + try: + import dpnp.dpnp_array as _mod + return getattr(_mod, "dpnp_array", None) + except Exception: + return None + + + _DPNP_ARRAY_IMPL = _resolve_dpnp_impl() + + + class _CuPyNdarrayMeta(ABCMeta): + """Supports `cupy.ndarray(shape, dtype=..., memptr=buf.data)` -- + the `memptr=` kwarg is CuPy-specific; dpnp uses `buffer=` and + doesn't accept our raw CuPy MemoryPointer shim, so we unwrap it.""" + + def __call__(cls, shape, dtype=np.float64, memptr=None): + if memptr is not None and hasattr(memptr, "get_array"): + memptr = memptr.get_array() + if isinstance(shape, (tuple, list)): + shape = tuple(int(s) for s in shape) + else: + shape = (int(shape),) + if memptr is not None: + return dpnp.ndarray(shape, dtype=dtype, buffer=memptr) + return dpnp.ndarray(shape, dtype=dtype) + + def __instancecheck__(cls, obj): + if isinstance(obj, dpnp.ndarray): + return True + if _DPNP_ARRAY_IMPL and isinstance(obj, _DPNP_ARRAY_IMPL): + return True + return False + + def __subclasscheck__(cls, sub): + try: + bases = [dpnp.ndarray] + if _DPNP_ARRAY_IMPL: + bases.append(_DPNP_ARRAY_IMPL) + return any(issubclass(sub, b) for b in bases) + except TypeError: + return False + + + class _CuPyNdarray(dpnp.ndarray, metaclass=_CuPyNdarrayMeta): + """Alias type for CuPy ndarray over dpnp arrays.""" + pass + + + # ----------------------------------------------------------------- + # Build the fake cupy module -- SEPARATE from the current module so + # setattr doesn't pollute our globals. See the module docstring. + # Give it package attributes so `from . import cuda` style imports + # resolve correctly when the fake is looked up as `cupy`. + # ----------------------------------------------------------------- + cupy_fake = types.ModuleType("cupy") + cupy_fake.__package__ = "cupy" + cupy_fake.__path__ = [os.path.dirname(os.path.abspath(__file__))] + cupy_fake.__gpu4pyscf_cupy_facade__ = True # marker for early short-circuit + + cupy_fake.ndarray = _CuPyNdarray + cupy_fake.asnumpy = dpnp.asnumpy + + # einsum: dpnp can return a non-contiguous result where numpy and cupy + # both return a C-contiguous one. For 'lkz,lxpq->kxpqz' on + # (3,2,2) x (3,1,6,6) numpy gives strides (576,1152,96,16,8) while dpnp + # gives (16,0,192,32,8) -- correct values (agreement to 2e-16), but a + # strided view with a 0-stride broadcast axis. + # + # Upstream code relies on the numpy/cupy layout. pbc/gto/int1e.py feeds + # the result of contract() straight into hermi_triu(), which asserts + # mat.flags.c_contiguous and aborts the test otherwise. Normalise here so + # every caller sees the layout it would get on the CUDA backend. + def _einsum(*args, **kwargs): + out = dpnp.einsum(*args, **kwargs) + if isinstance(out, dpnp.ndarray) and not out.flags.c_contiguous: + out = dpnp.ascontiguousarray(out) + return out + + _einsum.__name__ = 'einsum' + _einsum.__doc__ = getattr(dpnp.einsum, '__doc__', None) + cupy_fake.einsum = _einsum + + + # ----------------------------------------------------------------- + # ndarray.dot(out=...) -- fix a shape-mismatch edge case CuPy permits + # but dpnp rejects. Guarded so a second execution is a no-op. + # ----------------------------------------------------------------- + if not getattr(dpnp.ndarray.dot, "__gpu4pyscf_patched__", False): + _original_ndarray_dot = dpnp.ndarray.dot + + def _ndarray_dot_method(self, b, out=None, _orig=_original_ndarray_dot): + if out is None: + return _orig(self, b, out=None) + + result = _orig(self, b, out=None) + if result.shape != out.shape: + if result.size == out.size: + result = result.squeeze() + if result.shape != out.shape: + result = result.reshape(out.shape) + else: + raise ValueError( + f"Cannot fit result {result.shape} into {out.shape}") + + out[:] = result + return out + + _ndarray_dot_method.__gpu4pyscf_patched__ = True + dpnp.ndarray.dot = _ndarray_dot_method + + cupy_fake.dot = dpnp.dot + + + # ----------------------------------------------------------------- + # ndarray.reshape(shape) -- CuPy/NumPy accept a bare ndarray as the + # shape arg (unpacked element-wise, e.g. `.reshape(cell.mesh)`), but + # dpnp's dpctl.tensor.reshape wraps a non-list/tuple shape into a + # 1-tuple instead of unpacking it, so `d` in its `operator.index(d)` + # loop ends up being the whole array -> TypeError. Normalise here. + # ----------------------------------------------------------------- + if not getattr(dpnp.ndarray.reshape, "__gpu4pyscf_patched__", False): + _original_ndarray_reshape = dpnp.ndarray.reshape + + def _ndarray_reshape_method(self, *shape, _orig=_original_ndarray_reshape, **kwargs): + if len(shape) == 1 and isinstance(shape[0], np.ndarray): + shape = tuple(int(d) for d in shape[0].tolist()) + return _orig(self, *shape, **kwargs) + + _ndarray_reshape_method.__gpu4pyscf_patched__ = True + dpnp.ndarray.reshape = _ndarray_reshape_method + + + # ----------------------------------------------------------------- + # Initial population of cupy_fake from dpnp (narrow, explicit list) + # ----------------------------------------------------------------- + for _attr in ( + "append", "max", "linalg", "concatenate", "zeros", "ones", + "empty", "eye", "view", "empty_like", "copyto", "cumsum", "any", "matmul", + "vstack", "full", "arange", "stack", "expand_dims", "unique", "double", + "sign", "argsort", "count_nonzero", "where", "split", "take", "tril", "log", + "complex128", "uint8", "int32", "int64", "float32", "float64", "ravel", + "random", "sum", "exp", "outer", "ix_", "pi", "square", "multiply", + "diag_indices", "repeat", "diag", "tril_indices_from", "ceil", "newaxis", + "ascontiguousarray", "nonzero", "array_equal", "isinf", "isnan", "dtype", + "asfortranarray", "abs", "shape", "argmax", "trace", "prod", + ): + try: + setattr(cupy_fake, _attr, getattr(dpnp, _attr)) + except AttributeError: + pass + + + # ----------------------------------------------------------------- + # cupy.random.seed compatibility + # + # CuPy accepts any array-like seed -- pyscf's own tests call + # cupy.random.seed(np.asarray(1, dtype=np.uint64)) -- while + # dpnp.random.seed only takes a plain scalar and otherwise fails with + # "Cannot construct a dtype from an array". Coerce 0-d array-likes. + # ----------------------------------------------------------------- + _dpnp_random_seed = dpnp.random.seed + + def _seed(seed=None, *args, **kwargs): + if seed is not None and getattr(seed, "ndim", None) == 0: + seed = int(seed) + return _dpnp_random_seed(seed, *args, **kwargs) + + _seed.__name__ = "seed" + _seed.__doc__ = getattr(_dpnp_random_seed, "__doc__", None) + dpnp.random.seed = _seed + cupy_fake.random = dpnp.random + + + # ----------------------------------------------------------------- + # cupy.cuda submodule -- creates master queues, installs creation-API + # wrappers on dpnp/dpt, installs the master queue cache (replacing + # dpctl's process-global queue cache), installs in-place op drain. + # See cupy/cuda.py for details. + # ----------------------------------------------------------------- + _cuda_mod = None + try: + from . import cuda as _cuda_mod + cupy_fake.cuda = _cuda_mod + except Exception as e: # was: except ImportError + raise ImportError( + "gpu4pyscf.cupy.cuda failed to initialize the master SYCL queue" + ) from e + + + # ----------------------------------------------------------------- + # Rebind everything that cuda.py patched AFTER the patches. + # + # cuda.py's _wrap_with_master_queue() monkey-patched dpnp.asarray, + # dpnp.zeros, etc., in place. The direct-alias loop above captured + # the PRE-patch references and is now stale. Refresh every name on + # dpnp so cupy.foo() reaches the patched (queue-injecting) version. + # + # Skip names that have custom cupy_fake shims later in this file -- + # those shims already call the patched dpnp.* internally and inherit + # queue injection that way. + # ----------------------------------------------------------------- + _CUSTOM_CUPY_FAKE_SHIMS = frozenset({ + "ndarray", + "zeros", "zeros_like", "empty_like", + "hstack", "vstack", + "allclose", "sqrt", "tril_indices", + "dot", "asarray", "array", + "einsum", + }) + + for _attr in dir(dpnp): + if _attr.startswith("_") or _attr in _CUSTOM_CUPY_FAKE_SHIMS: + continue + _fn = getattr(dpnp, _attr, None) + if callable(_fn): + try: + setattr(cupy_fake, _attr, _fn) + except Exception: + pass + + # Direct aliases for the most common creation APIs (post-patch). + cupy_fake.asarray = dpnp.asarray + cupy_fake.array = dpnp.array + + # cupy.add.at -- dpnp's ufuncs have no .at (unbuffered scatter-add with + # duplicate-index accumulation). Host round-trip keeps np.add.at + # semantics exactly; call sites (hessian, sem) use natm-scale arrays. + class _AddWithAt: + def __call__(self, *args, **kwargs): + return dpnp.add(*args, **kwargs) + + def __getattr__(self, attr): + return getattr(dpnp.add, attr) + + @staticmethod + def at(a, indices, b): + def _host(x): + if isinstance(x, tuple): + return tuple(_host(i) for i in x) + return dpnp.asnumpy(x) if isinstance(x, dpnp.ndarray) else x + host = dpnp.asnumpy(a) + np.add.at(host, _host(indices), _host(b)) + a[...] = host + + cupy_fake.add = _AddWithAt() + + # ================================================================= + # .get() / .set() -- CuPy-style host <-> device transfer + # ================================================================= + def _dpnp_set(self, host_array, stream=None): + # `stream` accepted for CuPy API compatibility; dpnp assignments are + # ordered on the array's SYCL queue, so it is ignored. + self[...] = host_array + + + def _dpnp_get(self, stream=None, order='C', out=None, blocking=True): + host = self.asnumpy() + if out is not None: + out[...] = host + return out + # Preserve 0-d arrays: np.ascontiguousarray / asfortranarray force + # ndim >= 1, turning a scalar `array(5)` into `array([5])`. CuPy's + # .get() keeps the 0-d shape, and downstream code (e.g. using the + # result as a reshape dimension) relies on it being a scalar index. + if host.ndim == 0: + return host + if order == 'C': + return np.ascontiguousarray(host) + if order == 'F': + return np.asfortranarray(host) + if order == 'A': + if host.flags['F_CONTIGUOUS'] and not host.flags['C_CONTIGUOUS']: + return np.asfortranarray(host) + return np.ascontiguousarray(host) + if order == 'K': + return np.array(host, order='K', copy=False) + return np.ascontiguousarray(host) + + + dpnp_array.set = _dpnp_set + dpnp_array.get = _dpnp_get + + + # ================================================================= + # bool() on a size-1 array of ndim > 0. + # + # NumPy (and therefore CuPy) allow truth-testing any array whose size is + # 1, regardless of ndim: `bool(np.array([[5.0]]))` is True. dpnp only + # accepts 0-d and otherwise raises "TypeError: only 0-dimensional arrays + # can be converted to Python scalars". + # + # tdscf/math_helper.py:407 relies on the NumPy behaviour: `xy_norm` comes + # out of `cp.dot(x_tmp, x_tmp.T)` with shape (1, 1) and is then used as + # `if xy_norm > 1e-14:`. + # + # Note we do NOT relax __float__/__int__ -- NumPy 2 raises there for + # ndim > 0 and dpnp already matches, so the two agree. + # ================================================================= + if not getattr(dpnp_array, "__gpu4pyscf_bool_patched__", False): + _orig_dpnp_bool = dpnp_array.__bool__ + + def __bool__(self): + if self.ndim != 0 and self.size == 1: + return bool(self.reshape(())) + return _orig_dpnp_bool(self) + + dpnp_array.__bool__ = __bool__ + dpnp_array.__gpu4pyscf_bool_patched__ = True + + + # ================================================================= + # Pickling. + # + # cupy.ndarray is picklable -- it round-trips through host memory -- and + # pyscf leans on that: dft/tests/test_rks.py::test_rks_lda does + # `pickle.loads(pickle.dumps(mf))` to check that a converged mean-field + # object serializes. dpnp_array is a Cython extension type with a + # non-trivial __cinit__ and no __reduce__, so pickling it raises + # "TypeError: no default __reduce__ due to non-trivial __cinit__". + # Known upstream gap: IntelPython/dpnp#2602 "Cannot serialize arrays" + # (open feature request), so this shim stands until that lands. + # + # Round-trip through NumPy, and rebuild on the master queue so the + # restored array obeys the single-queue invariant that cuda.py enforces. + # Tags on a CPArrayWithTag are carried across too, matching CuPy, where + # the subclass __dict__ is part of the pickled state. + # ================================================================= + if "__reduce__" not in dpnp_array.__dict__: + + def __reduce__(self): + from gpu4pyscf.cupy.cuda import rebuild_dpnp_array + host = dpnp.asnumpy(self) + state = dict(getattr(self, "__dict__", None) or {}) + return (rebuild_dpnp_array, (host, type(self), state)) + + dpnp_array.__reduce__ = __reduce__ + + + # ================================================================= + # hstack / vstack -- cast numpy inputs to dpnp (CuPy does this, dpnp doesn't) + # ================================================================= + def _to_dpnp_seq(seq): + out = [] + for s in seq: + if isinstance(s, (np.ndarray, np.generic)) and not isinstance(s, dpnp.ndarray): + out.append(dpnp.asarray(s)) + else: + out.append(s) + return out + + + def _hstack(tup, *, dtype=None, casting="same_kind"): + return dpnp.hstack(_to_dpnp_seq(tup), dtype=dtype, casting=casting) + + + def _vstack(tup, *, dtype=None, casting="same_kind"): + return dpnp.vstack(_to_dpnp_seq(tup), dtype=dtype, casting=casting) + + + cupy_fake.hstack = _hstack + cupy_fake.vstack = _vstack + + + # ================================================================= + # zeros wrapper -- CuPy allows positional dtype; dpnp requires kwarg + # ================================================================= + def _cupy_zeros(shape, dtype=None, order='C'): + return dpnp.zeros(shape, dtype=dtype, order=order) + + + cupy_fake.zeros = _cupy_zeros + + + # ================================================================= + # zeros_like / empty_like -- CuPy accepts np.ndarray input; dpnp doesn't + # ================================================================= + def _norm_order(order): + return 'C' if order in (None, 'K', 'A') else order + + + def _shape_dtype_from(a, shape=None, dtype=None): + if shape is None: + try: + shape = a.shape + except Exception: + shape = np.asarray(a).shape + if dtype is None: + try: + dtype = a.dtype + except Exception: + dtype = np.asarray(a).dtype + shape = tuple(int(s) for s in shape) + return shape, np.dtype(dtype) + + + def _zeros_like(a, dtype=None, order='K', subok=False, shape=None): + if isinstance(a, np.ndarray): + shape, dtype = _shape_dtype_from(a, shape, dtype) + return dpnp.zeros(shape, dtype=dtype, order=_norm_order(order)) + return dpnp.zeros_like(a, dtype=dtype, order=_norm_order(order)) + + + def _empty_like(a, dtype=None, order='K', subok=False, shape=None): + if isinstance(a, np.ndarray): + shape, dtype = _shape_dtype_from(a, shape, dtype) + return dpnp.empty(shape, dtype=dtype, order=_norm_order(order)) + return dpnp.empty_like(a, dtype=dtype, order=_norm_order(order)) + + + cupy_fake.zeros_like = _zeros_like + cupy_fake.empty_like = _empty_like + + + # ================================================================= + # allclose -- CuPy accepts Python scalars; dpnp.allclose does not + # Upstream: https://github.com/IntelPython/dpnp/issues/2566 + # ================================================================= + def _cupy_allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False): + if np.isscalar(a) and np.isscalar(b): + return bool(np.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)) + + a_is_dp = isinstance(a, dpnp.ndarray) + b_is_dp = isinstance(b, dpnp.ndarray) + a_is_np = isinstance(a, np.ndarray) + b_is_np = isinstance(b, np.ndarray) + + if (a_is_np or b_is_np) and not (a_is_dp and b_is_dp): + if a_is_dp and b_is_np: + return bool(np.allclose(a.asnumpy(), b, rtol=rtol, atol=atol, equal_nan=equal_nan)) + if a_is_np and b_is_dp: + return bool(np.allclose(a, b.asnumpy(), rtol=rtol, atol=atol, equal_nan=equal_nan)) + return bool(np.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)) + + if a_is_dp and b_is_dp: + return bool(dpnp.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)) + + + cupy_fake.allclose = _cupy_allclose + + + # ================================================================= + # sqrt -- CuPy accepts Python scalars; dpnp doesn't. Guarded. + # ================================================================= + if not getattr(dpnp.sqrt, "__gpu4pyscf_patched__", False): + _orig_dpnp_sqrt = dpnp.sqrt + _SCALAR_TYPES = frozenset({int, float, complex, bool}) + + def _patched_dpnp_sqrt(x, _orig=_orig_dpnp_sqrt, **kwargs): + if type(x) in _SCALAR_TYPES: + x = dpnp.array(x) + return _orig(x, **kwargs) + + _patched_dpnp_sqrt.__gpu4pyscf_patched__ = True + dpnp.sqrt = _patched_dpnp_sqrt + + cupy_fake.sqrt = dpnp.sqrt + + + # ================================================================= + # numpy.einsum / numpy.dot -- auto-dispatch to dpnp when any arg is + # dpnp. Guarded. NOTE: `any` here is the Python builtin -- we keep + # cupy_fake separate from this module's globals specifically so that + # stays true; see module docstring. + # ================================================================= + def _convert_np_to_dpnp(arg): + if isinstance(arg, str): + return arg + if isinstance(arg, np.ndarray) and not isinstance(arg, dpnp.ndarray): + return dpnp.asarray(arg) + return arg + + + if not getattr(np.einsum, "__gpu4pyscf_patched__", False): + _original_numpy_einsum = np.einsum + + def _numpy_einsum_with_dpnp(*args, _orig=_original_numpy_einsum, **kwargs): + if any(isinstance(a, dpnp.ndarray) for a in args): + return dpnp.einsum(*(_convert_np_to_dpnp(a) for a in args), **kwargs) + return _orig(*args, **kwargs) + + _numpy_einsum_with_dpnp.__gpu4pyscf_patched__ = True + np.einsum = _numpy_einsum_with_dpnp + + if not getattr(np.dot, "__gpu4pyscf_patched__", False): + _original_numpy_dot = np.dot + + def _numpy_dot_with_dpnp(*args, _orig=_original_numpy_dot, **kwargs): + if any(isinstance(a, dpnp.ndarray) for a in args): + return dpnp.dot(*(_convert_np_to_dpnp(a) for a in args), **kwargs) + return _orig(*args, **kwargs) + + _numpy_dot_with_dpnp.__gpu4pyscf_patched__ = True + np.dot = _numpy_dot_with_dpnp + + # numpy.diag -- same NEP 18 gap as einsum/dot above. cupy.ndarray + # implements __array_function__, so on CUDA `numpy.diag()` + # dispatches to cupy.diag and stays on the device. dpnp.ndarray implements + # __array__ instead, which refuses implicit host conversion, so the same + # call raises "Implicit conversion to a NumPy array is not allowed" under + # SYCL. Dispatching to dpnp.diag keeps the result on the device, matching + # the CUDA path (a host round-trip would also work but costs ~22us). + # + # REVISIT: this patches numpy globally to work around unmodified pyscf + # call sites. Known sites reached from gpu4pyscf (pyscf 2.14.0), all + # `fock = numpy.diag(mo_energy)`-shaped, none of which gpu4pyscf overrides: + # pyscf/mp/mp2.py:770 _ChemistsERIs._common_init_ (canonical/converged + # fast path) -- hit via gpu4pyscf/mp/mp2.py + # _make_eris:123 and gpu4pyscf/mp/dfmp2_old.py + # DFMP2.ao2mo:146; covers test_mp2.py's test_mp2, + # test_mp2_frozen, test_mp2_with_df, test_to_cpu + # pyscf/mp/ump2.py:530-531, pyscf/mp/gmp2.py:264 same, U/G variants + # pyscf/mp/mp2.py:136-137, ump2.py:140-143, gmp2.py:79-80 non-canonical + # `fock[:nocc,:nocc] - numpy.diag(mo_e_o)` + # pyscf/cc/rccsd.py:394,423, pyscf/cc/ccsd_lambda.py:437 + # If these are ever fixed upstream (or gpu4pyscf grows its own overrides), + # this patch can go. + if not getattr(np.diag, "__gpu4pyscf_patched__", False): + _original_numpy_diag = np.diag + + def _numpy_diag_with_dpnp(v, k=0, _orig=_original_numpy_diag): + if isinstance(v, dpnp.ndarray): + return dpnp.diag(v, k=k) + return _orig(v, k=k) + + _numpy_diag_with_dpnp.__gpu4pyscf_patched__ = True + np.diag = _numpy_diag_with_dpnp + + # ================================================================= + # tril_indices -- accept numpy.int64 etc. + # ================================================================= + def _cupy_tril_indices(n, k=0, m=None): + n = int(n) + k = int(k) + m = None if m is None else int(m) + return dpnp.tril_indices(n, k=k, m=m) + + + cupy_fake.tril_indices = _cupy_tril_indices + + + # ================================================================= + # _LazyModule + cupy_backends stubs -- defer loading onemkl_lapack + # ================================================================= + class _LazyModule(types.ModuleType): + def __init__(self, name, loader_func): + super().__init__(name) + self._loader_func = loader_func + self._loaded = False + self._real_module = None + self.__path__ = [] + + def _load(self): + if not self._loaded: + self._real_module = self._loader_func() + self._loaded = True + return self._real_module + + def __getattr__(self, name): + if name.startswith('_'): + raise AttributeError(name) + real = self._load() + if real is None: + raise AttributeError(f"module has no attribute '{name}'") + return getattr(real, name) + + def __dir__(self): + real = self._load() + return dir(real) if real is not None else [] + + + def _load_onemkl_lapack(): + try: + from gpu4pyscf.lib import onemkl_lapack + return onemkl_lapack + except ImportError as e: + import warnings + warnings.warn(f"Could not import onemkl_lapack: {e}") + return None + + + def _setup_cupy_backends(): + if 'cupy_backends' in sys.modules: + return + + cupy_backends = types.ModuleType('cupy_backends') + cupy_backends.__path__ = [] + _cuda_submod = types.ModuleType('cupy_backends.cuda') + _cuda_submod.__path__ = [] + cupy_backends.cuda = _cuda_submod + libs = types.ModuleType('cupy_backends.cuda.libs') + libs.__path__ = [] + _cuda_submod.libs = libs + + cublas = types.ModuleType('cupy_backends.cuda.libs.cublas') + cublas.CUBLAS_FILL_MODE_LOWER = 0 + cublas.CUBLAS_FILL_MODE_UPPER = 1 + cublas.CUBLAS_OP_N = 0 + cublas.CUBLAS_OP_T = 1 + cublas.CUBLAS_OP_C = 2 + + cusolver = _LazyModule('cupy_backends.cuda.libs.cusolver', _load_onemkl_lapack) + libs.cusolver = cusolver + libs.cublas = cublas + + sys.modules['cupy_backends'] = cupy_backends + sys.modules['cupy_backends.cuda'] = _cuda_submod + sys.modules['cupy_backends.cuda.libs'] = libs + sys.modules['cupy_backends.cuda.libs.cusolver'] = cusolver + sys.modules['cupy_backends.cuda.libs.cublas'] = cublas + + gpu4pyscf_cusolver = _LazyModule('gpu4pyscf.lib.cusolver', _load_onemkl_lapack) + sys.modules['gpu4pyscf.lib.cusolver'] = gpu4pyscf_cusolver + + # gpu4pyscf.lib.cupy_helper -> gpu4pyscf.lib.dpnp_helper. + # + # cupy_helper.py reaches for cusolver, cutensor and cupy.cuda.runtime + # APIs that do not exist here, so importing it under SYCL fails. Alias + # it to its dpnp counterpart, which exposes the same public surface. + # + # Done here rather than in gpu4pyscf/lib/__init__.py so that file stays + # identical to upstream: gpu4pyscf/__init__.py imports _patch_pyscf + # (which imports cupy, loading this shim) before it imports .lib, so + # the alias is already registered by the time lib/__init__.py runs. + # Lazy, to avoid importing dpnp_helper before the master queue exists. + def _load_dpnp_helper(): + import importlib + return importlib.import_module('gpu4pyscf.lib.dpnp_helper') + + sys.modules['gpu4pyscf.lib.cupy_helper'] = _LazyModule( + 'gpu4pyscf.lib.cupy_helper', _load_dpnp_helper) + + + _setup_cupy_backends() + del _setup_cupy_backends + + + # ================================================================= + # Memory pool -- reports actual SYCL device memory. + # + # Reference cuda through cupy_fake.cuda (closure-captured) rather + # than re-importing -- this works regardless of which name the cuda + # module ended up registered under in sys.modules. + # ================================================================= + _cuda_ref = _cuda_mod # captured for the pool methods below + + + class _MemoryPool: + """Reports actual SYCL device memory usage via the cuda shim. + + used_bytes = total - free + free_bytes = free memory + total_bytes = total HBM/VRAM capacity + + All other methods are no-ops -- dpnp has no user-managed memory pool. + """ + + def free_all_blocks(self): + # CuPy hands pooled device memory back to the driver here. + # The analogue is releasing the deferred-free batches. + from gpu4pyscf.cupy.cuda import release_deferred_frees + release_deferred_frees() + + # Deprecated CuPy alias, kept for API parity. + free_all_free = free_all_blocks + def set_limit(self, size=None, fraction=None): pass + def get_limit(self): return 0 + def n_free_blocks(self): return 0 + + def used_bytes(self): + try: + return _cuda_ref.get_total_memory() - _cuda_ref.get_free_memory() + except Exception: + return 0 + + def free_bytes(self): + try: + return _cuda_ref.get_free_memory() + except Exception: + return 0 + + def total_bytes(self): + try: + return _cuda_ref.get_total_memory() + except Exception: + return 0 + + + _memory_pool = _MemoryPool() + + + def _get_default_memory_pool(): + return _memory_pool + + + def _get_default_pinned_memory_pool(): + """Pinned memory has no equivalent under SYCL/dpnp; reuse the same pool.""" + return _memory_pool + + + cupy_fake.get_default_memory_pool = _get_default_memory_pool + cupy_fake.get_default_pinned_memory_pool = _get_default_pinned_memory_pool + + if _cuda_mod is not None: + _cuda_mod.PinnedMemoryPool = _MemoryPool + + + # ================================================================= + # cupy.fuse -- no-op under dpnp (CuPy kernel fusion not available) + # ================================================================= + def fuse(*args, **kwargs): + """No-op replacement for cupy.fuse. + + Supports: + @cupy.fuse + @cupy.fuse() + @cupy.fuse(kernel_name='foo') + """ + if len(args) == 1 and callable(args[0]) and not kwargs: + return args[0] # @cupy.fuse (no parens) + return lambda func: func # @cupy.fuse(...) form + + + cupy_fake.fuse = fuse + + + # ================================================================= + # cupy.RawKernel / cupy.RawModule -- runtime kernel compilation. + # + # Backed by dpctl.program.create_kernel_bundle_from_sycl_source (the + # DPC++ `kernel_compiler` extension). The CUDA kernel sources already + # embedded in gpu4pyscf compile unmodified once rawkernel.py's + # compatibility prelude is prepended -- see that module for the + # __global__ / threadIdx / dimension-order mapping, the __shared__ + # rewrite, and the one unsupported construct (dynamic shared memory). + # ================================================================= + from . import rawkernel as _rawkernel_mod + cupy_fake.RawKernel = _rawkernel_mod.RawKernel + cupy_fake.RawModule = _rawkernel_mod.RawModule + sys.modules["cupy.rawkernel"] = _rawkernel_mod + sys.modules["gpu4pyscf.cupy.rawkernel"] = _rawkernel_mod + + + # ================================================================= + # sys.modules aliasing -- make `cupy`, `gpu4pyscf.cupy`, and their + # `.cuda` submodules all resolve to the SAME module objects. + # + # This is what prevents the double-load: once both entries are set, + # a subsequent `import gpu4pyscf.cupy` (or `import cupy`) finds the + # facade in sys.modules and returns it without re-executing this file. + # + # We also redirect sys.modules[__name__] to cupy_fake so whoever is + # currently waiting on this import gets the facade with all the + # attributes (ndarray, asarray, cuda, ...) rather than the partially + # populated current-module object. + # ================================================================= + sys.modules["cupy"] = cupy_fake + sys.modules["gpu4pyscf.cupy"] = cupy_fake + sys.modules[__name__] = cupy_fake # redirect current name too + if _cuda_mod is not None: + sys.modules["cupy.cuda"] = _cuda_mod + sys.modules["gpu4pyscf.cupy.cuda"] = _cuda_mod + + +# ================================================================= +# cupy.fft submodule -- direct aliases, dpnp signature is a superset +# ================================================================= +_fft_mod = types.ModuleType("cupy.fft") +_fft_mod.__package__ = "cupy" +_fft_mod.__path__ = [] + +for _fname in ( + "fft", "ifft", + "fft2", "ifft2", + "fftn", "ifftn", + "rfft", "irfft", + "rfft2", "irfft2", + "rfftn", "irfftn", + "hfft", "ihfft", + "fftshift", "ifftshift", +): + _fn = getattr(dpnp.fft, _fname, None) + if _fn is not None: + setattr(_fft_mod, _fname, _fn) + +# fftfreq / rfftfreq -- dpnp demands `n` be a plain Python int and raises +# ValueError for a numpy integer (np.int32/np.int64), which CuPy/NumPy accept. +# Callers here commonly pass mesh sizes taken from `np.asarray(mesh)`, so +# coerce rather than push the constraint onto every call site. +for _fname in ("fftfreq", "rfftfreq"): + _fn = getattr(dpnp.fft, _fname, None) + if _fn is not None: + def _make_freq(_orig): + def _freq(n, *args, **kwargs): + return _orig(int(n), *args, **kwargs) + return _freq + setattr(_fft_mod, _fname, _make_freq(_fn)) + +cupy_fake.fft = _fft_mod +sys.modules["cupy.fft"] = _fft_mod +sys.modules["gpu4pyscf.cupy.fft"] = _fft_mod + + +# --------------------------------------------------------------------- +# cupy.testing submodule +# +# Upstream test files call cp.testing.assert_allclose(ref, test, atol=...) +# with a mix of host (numpy) and device (dpnp) arrays. dpnp has no +# `testing` namespace, so pull the operands back to the host and defer to +# numpy.testing, which is what CuPy's version does anyway. +# --------------------------------------------------------------------- +_testing_mod = types.ModuleType("cupy.testing") + + +def _to_host(a): + """dpnp/usm array -> numpy; anything else passes through untouched.""" + if isinstance(a, dpnp.ndarray): + return dpnp.asnumpy(a) + if hasattr(a, "asnumpy"): + return a.asnumpy() + return a + + +def _assert_allclose(actual, desired, *args, **kwargs): + return np.testing.assert_allclose( + _to_host(actual), _to_host(desired), *args, **kwargs) + + +def _assert_array_equal(x, y, *args, **kwargs): + return np.testing.assert_array_equal( + _to_host(x), _to_host(y), *args, **kwargs) + + +def _assert_array_almost_equal(x, y, *args, **kwargs): + return np.testing.assert_array_almost_equal( + _to_host(x), _to_host(y), *args, **kwargs) + + +_testing_mod.assert_allclose = _assert_allclose +_testing_mod.assert_array_equal = _assert_array_equal +_testing_mod.assert_array_almost_equal = _assert_array_almost_equal + +cupy_fake.testing = _testing_mod +sys.modules["cupy.testing"] = _testing_mod +sys.modules["gpu4pyscf.cupy.testing"] = _testing_mod diff --git a/gpu4pyscf/cupy/cuda.py b/gpu4pyscf/cupy/cuda.py new file mode 100644 index 000000000..345158774 --- /dev/null +++ b/gpu4pyscf/cupy/cuda.py @@ -0,0 +1,1608 @@ +""" +Single-queue-per-device SYCL runtime shim for gpu4pyscf. + +Design invariant +---------------- +Exactly ONE dpctl.SyclQueue lives per GPU for the lifetime of the process. +Every allocation (Python-side via dpnp/dpctl, C++-side via libgsycl.so) +must land on that master queue. + +Enforcement layers (defence in depth) +------------------------------------- +1. Master queue registry -- `_master_queue(d)` creates the singleton + in-order queue for device `d` on first call, registers its native + pointer with libgsycl.so, and caches it forever. + +2. Global queue-cache replacement -- on dpctl/dpnp master, + `_global_device_queue_cache` is a plain process-global object whose + `get_or_create(key)` returns a SyclQueue. We replace it with a cache + that always returns the per-device master queue. Being process-global + (not a ContextVar), it is also visible to ThreadPoolExecutor worker + threads, so every thread sees the master queue. + +3. Creation-API wrappers -- every dpnp and dpctl.tensor array-creation + function is wrapped to inject `sycl_queue=master` unless the caller + has explicitly placed the allocation. + +Idempotency / reload-safety +--------------------------- +This module can be imported under two dotted names: `cupy.cuda` (when +we're loaded through the gpu4pyscf.cupy facade that re-exports as +`cupy`) and `gpu4pyscf.cupy.cuda` (the real dotted path). Both names +are aliased in gpu4pyscf/cupy/__init__.py, but as belt-and-suspenders +this file stashes its mutable state (master queue registry, device +cache, stream cache) on the `dpnp` module -- which is guaranteed to +load exactly once -- so even if we execute twice we don't duplicate +the master queue or install the wrappers twice. + +Verification +------------ +`_verify_single_queue_invariant()` runs once at import and proves: + - libgsycl's queue pointer matches the Python master per device, + - main-thread dpnp allocations land on master, + - worker-thread dpnp allocations land on master (catches regressions + in layer 2). + +The invariant is checked by native-handle equality (`addressof_ref()`), +not Python-object identity, because dpnp internals may reconstruct a +fresh SyclQueue Python wrapper around the same underlying sycl::queue. +""" +import atexit +import ctypes +import functools +import os +import threading +import time +import types +import warnings +import weakref + +import dpctl +import dpctl.memory as dpmem +import dpctl._sycl_queue_manager as qmgr +import dpnp + + +_DEFERRED_FREE_THRESHOLD = int( + os.environ.get("GPU4PYSCF_DEFERRED_FREE_THRESHOLD", "256") +) + +# Layer 4: queue-ordered deferred free of dpnp USM buffers. +# +# dpctl frees device USM EAGERLY on GC (synchronous sycl::free, not +# queue-ordered). gpu4pyscf launches raw SYCL kernels (lib/*/*.cu under +# USE_SYCL) fire-and-forget on the singleton in-order queue that read +# those buffers; an eager free of a still-in-use buffer -> GPU page +# fault (use-after-free). To match CUDA's stream-ordered free semantics +# we intercept dpnp array creation and, when the array is garbage +# collected, defer the actual release behind a host task gated on a +# queue barrier event, so the free happens only after all pending +# kernels complete. +# +# Set GPU4PYSCF_DEFER_FREE=0 to disable (falls back to eager free). +_DEFER_FREE_ENABLED = os.environ.get("GPU4PYSCF_DEFER_FREE", "1") != "0" +# Only defer allocations at least this many bytes. 0 = defer everything. +# +# This defaults to 0 (defer everything). It previously defaulted to 4096 on the +# rationale that "tiny scalar buffers are rarely the ones handed to raw +# kernels" -- that rationale is FALSE. The DF path hands several sub-4KB index +# arrays straight to raw SYCL kernels as borrowed pointers, e.g. in +# gpu4pyscf/df/int3c2e_bdiv.py: gout_stride (256 B), ksh_offsets_gpu (1 KB), +# shl_pair_offsets (2 KB). Under a 4096 B threshold all three were freed +# EAGERLY (synchronous sycl::free, not queue-ordered) while kernels reading +# them could still be in flight. CuPy frees every allocation stream-ordered +# regardless of size; matching that is the whole point of Layer 4. +_DEFER_FREE_MIN_BYTES = int( + os.environ.get("GPU4PYSCF_DEFER_FREE_MIN_BYTES", "0") +) +# Coalesce this many freed buffers into a single keep-alive host task to +# amortize per-free enqueue/GIL cost. 1 = submit a host task per free. +_DEFER_FREE_BATCH = max(1, int( + os.environ.get("GPU4PYSCF_DEFER_FREE_BATCH", "64") +)) +# Flush a lingering partial batch at most once per this many allocations, +# so tail buffers are not pinned indefinitely without flushing a host task +# on every single allocation. +_DEFER_FREE_FLUSH_STRIDE = max(1, int( + os.environ.get("GPU4PYSCF_DEFER_FREE_FLUSH_STRIDE", "128") +)) +# Confine Layer 4's SYCL scheduler operations (submit_barrier / queue.wait) +# to the thread that imported this module. Guards against a reader/writer +# deadlock inside libsycl's Scheduler -- see _on_scheduler_safe_thread(). +# Set to 0 ONLY to reproduce the pre-fix deadlock for debugging. +_DEFER_FREE_MAIN_THREAD_ONLY = ( + os.environ.get("GPU4PYSCF_DEFER_FREE_MAIN_THREAD_ONLY", "1") != "0" +) +# Safety valve for the guard above: if frees keep arriving on non-owning +# threads and no main-thread activity flushes the batch, stop growing it +# past this many buffers (release the oldest eagerly instead). Only reachable +# in a sustained all-off-thread free burst; normal runs flush long before. +_DEFER_FREE_OFFTHREAD_CAP = max( + _DEFER_FREE_BATCH, + int(os.environ.get("GPU4PYSCF_DEFER_FREE_OFFTHREAD_CAP", "4096")), +) + + +# ===================================================================== +# Shared, reload-safe state -- stashed on dpnp (which loads once). +# +# If this file gets executed twice (two distinct module objects under +# two names), both copies share the same registry, the same device +# cache, and the same "bootstrapped" flag, so _bootstrap() runs its +# side effects exactly once. +# ===================================================================== +_STATE_ATTR = "__gpu4pyscf_cuda_state__" +_state = getattr(dpnp, _STATE_ATTR, None) +if _state is None: + _state = { + # RLock, not Lock: a GC pass can fire mid-critical-section (any + # allocation can cross the threshold) and run a weakref finalizer + # that calls back into _master_queue() from the same thread -- + # _deferred_release -> _flush_deferred_frees_locked -> _master_queue(). + # A plain Lock self-deadlocks there (caught with gdb on a hung + # test_pbc_df_grad.py). + "master_lock": threading.RLock(), + "master_queues": {}, # int -> dpctl.SyclQueue + "gpu_devices": None, # cached device list + "stream_cache": {}, # int -> Stream + "stream_cache_lock": threading.Lock(), + "device_cache": {}, # int -> Device + "device_cache_lock": threading.Lock(), + "bootstrapped": False, + "verified": False, + "shutting_down": False, + # Layer 4 deferred-free bookkeeping. Holds keys: + # "batch" -> list[_Memory] pending queue-ordered release + # "tagged" -> list[(SyclEvent, list[_Memory])] awaiting completion + # "alloc_count" -> int, throttles partial-batch flushing + "defer_free_lock": threading.Lock(), + "pending_frees": {}, + # Thread that imported this module. Layer 4 only ever touches the + # SYCL scheduler (submit_barrier / queue.wait) from this thread -- + # see _on_scheduler_safe_thread() for why. + "owner_thread_id": threading.get_ident(), + } + setattr(dpnp, _STATE_ATTR, _state) + +# Reload-safety: a pre-existing _state (from an earlier load of this +# module under a different dotted name) may predate the Layer 4 keys. +if "defer_free_lock" not in _state: + _state["defer_free_lock"] = threading.Lock() +if "pending_frees" not in _state: + _state["pending_frees"] = {} +if "owner_thread_id" not in _state: + _state["owner_thread_id"] = threading.get_ident() + +_master_lock = _state["master_lock"] +_master_queues = _state["master_queues"] +_stream_cache = _state["stream_cache"] +_stream_cache_lock = _state["stream_cache_lock"] +_device_cache = _state["device_cache"] +_device_cache_lock = _state["device_cache_lock"] +_defer_free_lock = _state["defer_free_lock"] +_pending_frees = _state["pending_frees"] +_owner_thread_id = _state["owner_thread_id"] + + +def _on_scheduler_safe_thread(): + """True only on the thread that imported this module. + + Layer 4's flush/reap paths call into the SYCL scheduler + (`queue.submit_barrier()`, `queue.wait()`). Both take the scheduler's + global reader/writer lock: `submit_barrier` -> `Scheduler::addCG` needs + it EXCLUSIVELY, while a `wait` sits inside + `Scheduler::GraphProcessor::waitForEvent` holding it SHARED for the whole + duration of the wait. + + These calls originate from a weakref finalizer, so they run on whatever + thread happened to drop the last reference. That is frequently NOT the + main thread: gpu4pyscf's DFT path (`dft/numint.py`) runs XC evaluation + inside a ThreadPoolExecutor even for num_devices == 1, and dpnp's own + `keep_args_alive` host tasks drop their kept references from SYCL + thread-pool worker threads. Measured on a single short RKS+newton run: + barriers submitted from 13 distinct non-main threads, plus 15 off-main + `queue.wait()` calls. + + That is enough for a reader/writer deadlock inside libsycl, observed + live under gdb on the full scf/tests/test_soscf.py suite: + + Thread 1 (main): event_impl::wait() + -> GraphProcessor::waitForEvent(shared_lock&) + -> blocked in waitInternal(), STILL HOLDING the + shared lock + + Thread 7 (worker): our finalizer -> SyclQueue.submit_barrier() + -> Scheduler::addCG + -> pthread_rwlock_wrlock BLOCKED behind Thread 1 + + The pending writer then blocks any further readers, the queue never + drains, Thread 1's event never signals. GPU utilization sits at 0% with + every thread parked in futex_wait. `ZE_SERIALIZE=2` masks it only by + changing when frees land relative to scheduler activity. + + So: off-thread frees are still COLLECTED (correctness of the deferral is + unchanged -- the _Memory strong ref keeps the buffer alive, which is the + entire point of Layer 4), they are just not the ones to drive a scheduler + operation. The next main-thread allocation or free flushes them. The + batch is bounded by _DEFER_FREE_BATCH and drained on every main-thread + flush, so this defers reclamation slightly; it does not leak. + """ + if not _DEFER_FREE_MAIN_THREAD_ONLY: + return True + return threading.get_ident() == _owner_thread_id + + +# ===================================================================== +# libgsycl.so -- the C++ side's master-queue registry +# ===================================================================== +_lib_path = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../lib/libgsycl.so")) +libgpu = ctypes.CDLL(_lib_path) + +# Bindings must match sycl_api_python.cpp exactly. +libgpu.sycl_get_device_id.argtypes = [] +libgpu.sycl_get_device_id.restype = ctypes.c_int +libgpu.sycl_get_queue_ptr.argtypes = [] +libgpu.sycl_get_queue_ptr.restype = ctypes.c_void_p +libgpu.sycl_set_queue_ptr.argtypes = [ctypes.c_int, ctypes.c_void_p] +libgpu.sycl_set_queue_ptr.restype = None +libgpu.sycl_set_device.argtypes = [ctypes.c_int] +libgpu.sycl_set_device.restype = None +libgpu.sycl_get_total_memory.argtypes = [] +libgpu.sycl_get_total_memory.restype = ctypes.c_size_t +libgpu.sycl_get_shared_memory.argtypes = [] +libgpu.sycl_get_shared_memory.restype = ctypes.c_size_t +libgpu.sycl_get_compute_units.argtypes = [] +libgpu.sycl_get_compute_units.restype = ctypes.c_int +libgpu.sycl_get_device_name.argtypes = [ctypes.c_char_p, ctypes.c_int] +libgpu.sycl_get_device_name.restype = None +libgpu.sycl_get_free_memory.argtypes = [] +libgpu.sycl_get_free_memory.restype = ctypes.c_size_t +libgpu.sycl_memcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_size_t] +libgpu.sycl_memcpy.restype = ctypes.c_size_t + + +class classproperty: + def __init__(self, fget): + self.fget = fget + def __get__(self, obj, owner): + return self.fget(owner) + + +# ===================================================================== +# Queue pointer helper +# ===================================================================== +def _get_sycl_queue_ptr(q: dpctl.SyclQueue) -> int: + """Return the actual sycl::queue* as an integer. + + DPCTLSyclQueueRef is a typedef for sycl::queue*, and + SyclQueue.addressof_ref() returns its value cast to size_t -- + i.e. the sycl::queue* itself. sycl_set_queue_ptr does a direct + static_cast, so we pass the value as-is. + + q must remain alive for the lifetime of the stored pointer -- + _master_queues guarantees this for master queues. + """ + return int(q.addressof_ref()) + + +# ===================================================================== +# Master-queue registry +# ===================================================================== +def _gpu_devices(): + """Enumerate compute devices once (prefer level_zero GPUs). Cached in _state. + + Each candidate is tried in turn and the first non-empty result wins. A + query that succeeds but returns nothing has to fall through exactly like + one that raises: on a machine with no GPU -- a CPU-only CI runner, say -- + the GPU queries return [] without error, and stopping there would leave + the master-queue registry empty. Every later layer assumes device 0 + exists, so the shortfall would not surface here but as a std::terminate + out of libgsycl's sycl_set_device. + + The final candidate is unfiltered, so the OpenCL CPU device is used when + that is all there is. Where a GPU is present the first query answers and + the rest are never reached, leaving that path unchanged. + """ + if _state["gpu_devices"] is not None: + return _state["gpu_devices"] + + candidates = ( + lambda: dpctl.get_devices(backend="level_zero", device_type="gpu"), + lambda: dpctl.get_devices(device_type="gpu"), + lambda: dpctl.get_devices(), + ) + + devs = [] + for query in candidates: + try: + devs = query() or [] + except Exception: + devs = [] + if devs: + break + + _state["gpu_devices"] = devs + return devs + +def _master_queue(device_id=None): + """Return the singleton master in-order SyclQueue for a device. + + First call creates the queue and registers its native sycl::queue* + with libgsycl.so so low-level kernel launches run on the same + in-order queue dpnp/dpctl use. Cached for the process lifetime, which + keeps the pointer valid. dpctl defers USM frees on in-order queues + (queue-ordered host task), so allocations are not released while + kernels enqueued here -- including these C++ launches -- still use them. + """ + if device_id is None: + device_id = int(libgpu.sycl_get_device_id()) + with _master_lock: + q = _master_queues.get(device_id) + if q is not None: + return q + devs = _gpu_devices() + if device_id < 0 or device_id >= len(devs): + raise ValueError( + f"device_id {device_id} out of range (have {len(devs)} GPUs)") + q = dpctl.SyclQueue(devs[device_id], property="in_order") + libgpu.sycl_set_queue_ptr( + ctypes.c_int(device_id), + ctypes.c_void_p(_get_sycl_queue_ptr(q))) + _master_queues[device_id] = q # keeps q alive -> pointer stays valid + return q + + +def master_device(device_id=None): + """Public accessor for the master SyclQueue of a device. + + Pass this anywhere code needs an explicit ``sycl_queue=``. + """ + return _master_queue(device_id) + + +def _same_queue(q1, q2): + if q1 is None or q2 is None: + return False + if q1 is q2: + return True + try: + return _get_sycl_queue_ptr(q1) == _get_sycl_queue_ptr(q2) # compare sycl::queue* + except Exception: + return False + + +# ===================================================================== +# Layer 2 -- replace dpctl's process-global queue cache +# ===================================================================== +class _MasterQueueCache: + """Drop-in replacement for dpctl._DeviceDefaultQueueCache. + + On dpctl/dpnp master, `_global_device_queue_cache` is a plain + process-global object (NOT a ContextVar), and + `get_device_cached_queue(key)` calls + `_global_device_queue_cache.get_or_create(key)` directly, expecting a + bare dpctl.SyclQueue in return. + + We resolve every key to the per-device master in-order queue so all + dpnp/dpctl allocations land on the singleton queue for that GPU. + Because this object is process-global rather than a ContextVar, + ThreadPoolExecutor worker threads observe it too -- fixing the + worker-thread allocation escape that motivated the original shim. + + Accepted key types (per dpctl): a SyclDevice, a (SyclContext, + SyclDevice) 2-tuple, or a oneAPI filter-selector string. Unknown key + types or devices not present among the enumerated GPUs raise rather + than silently falling back to device 0. + """ + __slots__ = ("_lock",) + + def __init__(self): + self._lock = threading.Lock() + + def _device_from_key(self, key): + if isinstance(key, tuple) and len(key) == 2: + return key[1] + if isinstance(key, str): + return dpctl.SyclDevice(key) # may raise -> propagate + if isinstance(key, dpctl.SyclDevice): + return key + raise TypeError( + f"_MasterQueueCache.get_or_create: unsupported key type " + f"{type(key)!r}") + + def _device_id_for(self, dev): + devs = _gpu_devices() + # Exact device-object match against the same list used to build the + # master queues. + for i, d in enumerate(devs): + try: + if d == dev: + return i + except Exception: + pass + # Backup match by oneAPI filter string. + for i, d in enumerate(devs): + try: + if d.filter_string == dev.filter_string: + return i + except Exception: + pass + raise RuntimeError( + f"_MasterQueueCache: device {dev} not found among the " + f"{len(devs)} enumerated GPU(s); cannot map it to a master queue") + + def get_or_create(self, key): + with self._lock: + return _master_queue(self._device_id_for(self._device_from_key(key))) + + # dpctl internals may copy/update the cache; keep safe stubs. + def _update_map(self, *args, **kwargs): + return None + + def __copy__(self): + return self + +# ===================================================================== +# Layer 3 -- wrap every creation API so sycl_queue=master is injected +# ===================================================================== +_DPNP_CREATION = ( + "asarray", "array", "zeros", "ones", "empty", "full", + "zeros_like", "ones_like", "empty_like", "full_like", + "arange", "linspace", "logspace", "geomspace", + "eye", "identity", "tri", "frombuffer", "fromfunction", + "copy", +) + + +# ===================================================================== +# Layer 4 -- queue-ordered deferred free of dpnp USM buffers +# ===================================================================== +# +# Why this exists +# --------------- +# dpctl frees device USM EAGERLY on GC (synchronous sycl::free in +# _Memory.__dealloc__, NOT queue-ordered). gpu4pyscf launches raw SYCL +# kernels (lib/*/*.cu under USE_SYCL) fire-and-forget on the singleton +# in-order queue that read those buffers; the C++ kernel wrappers only +# receive BORROWED raw pointers and cannot own/keep the buffers alive. +# So an eager free of a still-in-use buffer -> GPU page fault. +# +# How ordering is achieved WITHOUT a keep-alive host task +# --------------------------------------------------------- +# The fix keeps the freed buffer's owning dpctl _Memory alive, submits a +# BARRIER on the master queue, and retains the (event, batch) pair on the +# Python side. A later call to _reap_completed_locked() drains the master +# queue with a REAL `.wait()` (batched/throttled -- not per kernel launch, +# see that function's docstring for why a non-blocking execution_status poll +# was tried first and found insufficient on this SYCL/PVC stack) before +# dropping any tagged batch, so the real sycl::free happens strictly after +# the kernels that read those buffers. `.wait()` releases the GIL while +# blocking, so this does not reintroduce the async_dec_ref/PyGILState_Ensure +# deadlock described below -- that required a SYCL WORKER THREAD to need the +# GIL back while the calling thread held it; here the calling thread gives up +# the GIL itself before blocking. +# +# Why NOT a keep-alive host task +# ------------------------------ +# The previous design handed the batch to SyclQueue._submit_keep_args_alive(), +# enqueueing dpctl's async_dec_ref host task, which calls PyGILState_Ensure() +# when it runs. That DEADLOCKS on a shared in-order queue: an in-order queue +# serializes ENQUEUE, so a Python thread submitting work while holding the GIL +# blocks in SYCL's scheduler behind the un-run host task, which can then never +# acquire the GIL. Polling an event needs no callback and never touches the +# GIL from a SYCL worker thread. (execution_status is a cheap status query and +# is safe on a barrier event even when the queue was idle at submit time -- +# unlike .wait(), which the Event class below documents as unreliable there.) +# +# Cost control: batching +# ---------------------- +# Each barrier is itself an enqueued command, so submitting one per freed +# array would be expensive under high allocation churn. We instead COALESCE +# freed _Memory objects and tag a whole batch with a SINGLE barrier event, +# amortizing the cost to ~1 barrier per _DEFER_FREE_BATCH frees. + + +def _flush_deferred_frees_locked(): + """Tag the pending batch with a barrier event and retain it. + + Caller must hold _defer_free_lock. + + NO host task is submitted (intel/llvm#22943 -- an in-order queue plus a + host task that takes a lock deadlocks). This previously called + SyclQueue._submit_keep_args_alive(), which enqueues dpctl's async_dec_ref + host task; that task calls PyGILState_Ensure() when it runs. On the shared + in-order master queue that deadlocks: an in-order queue serializes ENQUEUE, + so any Python thread that later submits work while holding the GIL blocks + inside SYCL's scheduler behind the un-run host task, which can then never + acquire the GIL to retire. (Observed as the test_rhf_hessian hang: a worker + thread inside oneMKL ddot holding the GIL, a SYCL thread-pool worker stuck + in async_dec_ref -> PyGILState_Ensure, and the main thread waiting on the + worker's future.) + + Instead we submit a barrier and retain the batch on the Python side. + `_reap_completed_locked()` now drains the master queue with a real + `.wait()` before releasing any tagged batch (see that function's + docstring for why the earlier non-blocking `execution_status` check was + replaced) -- the barrier submitted here is retained for bookkeeping + (probe stats, the `(event, batch)` tagging structure) but the actual + safety guarantee comes from the queue drain in `_reap_completed_locked`, + not from this event's status. The GIL is never touched from a SYCL + worker thread by anything in this function. + """ + batch = _pending_frees.get("batch") + if not batch: + return + # Never submit to the SYCL scheduler off the owning thread: addCG takes + # the scheduler's write lock and deadlocks against a main thread parked + # inside waitForEvent with the read lock held. Leave the batch pending; + # the next main-thread flush picks it up. See _on_scheduler_safe_thread. + if not _on_scheduler_safe_thread(): + return + _pending_frees["batch"] = [] + try: + ev = _master_queue().submit_barrier() + _pending_frees.setdefault("tagged", []).append((ev, batch)) + except Exception: + # On failure, dropping `batch` here frees eagerly (still correct + # if no kernel is mid-flight; worst case reproduces the original + # eager-free behavior only for this batch). + pass + + +def _reap_completed_locked(): + """Release every tagged batch, after a REAL wait for true completion. + + Caller must hold _defer_free_lock. + + HISTORY / WHY THIS CHANGED (DEFECT5 hypothesis 24, Finding 4) + --------------------------------------------------------------- + This previously polled `ev.execution_status == event_status_type.complete` + -- a non-blocking status query, never a host wait -- on the theory that a + completed barrier event proves every kernel submitted before it has + genuinely finished touching device memory, so the real sycl::free is safe. + + That assumption was independently disproved this session on the SAME + SYCL/PVC/Level-Zero stack, in a different code path: bisecting a + reproducible segfault in `RYS_build_jk`'s task loop showed that a bare + `queue.submit_barrier()` event -- checked without a host wait -- is NOT + sufficient to guarantee true completion, while the SAME barrier `.wait()`d + IS sufficient (3/3 clean runs each way; see + hang_analysis_evidence/DEFECT5_free_and_device_global_audit.md, section + 5e/5f). Since this reaper used the identical primitive + (submit_barrier() + a completion check with no host wait) to decide when + to run `sycl::free`, it was exposed to the same gap: a status query + reporting "complete" before the barrier has actually drained lets this + reaper free memory a still-running kernel is reading -- a genuine + read-after-free, producing exactly the NotPresent page fault this + investigation was chasing. Confirmed present and firing (before this fix) + on the exact failing test via `GPU4PYSCF_REAPER_PROBE` instrumentation + (section 5h): two batches (64 then ~16 items) reaped via the un-waited + status query on every single `get_jk` call, before the task loop that + later reports the fault even starts. + + THE FIX: a real, GIL-releasing wait + ------------------------------------ + `_master_queue().wait()` drains the ENTIRE in-order master queue -- a + documented-safe superset of waiting for any barrier submitted on that + queue (see the `Event`/`Device.synchronize()` comment above, which + already uses this exact pattern and documents why: `submit_barrier()` on + an IDLE queue can return a Level Zero "internal event" that cannot be + `.wait()`'d on directly, but `queue.wait()` has no such caveat and is + always safe). Both `SyclQueue.wait()` and `SyclEvent.wait()` release the + GIL while blocking (`with nogil: DPCTLQueue_Wait(...)` / + `DPCTLEvent_Wait(...)` in dpctl's Cython source) -- this is NOT the + `async_dec_ref`/`PyGILState_Ensure` deadlock this reaper design was + originally built to avoid. That deadlock required a SYCL WORKER THREAD to + need the GIL back while the main thread held it inside the driver's + enqueue path; here, the CALLING Python thread simply releases the GIL + itself before blocking, so any other thread (worker or otherwise) that + needs the GIL remains free to acquire it throughout the wait. + + Cost: this still only runs where `_flush_deferred_frees_locked` / + `_reap_completed_locked` were already being called (throttled to every + `_DEFER_FREE_FLUSH_STRIDE` allocations, or when a batch reaches + `_DEFER_FREE_BATCH`) -- NOT once per kernel launch. This keeps the fix + entirely inside this shim layer; no call site outside this file changes. + """ + tagged = _pending_frees.get("tagged") + if not tagged: + return + # queue.wait() parks inside Scheduler::GraphProcessor::waitForEvent while + # holding the scheduler's SHARED lock. Doing that from a worker thread + # adds a second reader that can outlive the main thread's own wait and + # starve a pending writer (submit_barrier from any thread). Confine the + # drain to the owning thread. See _on_scheduler_safe_thread. + if not _on_scheduler_safe_thread(): + return + # Release _defer_free_lock across the actual blocking wait. The SYCL + # host-task thread this wait drains runs Python DECREFs (dropping the + # tagged batch's dpnp arrays); if one of those DECREFs collects another + # usm_ndarray, its finalizer re-enters _deferred_release(), which needs + # this SAME lock to append to "batch". Holding the lock through the wait + # self-deadlocks: this (owning) thread parks in queue.wait() waiting for + # the host task to finish, while the host task blocks acquiring a lock + # this thread still holds. Confirmed via gdb on the real hang: main + # thread inside DPCTLQueue_Wait -> Scheduler::waitForEvent, a SYCL + # ThreadPool worker inside a DispatchHostTask DECREF chain blocked in + # PyThread_acquire_lock_timed. Dropping the lock here is safe -- only + # the owning thread ever reaches this function or _flush_deferred_frees_ + # locked (both gated on _on_scheduler_safe_thread), so nothing else can + # touch "tagged" while we wait; off-thread callers only ever append to + # "batch", which is unaffected by releasing this lock. + _defer_free_lock.release() + try: + _master_queue().wait() + except Exception: + # Cannot drain the queue -- release rather than leak. + pass + finally: + _defer_free_lock.acquire() + # Every tagged batch was submitted strictly before this wait (the master + # queue is in-order and _flush_deferred_frees_locked always submits its + # barrier before returning), so draining the queue proves every one of + # them is now genuinely safe to release. Nothing stays "still pending". + _pending_frees["tagged"] = [] + + +def _deferred_release(mem): + """Finalizer body: queue the freed USM `_Memory` for batched, + queue-ordered release. + + `mem` is a strong reference to the dpctl _Memory owner; holding it + here means the eager sycl::free in _Memory.__dealloc__ has NOT run + yet. We append it to the pending batch and flush when the batch is + large enough. + """ + # During interpreter shutdown, host tasks acquiring the GIL are unsafe + # (dpctl warns). Returning drops `mem` -> eager free, which is fine at + # exit since no new kernels are being launched. + if _state.get("shutting_down"): + return + with _defer_free_lock: + batch = _pending_frees.setdefault("batch", []) + batch.append(mem) + if len(batch) >= _DEFER_FREE_BATCH: + # No-op off the owning thread (scheduler deadlock guard). The + # batch then keeps growing until a main-thread free/alloc flushes + # it, which is the common case -- but a long run of purely + # off-thread frees with no intervening main-thread activity would + # pin memory without bound. Cap it: past the safety limit, release + # the overflow eagerly rather than grow forever. Eager release is + # the pre-Layer-4 behavior (a correctness risk only if that exact + # buffer is being read by a raw kernel right now), which is + # strictly better than an unbounded hold. + if _on_scheduler_safe_thread(): + _flush_deferred_frees_locked() + elif len(batch) >= _DEFER_FREE_OFFTHREAD_CAP: + del batch[:-_DEFER_FREE_BATCH] + + +def _register_deferred_free(arr): + """Register a finalizer on a freshly created dpnp array so that, when + it is garbage collected, its USM allocation is released in a batched, + queue-ordered manner instead of eagerly. + + No-op (returns arr unchanged) if deferral is disabled, the object is + not a dpnp array, it is too small, or the underlying USM handles are + unavailable. + """ + if not _DEFER_FREE_ENABLED: + return arr + try: + get_array = getattr(arr, "get_array", None) + if get_array is None: + return arr + usm = get_array() # weak-referenceable usm_ndarray + mem = usm.usm_data # strong ref to _Memory owner + nbytes = getattr(mem, "nbytes", 0) + if nbytes < _DEFER_FREE_MIN_BYTES: + return arr + # Flush a lingering partial batch so buffers freed during a burst + # of frees followed by pure compute (no more frees to trigger a + # batch flush) do not stay pinned indefinitely. This piggybacks on + # allocation activity and is throttled by _flush_stride so it does + # not submit a host task on every allocation. + cnt = _pending_frees.get("alloc_count", 0) + 1 + _pending_frees["alloc_count"] = cnt + if (cnt % _DEFER_FREE_FLUSH_STRIDE) == 0: + with _defer_free_lock: + _flush_deferred_frees_locked() + _reap_completed_locked() + # weakref.finalize on the usm_ndarray fires when it is collected; + # `mem` captured in the finalizer keeps _Memory alive past that, + # letting us order the real free behind queue work. + weakref.finalize(usm, _deferred_release, mem) + except Exception: + # Never let lifetime-management bookkeeping break array creation. + return arr + return arr + + +# Names of dpnp_array methods/operators whose result is a NEW device +# allocation that Layer 3's creation-API wrapping does not cover. Arithmetic +# results and dtype conversions are the important ones: before this, `a * 2` +# and `a.astype(...)` produced buffers that were freed EAGERLY no matter how +# large, because only the creation APIs registered a finalizer. CuPy is +# stream-ordered for every array however it was produced; this closes the gap. +_DPNP_ARRAY_PRODUCERS = ( + # binary arithmetic (and their reflected forms) + "__add__", "__radd__", "__sub__", "__rsub__", + "__mul__", "__rmul__", "__truediv__", "__rtruediv__", + "__floordiv__", "__rfloordiv__", "__mod__", "__rmod__", + "__pow__", "__rpow__", "__matmul__", "__rmatmul__", + # unary + "__neg__", "__pos__", "__abs__", + # dtype / layout conversions that allocate + "astype", "conj", "conjugate", +) + + +# Module-level dpnp functions that ALLOCATE a new device buffer but are +# neither creation APIs (Layer 3 / _DPNP_CREATION) nor ndarray dunders +# (Layer 4b / _DPNP_ARRAY_PRODUCERS). Before Layer 4c these escaped deferred +# free entirely: dpctl released them via the eager, NON-queue-ordered +# sycl::free in _Memory.__dealloc__ the moment the last Python reference +# dropped -- even with raw kernels still reading the buffer. +# +# This is not hypothetical. In `_VHFOpt.get_jk` with hermi==0 (the default, +# and what test_j_engine_integral_screen exercises): +# dms = cp.vstack([dms, dms.transpose(0,2,1)]) +# `dms` -- the density matrix whose RAW POINTER is handed to all 28 +# RYS_build_jk launches -- was produced by `vstack` and therefore had no +# deferred-free finalizer. The Level Zero loader trace (UR_L0_DEBUG) shows +# exactly one zeMemFree during compute, firing 1 ms after the final kernel +# launch while that kernel's event still reported ZE_RESULT_NOT_READY, with +# the GPU page fault landing 226 ms later INSIDE dm[0] (offsets 1.34-1.81 MB +# of the 2.70 MB slab). See DEFECT5_free_and_device_global_audit.md. +# +# CuPy has no equivalent bug because cudaFree() implicitly synchronizes the +# device; sycl::free() does not. Layer 4 exists to close precisely that gap, +# and this list closes the part of it Layers 3 and 4b did not reach. +_DPNP_ALLOCATING_FUNCS = ( + # shape / joining -- these produce the buffers most likely to be handed + # to a raw kernel as a borrowed pointer + "vstack", "hstack", "dstack", "column_stack", "row_stack", + "concatenate", "stack", "append", "repeat", "tile", + # linear algebra / reductions producing fresh buffers + "outer", "sum", "prod", "cumsum", "trace", + # elementwise ufuncs (module-level forms; the operator forms are 4b) + "exp", "log", "sqrt", "square", "abs", "sign", + "multiply", "add", "subtract", "divide", + # selection / construction + "where", "take", "tril", "triu", "unique", +) + + +def _wrap_allocating_funcs(): + """Layer 4c -- attach the deferred-free finalizer to module-level dpnp + functions that allocate. + + Same rationale as Layer 4b (`_wrap_array_producers`), but for functions + reached as `dpnp.foo(...)` rather than as an operator on an ndarray. + Deliberately does NOT inject `sycl_queue=` -- these are compute-follows- + data operations that correctly inherit their queue from their inputs + (verified: a global queue-identity probe over this whole code path found + zero divergence). The ONLY thing being added is the queue-ordered free. + + Idempotent; failures are swallowed so a dpnp build missing any one of + these names cannot break import. + """ + for name in _DPNP_ALLOCATING_FUNCS: + orig = getattr(dpnp, name, None) + if orig is None or getattr(orig, "__master_q_wrapped__", False): + continue + + if isinstance(orig, types.FunctionType): + @functools.wraps(orig) + def wrapper(*args, _orig=orig, **kwargs): + return _register_deferred_free(_orig(*args, **kwargs)) + + wrapper.__master_q_wrapped__ = True + wrapper.__wrapped__ = orig + else: + # `multiply`/`add`/`subtract`/`divide` are DPNPBinaryFunc + # objects (ufunc-like), not plain functions -- they carry + # callable attributes such as `.outer` that dpnp's own + # implementations reach through the module attribute (e.g. + # `dpnp.outer` calls `dpnp.multiply.outer(...)`). A + # `functools.wraps` closure is a bare function and has no + # `.outer`, so replacing the module attribute with one broke + # any internal dpnp call that goes through it + # (AttributeError: 'function' object has no attribute + # 'outer', hit by dpnp_helper.krylov's QR step via + # dpnp.outer -> dpnp.multiply.outer). Use an + # attribute-forwarding proxy instead so calling the object + # still hits the deferred-free path while every other + # attribute resolves straight through to `orig`. + class _DeferredFreeProxy: + def __init__(self, orig): + self._orig = orig + self.__master_q_wrapped__ = True + + def __call__(self, *args, **kwargs): + return _register_deferred_free(self._orig(*args, **kwargs)) + + def __getattr__(self, attr): + return getattr(self._orig, attr) + + wrapper = _DeferredFreeProxy(orig) + try: + setattr(dpnp, name, wrapper) + except (TypeError, AttributeError): + continue + + +def _wrap_array_producers(): + """Attach the deferred-free finalizer to arithmetic / astype results. + + Layer 3 only wraps dpnp's *creation* functions, so any array produced by + an operator (`a * 2`) or a conversion (`a.astype(...)`) escaped Layer 4 + entirely and was released by dpctl's eager synchronous sycl::free. If such + a buffer had been handed to a raw SYCL kernel as a borrowed pointer, that + is a use-after-free. + + Wrapping the dunder on the *type* is required -- Python looks up operators + on the type, not the instance. Idempotent, and failures are swallowed so a + dpnp version without one of these names cannot break import. + """ + try: + arr_cls = dpnp.ndarray + except AttributeError: + return + for name in _DPNP_ARRAY_PRODUCERS: + orig = getattr(arr_cls, name, None) + if orig is None or getattr(orig, "__master_q_wrapped__", False): + continue + + @functools.wraps(orig) + def wrapper(self, *args, _orig=orig, _name=name, **kwargs): + return _register_deferred_free(_orig(self, *args, **kwargs)) + + wrapper.__master_q_wrapped__ = True + wrapper.__wrapped__ = orig + try: + setattr(arr_cls, name, wrapper) + except (TypeError, AttributeError): + # Immutable/extension type -- skip rather than fail import. + continue + + +def _wrap_with_master_queue(mod, names): + """Inject sycl_queue=master into every creation call on `mod`, and + register a queue-ordered deferred-free finalizer on the result. + + Idempotent: re-wrapping a wrapped function is a no-op. + """ + for name in names: + orig = getattr(mod, name, None) + if orig is None or getattr(orig, "__master_q_wrapped__", False): + continue + + @functools.wraps(orig) + def wrapper(*args, _orig=orig, **kwargs): + if "sycl_queue" not in kwargs and "device" not in kwargs: + kwargs["sycl_queue"] = _master_queue() + return _register_deferred_free(_orig(*args, **kwargs)) + + wrapper.__master_q_wrapped__ = True + wrapper.__wrapped__ = orig + setattr(mod, name, wrapper) + +# ===================================================================== +# Layer 5: drain the queue before dpnp's blocking native math calls. +# +# Workaround for intel/llvm#22943 -- "[SYCL][UR] Hangs when using +# `in-order` and SYCL `host_task` under multi-threading" (open as of +# 2026-08-14; reproduced on PVC 1550 with both Level-Zero and OpenCL, +# not seen on CUDA/HIP). An in-order queue plus a host task that takes a +# lock deadlocks; out-of-order queues do not. +# +# Here the lock is the GIL. dpnp's BLAS/LAPACK pybind11 extensions call +# into oneMKL, which blocks on sycl::event::wait() internally and does +# NOT release the GIL while it does so. On the in-order master queue +# that wait transitively covers every command submitted earlier -- +# including the keep-alive host tasks dpctl attaches to Python operands. +# Such a host task runs on a SYCL worker thread and needs the GIL to +# DECREF, but the caller blocked inside oneMKL is still holding it. +# Permanent deadlock. +# +# The workaround the issue recommends -- switch to an out-of-order queue +# -- is not available to us: libgint/libgvhf/libgdft are handed the raw +# sycl::queue* and launch kernels on it with no event plumbing across +# the ctypes boundary, and the deferred-free reaper below tags batches +# with barriers on the assumption of in-order semantics. So we instead +# make sure no GIL-needing host task is ever pending when oneMKL blocks. +# +# Two live instances were diagnosed with gdb/py-spy: +# df.DF.build() -> cholesky() -> mkl::lapack::potrf_dispatch +# -> event_impl::waitInternal [holds GIL] +# vs ThreadPool worker in take_gil +# int3c2e.get_j_int3c2e_pass1() -> coeff @ dm0 -> bi._gemm [same shape] +# +# The wait has to happen at the *native* boundary, not at the public +# dpnp.linalg/dpnp.matmul entry point: those first make copies and +# temporaries, each of which registers a fresh keep-alive host task, so +# a drain performed before them is already stale by the time oneMKL is +# reached. dpctl's SyclQueue.wait() is declared `with nogil`, so the +# drain lets any pending host task retire first. +# +# Measured cost of the added drain: none (0.240 vs 0.250 ms per 512x512 +# matmul). dpnp is already effectively synchronous per operation on an +# in-order queue with host-task keep-alives, so there is no pipelining +# to lose. +# ===================================================================== +_DPNP_NATIVE_BLOCKING = { + "dpnp.backend.extensions.blas._blas_impl": ( + "_dot", "_dotc", "_dotu", "_gemm", "_gemm_batch", "_gemv", "_syrk", + ), + "dpnp.backend.extensions.lapack._lapack_impl": ( + "_geqrf", "_geqrf_batch", "_gesv", "_gesv_batch", "_gesvd", + "_gesvd_batch", "_getrf", "_getrf_batch", "_getri_batch", "_getrs", + "_getrs_batch", "_heevd", "_heevd_batch", "_orgqr", "_orgqr_batch", + "_potrf", "_potrf_batch", "_syevd", "_syevd_batch", "_ungqr", + "_ungqr_batch", + ), +} + + +def _drain_then(orig): + def wrapper(*args, **kwargs): + try: + _master_queue().wait() + except Exception: + pass + return orig(*args, **kwargs) + + wrapper.__name__ = getattr(orig, "__name__", "wrapped") + wrapper.__doc__ = getattr(orig, "__doc__", None) + wrapper.__gil_drained__ = True + wrapper.__wrapped__ = orig + return wrapper + + +def _wrap_blocking_setitem(): + """Same treatment for dpnp's `arr[idx] = value` (intel/llvm#22943). + + `usm_ndarray.__setitem__` enqueues the copy then blocks on + `event_impl::wait()` in `dpnp/tensor/_tensor_impl`; on the in-order queue + that reaches `Scheduler::GraphProcessor::waitForEvent`, which blocks while + holding the graph read lock, so an in-flight host task can never be + enqueued. Caught with gdb on a hung test_df_int3c2e.py::test_int3c2e_rsh. + Overhead ~7-10% per setitem; accepted because the alternative is a hang. + """ + from dpnp.dpnp_array import dpnp_array as _arr + orig = _arr.__setitem__ + if getattr(orig, "__gil_drained__", False): + return + _arr.__setitem__ = _drain_then(orig) + + +def _wrap_blocking_lapack(): + import importlib + + for mod_name, fn_names in _DPNP_NATIVE_BLOCKING.items(): + try: + mod = importlib.import_module(mod_name) + except Exception as e: + warnings.warn( + f"Layer 5: could not import {mod_name} to guard against the " + f"oneMKL/GIL host-task deadlock: {e}", + RuntimeWarning) + continue + for name in fn_names: + orig = getattr(mod, name, None) + if orig is None or getattr(orig, "__gil_drained__", False): + continue + try: + setattr(mod, name, _drain_then(orig)) + except Exception: + # A pybind11 module that refuses attribute assignment would + # leave the deadlock in place; say so rather than fail + # silently. + warnings.warn( + f"Layer 5: {mod_name}.{name} is not patchable; the " + "oneMKL/GIL host-task deadlock remains reachable there.", + RuntimeWarning) + + +# ===================================================================== +# Layer 6: keep dpctl's order managers alive. +# +# Second face of intel/llvm#22943. dpctl keeps its +# `_SequentialOrderManager` instances in a **thread-local** map +# (`SyclQueueToOrderManagerMap._get_map`), and the manager's `__del__` +# does +# +# SyclEvent.wait_for(_local.get_submitted_events()) +# SyclEvent.wait_for(_local.get_host_task_events()) +# +# So when *any* worker thread exits, its thread-local dict is torn down +# and a blocking SYCL event wait runs from inside a garbage-collection +# finalizer. That wait enters +# `Scheduler::GraphProcessor::waitForEvent`, which blocks while holding +# the graph read lock; if a host task is in flight it can never be +# enqueued, and the process wedges. Captured with gdb on a hung +# `dft/tests/test_numint.py`: +# +# Thread 9 : slot_tp_finalize -> SyclEvent.wait_for -> DPCTLEvent_Wait +# -> Scheduler::waitForEvent -> enqueueCommand(BLOCKING) +# -> event_impl::waitInternal [holds GraphReadLock] +# Thread 3 : DispatchHostTask::waitForEvents -> urEventWait +# Thread 1 : blocked on a Python lock held by thread 9 +# +# Pinning every manager with a process-lifetime strong reference means +# `__del__` never runs before interpreter shutdown, where dpctl's own +# `sys.is_finalizing()` guard already short-circuits the waits. Nothing +# else changes: the managers stay functional and keep ordering work +# exactly as before. Cost is a few small objects per thread. +# ===================================================================== +_pinned_order_managers = _state.setdefault("pinned_order_managers", []) + + +def _pin_order_managers(): + try: + from dpctl.utils import _order_manager as _om + except Exception as e: + warnings.warn( + f"Layer 6: could not import dpctl.utils._order_manager; the " + f"finalizer-driven variant of intel/llvm#22943 remains " + f"reachable: {e}", RuntimeWarning) + return + + cls = getattr(_om, "_SequentialOrderManager", None) + if cls is None or getattr(cls, "__gpu4pyscf_pinned__", False): + return + + _orig_init = cls.__init__ + + @functools.wraps(_orig_init) + def __init__(self, *args, **kwargs): + _orig_init(self, *args, **kwargs) + # Strong reference -> never finalized mid-run. + _pinned_order_managers.append(self) + + cls.__init__ = __init__ + cls.__gpu4pyscf_pinned__ = True + + +# ===================================================================== +# Bootstrap -- install layers 1-3. Guarded by _state["bootstrapped"] +# so a second execution of this file is a no-op. +# ===================================================================== +def _bootstrap(): + if _state["bootstrapped"]: + return + + # Layer 1: create master queues eagerly for every GPU. + for d in range(len(_gpu_devices())): + try: + _master_queue(d) + except Exception as e: + warnings.warn( + f"Failed to install master queue for device {d}: {e}", + RuntimeWarning) + + # Layer 2: replace dpctl's process-global queue cache with one that + # always returns the per-device master queue -- but only if not already + # replaced by a previous load. + try: + existing = qmgr._global_device_queue_cache + if not isinstance(existing, _MasterQueueCache): + qmgr._global_device_queue_cache = _MasterQueueCache() + probe = dpnp.zeros(4) + if not _same_queue(probe.sycl_queue, _master_queue(0)): + warnings.warn( + "Layer 2: dpnp allocation did NOT land on the master queue. " + "Queue-cache shim install may have failed.", + RuntimeWarning, + ) + except Exception as e: + warnings.warn( + f"Failed to replace dpctl device queue cache: {e}", + RuntimeWarning) + + # Layer 3: wrap creation APIs. + _wrap_with_master_queue(dpnp, _DPNP_CREATION) + # Layer 4b: arithmetic / astype results also allocate, and Layer 3's + # creation-API wrapping does not see them. Without this they are freed + # eagerly regardless of size. + _wrap_array_producers() + # Layer 4c: module-level allocating functions (vstack/concatenate/sum/...) + # are seen by neither Layer 3 nor 4b. This is the gap that left `dms` -- + # produced by cp.vstack() in _VHFOpt.get_jk's hermi==0 path and handed to + # RYS_build_jk as a raw pointer -- eligible for eager, non-queue-ordered + # sycl::free while kernels were still reading it. + _wrap_allocating_funcs() + # Layer 5: keep oneMKL's internal, GIL-holding waits away from + # pending GIL-needing host tasks. + _wrap_blocking_lapack() + _wrap_blocking_setitem() + # Layer 6: stop dpctl's thread-local order managers from running a + # blocking event wait inside a finalizer when a worker thread exits. + _pin_order_managers() + + _state["bootstrapped"] = True + + +_bootstrap() + + +# ===================================================================== +# Runtime verification -- catches regressions early. +# Uses native-handle equality (not `is`) because dpnp may rewrap a +# SyclQueue Python object around the same underlying sycl::queue. +# Runs once per process (guarded by _state["verified"]). +# ===================================================================== +def _verify_single_queue_invariant(): + if _state["verified"]: + return + + + # (1) libgsycl pointer parity per device. + for d in range(len(_gpu_devices())): + q = _master_queue(d) + libgpu.sycl_set_device(ctypes.c_int(d)) + if int(libgpu.sycl_get_queue_ptr() or 0) != _get_sycl_queue_ptr(q): # <- fixed + raise RuntimeError( + f"libgsycl queue pointer diverges from Python master on device {d}") + + # (2) main-thread dpnp allocation lands on master. + libgpu.sycl_set_device(ctypes.c_int(0)) + if not _same_queue(dpnp.zeros(4).sycl_queue, _master_queue(0)): + raise RuntimeError("main-thread dpnp allocation escaped master queue") + + # (3) worker-thread dpnp allocation lands on master. + from concurrent.futures import ThreadPoolExecutor + def _probe(): + return dpnp.zeros(4).sycl_queue + with ThreadPoolExecutor(max_workers=1) as ex: + worker_q = ex.submit(_probe).result() + if not _same_queue(worker_q, _master_queue(0)): + raise RuntimeError( + "worker-thread dpnp allocation escaped master queue -- " + "ContextVar replacement regressed" + ) + + _state["verified"] = True + + +_verify_single_queue_invariant() + + +# ===================================================================== +# Shutdown guard +# ===================================================================== +@atexit.register +def _mark_shutdown(): + """Flush and release outstanding deferred frees, then mark shutdown. + + With the polled reaper there is no host task to self-guard against + interpreter teardown (dpctl's async_dec_ref used to skip its DECREF via + Py_IsFinalizing()), so retained _Memory objects would otherwise survive to + process exit. We drain explicitly: flush any partial batch, wait once on + each master queue, then release everything. This is the only wait added by + the reaper design and it runs at process exit only, never on a hot path. + """ + try: + with _defer_free_lock: + _flush_deferred_frees_locked() + for q in list(_master_queues.values()): + q.wait() + with _defer_free_lock: + _reap_completed_locked() + except Exception: + pass + _state["shutting_down"] = True + +def _shutting_down(): + return _state["shutting_down"] + +def rebuild_dpnp_array(host, cls, state): + """Unpickle reconstructor for dpnp arrays -- see gpu4pyscf/cupy/__init__.py. + + Lives here rather than in the `cupy` shim package because pickle has to + import the reconstructor by qualified name, and the shim is registered + under a synthetic module name. + """ + arr = dpnp.asarray(host, sycl_queue=_master_queue()) + if cls is not dpnp.ndarray: + arr = arr.view(cls) + if state: + arr.__dict__.update(state) + return arr + + +def release_deferred_frees(): + """Return every deferred-free allocation to the driver, now. + + Backs `cupy.get_default_memory_pool().free_all_blocks()`. The ~30 call + sites for that in gpu4pyscf are memory-pressure relief points; with the + deferred-free reaper they are exactly where the retained batches should be + handed back. Same flush -> wait -> reap sequence as `_mark_shutdown()`, + minus the shutdown flag. + """ + try: + with _defer_free_lock: + _flush_deferred_frees_locked() + for q in list(_master_queues.values()): + q.wait() + with _defer_free_lock: + _reap_completed_locked() + except Exception: + pass + + + +# ===================================================================== +# Stream -- singleton per device, wraps master SyclQueue. +# Uses the shared _stream_cache on _state so both module copies (if any) +# hand out the same Stream instance per device. +# ===================================================================== +class Stream: + """CuPy-compatible singleton Stream wrapping the master SyclQueue. + + The constructor arguments (null, non_blocking, ptds) are accepted + for CuPy API parity but ignored -- every Stream for a given device + returns the same object, backed by the master queue. If you need + true stream-level concurrency you must step outside this shim and + create a dpctl queue directly, which voids the single-queue + invariant on that code path. + """ + + def __new__(cls, null=False, non_blocking=False, ptds=False, + *, device_id=None): + if device_id is None: + device_id = int(libgpu.sycl_get_device_id()) + with _stream_cache_lock: + s = _stream_cache.get(device_id) + if s is not None: + return s + s = object.__new__(cls) + s._device_id = device_id + s._sycl_queue = _master_queue(device_id) + s._ptr = _get_sycl_queue_ptr(s._sycl_queue) + _stream_cache[device_id] = s + return s + + def __init__(self, *a, **kw): + return + + @property + def ptr(self): + return self._ptr + + @property + def sycl_queue(self): + return self._sycl_queue + + def __int__(self): + return self._ptr + + def __enter__(self): + libgpu.sycl_set_device(ctypes.c_int(self._device_id)) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + def synchronize(self): + self._sycl_queue.wait() + + def wait_event(self, event): + # Every Stream is the same in-order master queue, so any work the + # event was recorded after is already ordered before later + # submissions on this "stream" -- nothing to wait for. + pass + + @classproperty + def null(cls): + return get_current_stream() + + +class _StreamNS: + Stream = Stream + + @staticmethod + def get_current_stream(): + return get_current_stream() + + +stream = _StreamNS() + + +def get_current_stream(): + return Stream() + + +def get_device_count(): + return len(_gpu_devices()) + + +def get_total_memory(): + return libgpu.sycl_get_total_memory() + + +def get_shared_memory(): + return libgpu.sycl_get_shared_memory() + + +def get_free_memory(): + return libgpu.sycl_get_free_memory() + +def get_compute_units(): + """Number of compute units (maps to CUDA multiProcessorCount). + + Queries the registered SYCL queue's device. + """ + return int(libgpu.sycl_get_compute_units()) + +def get_device_name(): + """Device name (maps to CUDA cudaDeviceProp::name). + + Queries the registered SYCL queue's device. + """ + buf = ctypes.create_string_buffer(256) + libgpu.sycl_get_device_name(buf, ctypes.c_int(len(buf))) + return buf.value.decode('utf-8', errors='replace') + +def get_device_id() -> int: + """Current SYCL device id (maps to CUDA cudaGetDevice). + + Queries the registered SYCL queue's device. + """ + return int(libgpu.sycl_get_device_id()) + + +# ===================================================================== +# Device -- singleton per id, backed by the shared _device_cache on _state. +# ===================================================================== +class Device: + """Singleton-per-id Device wrapper -- CuPy Device(0) semantics.""" + + def __new__(cls, device=None): + if device is None: + device = int(libgpu.sycl_get_device_id()) + elif not isinstance(device, int): + raise TypeError("device must be None or an integer device ID") + count = len(_gpu_devices()) + if device < 0 or device >= count: + raise ValueError( + f"Device index {device} out of range (available: {count})") + with _device_cache_lock: + d = _device_cache.get(device) + if d is not None: + return d + d = object.__new__(cls) + d._id = device + _master_queue(device) # ensure master exists + _device_cache[device] = d + return d + + def __init__(self, device=None): + return + + @classmethod + def get_device_id(cls) -> int: + return int(libgpu.sycl_get_device_id()) + + @property + def id(self): + return self._id + + def __enter__(self): + libgpu.sycl_set_device(ctypes.c_int(self._id)) + return self + + def __exit__(self, exc_type, exc_value, tb): + pass + + def synchronize(self): + """Drain the device's master queue -- superset of cudaDeviceSynchronize.""" + _master_queue(self._id).wait() + + @property + def mem_info(self): + return (get_free_memory(), get_total_memory()) + + +device = Device + + +# ===================================================================== +# Event -- wall-clock timing + queue.wait() sync +# +# submit_barrier() on an idle in-order queue can return a Level Zero +# 'internal event' that cannot be .wait()'d on, so we use host-clock +# for elapsed-time math and queue.wait() as the sync primitive. On +# an in-order queue, queue.wait() is a strict superset of 'wait for +# the barrier we would have submitted'. +# ===================================================================== +class Event: + """CuPy-compatible GPU timing Event.""" + + def __init__(self): + self._queue = None + self._timestamp = None + self._recorded = False + self._synced = False + + def record(self, stream=None): + if stream is not None and hasattr(stream, "sycl_queue"): + self._queue = stream.sycl_queue + else: + self._queue = _master_queue() + self._timestamp = time.perf_counter() + self._recorded = True + self._synced = False + + def synchronize(self): + if self._recorded and not self._synced and self._queue is not None: + try: + self._queue.wait() + except Exception: + pass + self._synced = True + + def query(self): + if not self._recorded: + return True + self.synchronize() + return True + + def __del__(self): + # Finalizer never touches GPU work -- queue may be in teardown. + self._queue = None + + +def get_elapsed_time(start_event, end_event): + """Elapsed wall-clock time between two recorded Events, in ms.""" + if not isinstance(start_event, Event) or not isinstance(end_event, Event): + raise TypeError("Both arguments must be cuda.Event instances.") + if not (start_event._recorded and end_event._recorded): + raise ValueError("Both events must be recorded.") + end_event.synchronize() + return (end_event._timestamp - start_event._timestamp) * 1000.0 + + +# ===================================================================== +# Address helper -- used by _Runtime.memcpy +# ===================================================================== +def _addr_of(obj) -> int: + if isinstance(obj, int): + return obj + if isinstance(obj, ctypes.c_void_p): + return int(obj.value) + ai = getattr(obj, "__sycl_usm_array_interface__", None) + if isinstance(ai, dict) and "data" in ai: + return int(ai["data"][0]) + ai = getattr(obj, "__array_interface__", None) + if isinstance(ai, dict) and "data" in ai: + return int(ai["data"][0]) + try: + return int(obj) + except Exception: + pass + if hasattr(obj, "ctypes") and hasattr(obj.ctypes, "data"): + try: + return int(obj.ctypes.data) + except Exception: + pass + raise TypeError(f"Cannot obtain address from object of type {type(obj)}") + + +# ===================================================================== +# CUDA-compat Runtime shim +# ===================================================================== +class _Runtime: + memcpyHostToHost = 0 + memcpyHostToDevice = 1 + memcpyDeviceToHost = 2 + memcpyDeviceToDevice = 3 + memcpyDefault = 4 + hostAllocMapped = 0x02 + + @staticmethod + def getDeviceCount() -> int: + return get_device_count() + + @staticmethod + def memGetInfo(): + return (get_free_memory(), get_total_memory()) + + @staticmethod + def memcpy(dst, src, nbytes, kind): + libgpu.sycl_memcpy( + ctypes.c_void_p(_addr_of(dst)), + ctypes.c_void_p(_addr_of(src)), + ctypes.c_size_t(int(nbytes))) + + @staticmethod + def getDeviceProperties(device_id: int) -> dict: + devices = _gpu_devices() + if not devices or device_id < 0 or device_id >= len(devices): + compute_units = get_compute_units() + return { + 'totalGlobalMem': get_total_memory(), + 'sharedMemPerBlock': get_shared_memory(), + 'sharedMemPerBlockOptin': get_shared_memory(), + 'name': get_device_name(), + 'maxThreadsPerBlock': 1024, + 'maxWorkGroupSize': 1024, + 'maxComputeUnits': compute_units, + 'major': 8, 'minor': 0, + 'warpSize': 32, + 'multiProcessorCount': compute_units, + } + dev = devices[device_id] + try: + warp_size = dev.sub_group_sizes[0] if dev.sub_group_sizes else 32 + except Exception: + warp_size = 32 + compute_units = dev.max_compute_units + if not compute_units or compute_units < 1: + compute_units = get_compute_units() + return { + 'totalGlobalMem': dev.global_mem_size, + 'sharedMemPerBlock': dev.local_mem_size, + 'sharedMemPerBlockOptin': dev.local_mem_size, + 'name': dev.name, + 'maxThreadsPerBlock': dev.max_work_group_size, + 'maxWorkGroupSize': dev.max_work_group_size, + 'maxComputeUnits': compute_units, + 'major': 8, 'minor': 0, + 'warpSize': warp_size, + 'multiProcessorCount': compute_units, + 'localMemSize': dev.local_mem_size, + 'globalMemSize': dev.global_mem_size, + } + + @staticmethod + def deviceCanAccessPeer(src: int, dst: int) -> bool: + return True + + +runtime = _Runtime() + + +# ===================================================================== +# Pinned-memory allocator (attached to master queue) +# ===================================================================== +def alloc_pinned_memory(nbytes, flags=None): + nbytes = int(nbytes) + q = _master_queue() + mapped = True + if flags is not None: + try: + mapped = bool(flags & runtime.hostAllocMapped) + except Exception: + mapped = True + Mem = dpmem.MemoryUSMShared if mapped else dpmem.MemoryUSMHost + return Mem(nbytes, queue=q) + + +# ===================================================================== +# cupy.cuda.memory shim +# +# CuPy exposes allocation failures as `cupy.cuda.memory.OutOfMemoryError`. +# Code that only ever *catches* that class (e.g. lib/cutensor.py) would +# otherwise raise AttributeError while unwinding, masking whatever the +# original exception actually was. Provide the name, backed by dpctl's +# USM allocation error plus the builtin MemoryError. +# ===================================================================== +import sys as _sys + +_memory_mod = types.ModuleType('cupy.cuda.memory') + +# A tuple is a valid `except` target, so this stays usable as +# `except cupy.cuda.memory.OutOfMemoryError:` while covering both the +# dpctl USM failure and a plain host MemoryError. +_oom_types = [MemoryError] +if hasattr(dpmem, 'USMAllocationError'): + _oom_types.insert(0, dpmem.USMAllocationError) +OutOfMemoryError = tuple(_oom_types) + +_memory_mod.OutOfMemoryError = OutOfMemoryError +memory = _memory_mod +_sys.modules.setdefault('cupy.cuda.memory', _memory_mod) diff --git a/gpu4pyscf/cupy/rawkernel.py b/gpu4pyscf/cupy/rawkernel.py new file mode 100644 index 000000000..f6bd9cfce --- /dev/null +++ b/gpu4pyscf/cupy/rawkernel.py @@ -0,0 +1,505 @@ +""" +`cupy.RawKernel` / `cupy.RawModule` replacement for the SYCL/dpnp backend. + +Background +---------- +gpu4pyscf embeds literal CUDA C++ kernel sources in Python strings and +JIT-compiles them with `cupy.RawKernel` (NVRTC under the hood). dpnp has no +equivalent, so those call sites die with + + AttributeError: module 'cupy' has no attribute 'RawKernel' + +dpctl 0.23 exposes the DPC++ ``kernel_compiler`` extension through +``dpctl.program.create_kernel_bundle_from_sycl_source(q, source, ...)``, which +runtime-compiles a *SYCL* (not CUDA) source string into a kernel bundle. The +kernels it can expose are "free function kernels": ``extern "C"`` functions +annotated with ``SYCL_EXTERNAL`` plus +``SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((sycl::ext::oneapi::experimental::nd_range_kernel))``, +which obtain their work-item index from +``sycl::ext::oneapi::this_work_item::get_nd_item()`` rather than from a +kernel-lambda parameter. + +That is close enough to CUDA's ``__global__`` model that the *existing CUDA +kernel bodies compile unmodified* once a small compatibility prelude is +prepended: ``__global__`` becomes the annotation pair, and +``threadIdx``/``blockIdx``/``blockDim``/``gridDim`` become tiny structs that +pull from ``get_nd_item<3>()``. So this shim does not require rewriting any +kernel source in the codebase -- it prepends the prelude and hands the result +to dpctl. + +Dimension convention +-------------------- +CUDA's fastest-varying dimension is ``x``; SYCL's is the *last* index of an +``nd_range``. A 3D ``nd_item`` is therefore indexed as + + CUDA .x -> sycl dim 2 + CUDA .y -> sycl dim 1 + CUDA .z -> sycl dim 0 + +and a CUDA ``grid=(gx,gy,gz)``, ``block=(bx,by,bz)`` launch becomes a SYCL +global range ``[gz*bz, gy*by, gx*bx]`` with local range ``[bz, by, bx]``. Both +the prelude and :func:`_launch` implement exactly this mapping; every kernel is +compiled as ``nd_range_kernel<3>`` so one prelude covers 1D/2D/3D launches. + +Static shared memory +-------------------- +``__shared__ T name[N];`` maps to +``static sycl::ext::oneapi::experimental::work_group_static name;``. The +array bound moves *inside* the template argument, so an object-like macro +cannot express it; :func:`_rewrite_shared` does the rewrite textually instead. +Dynamic shared memory (CuPy's ``shared_mem=`` launch argument) has no +equivalent here and is rejected. +""" + +import ctypes +import re +import threading + +import numpy as np + +import dpctl +import dpctl.memory as dpmem +import dpctl.program as dpprog + + +__all__ = ["RawKernel", "RawModule", "is_available"] + + +# --------------------------------------------------------------------- +# CUDA -> SYCL free-function-kernel compatibility prelude. +# +# Prepended verbatim to every source string. Keeps the CUDA kernel bodies +# already in the codebase compilable as-is. +# --------------------------------------------------------------------- +_CUDA_COMPAT_PRELUDE = r''' +#include +#include +#include + +namespace syclext = sycl::ext::oneapi::experimental; + +namespace g4p_compat { + +static inline sycl::nd_item<3> _it() { + return sycl::ext::oneapi::this_work_item::get_nd_item<3>(); +} + +// CUDA .x is the fastest-varying dim; in SYCL that is the LAST index. +struct _ThreadIdx { + int x, y, z; + _ThreadIdx() { auto i = _it(); + x = (int)i.get_local_id(2); y = (int)i.get_local_id(1); z = (int)i.get_local_id(0); } +}; +struct _BlockIdx { + int x, y, z; + _BlockIdx() { auto i = _it(); + x = (int)i.get_group(2); y = (int)i.get_group(1); z = (int)i.get_group(0); } +}; +struct _BlockDim { + int x, y, z; + _BlockDim() { auto i = _it(); + x = (int)i.get_local_range(2); y = (int)i.get_local_range(1); z = (int)i.get_local_range(0); } +}; +struct _GridDim { + int x, y, z; + _GridDim() { auto i = _it(); + x = (int)i.get_group_range(2); y = (int)i.get_group_range(1); z = (int)i.get_group_range(0); } +}; + +} // namespace g4p_compat + +#define threadIdx (g4p_compat::_ThreadIdx()) +#define blockIdx (g4p_compat::_BlockIdx()) +#define blockDim (g4p_compat::_BlockDim()) +#define gridDim (g4p_compat::_GridDim()) +#define __syncthreads() sycl::group_barrier(g4p_compat::_it().get_group()) +#define __syncwarp() sycl::group_barrier(g4p_compat::_it().get_sub_group()) +#define __global__ SYCL_EXTERNAL SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclext::nd_range_kernel<3>)) +#define __device__ inline +#define __forceinline__ inline +#define __restrict__ __restrict__ + +// cupy/complex.cuh provides `complex` in the global namespace; std::complex +// is device-usable under DPC++ and has the same interface for our uses. +using std::complex; +using std::exp; +using std::abs; +using std::fabs; +using std::sqrt; +using std::pow; +using std::log; +using std::sin; +using std::cos; + +// CUDA-flavoured math spellings used by the embedded kernels. Defined as +// macros, not functions: glibc's already declares a *narrowing* +// `float fsqrt(double)` (C2x TS-18661), so a `static inline double fsqrt` +// is a redeclaration error -- and picking up glibc's would silently halve +// the precision. +#define fsqrt(x) (sycl::sqrt((double)(x))) +#define rsqrt(x) (sycl::rsqrt((double)(x))) +#define rsqrtf(x) (sycl::rsqrt((float)(x))) +#define __fdividef(a, b) ((float)(a) / (float)(b)) + +// CUDA built-in vector types used by the embedded kernels. Plain structs +// with the same layout/member names as the CUDA originals -- no alignment +// annotations, since these kernels only ever access them element-wise. +struct double2 { double x, y; }; +struct float2 { float x, y; }; +struct int2 { int x, y; }; +static inline double2 make_double2(double x, double y) { return double2{x, y}; } +static inline float2 make_float2(float x, float y) { return float2{x, y}; } +static inline int2 make_int2(int x, int y) { return int2{x, y}; } + +// CUDA atomicAdd -> sycl::atomic_ref (relaxed, device scope) compound add. +template +static inline T atomicAdd(T *addr, T val) { + sycl::atomic_ref ref(*addr); + return ref.fetch_add(val); +} + +// CUDA warp-shuffle-down -> SYCL sub-group shuffle-down. The mask argument +// is accepted for source compatibility and ignored (matches full-warp use +// in every embedded kernel; sub-group size is fixed by the compiled device). +template +static inline T __shfl_down_sync(unsigned mask, T val, unsigned delta) { + (void)mask; + return sycl::shift_group_left(g4p_compat::_it().get_sub_group(), val, delta); +} +template +static inline T __shfl_up_sync(unsigned mask, T val, unsigned delta) { + (void)mask; + return sycl::shift_group_right(g4p_compat::_it().get_sub_group(), val, delta); +} +template +static inline T __shfl_xor_sync(unsigned mask, T val, int lane_mask) { + (void)mask; + return sycl::permute_group_by_xor(g4p_compat::_it().get_sub_group(), val, lane_mask); +} +template +static inline T __shfl_sync(unsigned mask, T val, int src_lane) { + (void)mask; + return sycl::select_from_group(g4p_compat::_it().get_sub_group(), val, src_lane); +} +''' + +# Source lines matched here are stripped before compiling: CUDA-only headers +# that have no SYCL counterpart but whose contents the prelude already covers. +_STRIP_INCLUDES = ( + "#include ", + "#include ", + "#include ", +) + + +def is_available(): + """True if this dpctl/DPC++ build can runtime-compile SYCL source.""" + try: + return bool(dpprog.is_sycl_source_compilation_available()) + except Exception: + return False + + +# `__shared__ [];` +# -> `static work_group_static<[]> ;` +# The array bound has to move inside the template argument, so this cannot be +# a macro. may be any constant expression (e.g. TILE*TILE), and +# any multi-word builtin ("unsigned short"). +_SHARED_ARRAY_RE = re.compile( + r"__shared__\s+([A-Za-z_][\w:\s]*?)\s+([A-Za-z_]\w*)\s*\[([^\]]+)\]\s*;" +) +# Scalar form: `__shared__ ;` +_SHARED_SCALAR_RE = re.compile( + r"__shared__\s+([A-Za-z_][\w:\s]*?)\s+([A-Za-z_]\w*)\s*;" +) + + +def _rewrite_shared(code): + """Translate CUDA static __shared__ declarations to work_group_static.""" + code = _SHARED_ARRAY_RE.sub( + r"static syclext::work_group_static<\1[\3]> \2;", code) + code = _SHARED_SCALAR_RE.sub( + r"static syclext::work_group_static<\1> \2;", code) + if "__shared__" in code: + raise NotImplementedError( + "unrecognised __shared__ declaration form; dynamic (extern) " + "shared memory is not supported by the SYCL " + "free-function-kernel backend" + ) + return code + + +def _preprocess(code): + """Strip CUDA-only includes, rewrite __shared__, prepend the prelude.""" + for inc in _STRIP_INCLUDES: + code = code.replace(inc, "") + code = _rewrite_shared(code) + return _CUDA_COMPAT_PRELUDE + "\n" + code + + +def _default_queue(): + """The queue dpnp allocations live on, so launches stay correctly ordered. + + gpu4pyscf.cupy.cuda installs a per-device singleton in-order master queue + and injects it into every dpnp allocation. Use it when available so a + JIT'd kernel is enqueued behind the dpnp work that produced its inputs. + """ + try: + from . import cuda as _cuda + return _cuda._master_queue() + except Exception: + return dpctl.SyclQueue() + + +# --------------------------------------------------------------------- +# Kernel-argument marshaling +# +# dpctl.SyclQueue.submit accepts (see dpctl/_sycl_queue.pyx _populate_args): +# ctypes c_char/c_uint8/c_short/c_ushort/c_int/c_uint/c_longlong/ +# c_ulonglong/c_float/c_double -> by-value scalars +# dpctl.memory._Memory -> USM pointer +# LocalAccessor / WorkGroupMemory / RawKernelArg -> unused here +# Notably c_long is NOT accepted, so 64-bit ints must be c_longlong. +# --------------------------------------------------------------------- +_NP_TO_CTYPES = { + np.dtype(np.int8): ctypes.c_char, + np.dtype(np.uint8): ctypes.c_uint8, + np.dtype(np.int16): ctypes.c_short, + np.dtype(np.uint16): ctypes.c_ushort, + np.dtype(np.int32): ctypes.c_int, + np.dtype(np.uint32): ctypes.c_uint, + np.dtype(np.int64): ctypes.c_longlong, + np.dtype(np.uint64): ctypes.c_ulonglong, + np.dtype(np.float32): ctypes.c_float, + np.dtype(np.float64): ctypes.c_double, +} + + +def _as_usm_memory(arr): + """USM buffer for a dpnp/dpctl array, honouring any view offset. + + `arr.usm_data` is the *base* allocation and ignores slicing offsets; + `dpctl.memory.as_usm_memory` goes through __sycl_usm_array_interface__ and + yields a buffer whose pointer is the array's own first element. + """ + usm = arr.get_array() if hasattr(arr, "get_array") else arr + flags = usm.flags + if not (flags["C_CONTIGUOUS"] or flags["F_CONTIGUOUS"]): + raise ValueError( + "RawKernel arguments must be contiguous; got an array with " + f"shape {usm.shape} strides {usm.strides}" + ) + return dpmem.as_usm_memory(usm) + + +def _marshal(arg): + # Already-marshaled USM buffers pass straight through. Checked FIRST: + # _Memory also exposes __sycl_usm_array_interface__, so the array branch + # below would otherwise claim it and then fail on the missing `.flags`. + if isinstance(arg, (dpmem.MemoryUSMDevice, dpmem.MemoryUSMShared, + dpmem.MemoryUSMHost)): + return arg + # Arrays (dpnp.ndarray, dpctl.tensor.usm_ndarray, anything USM-aware) + if hasattr(arg, "__sycl_usm_array_interface__") or hasattr(arg, "get_array"): + return _as_usm_memory(arg) + if isinstance(arg, ctypes._SimpleCData): + return arg + # numpy scalars and 0-d arrays carry their own C type + if isinstance(arg, np.generic) or (isinstance(arg, np.ndarray) and arg.ndim == 0): + dt = np.dtype(arg.dtype) + try: + return _NP_TO_CTYPES[dt](arg.item()) + except KeyError: + raise TypeError(f"unsupported RawKernel scalar dtype {dt}") + # Plain Python scalars: match CuPy, which widens int -> long long. + if isinstance(arg, bool): + return ctypes.c_char(int(arg)) + if isinstance(arg, int): + return ctypes.c_longlong(arg) + if isinstance(arg, float): + return ctypes.c_double(arg) + raise TypeError(f"unsupported RawKernel argument type {type(arg)}") + + +def _cuda_ranges_to_sycl(grid, block): + """CUDA (gx,gy,gz)/(bx,by,bz) -> SYCL global/local range lists [z,y,x].""" + def _pad(t): + t = tuple(int(v) for v in (t if isinstance(t, (tuple, list)) else (t,))) + return t + (1,) * (3 - len(t)) + gx, gy, gz = _pad(grid) + bx, by, bz = _pad(block) + return [gz * bz, gy * by, gx * bx], [bz, by, bx] + + +# --------------------------------------------------------------------- +# Compiled-bundle cache +# +# kernel_compiler invocation costs seconds, so bundles are memoised on +# (source, options, device). SYCL_CACHE_PERSISTENT additionally caches the +# device binary across processes. +# --------------------------------------------------------------------- +_bundle_cache = {} +_bundle_lock = threading.Lock() + + +def _get_bundle(code, options, queue): + key = (code, tuple(options), queue.sycl_device.filter_string) + with _bundle_lock: + bundle = _bundle_cache.get(key) + if bundle is not None: + return bundle + if not is_available(): + raise RuntimeError( + "Runtime SYCL kernel compilation is unavailable in this " + "DPC++/dpctl build (dpctl.program." + "is_sycl_source_compilation_available() is False). Kernels " + "that used cupy.RawKernel cannot be JIT-compiled; they must " + "be built ahead of time into gpu4pyscf's native libraries." + ) + bundle = dpprog.create_kernel_bundle_from_sycl_source( + queue, _preprocess(code), headers=[], registered_names=[], + copts=list(options)) + _bundle_cache[key] = bundle + return bundle + + +class _KernelBase: + """Shared launch machinery for RawKernel and RawModule.get_function.""" + + def _resolve(self, queue): + raise NotImplementedError + + def __call__(self, grid, block, args, shared_mem=0, stream=None, + enable_cooperative_groups=False): + if shared_mem: + raise NotImplementedError( + "dynamic shared memory (shared_mem=) is not supported by the " + "SYCL free-function-kernel backend") + queue = None + for a in args: + q = getattr(a, "sycl_queue", None) + if q is not None: + queue = q + break + if queue is None: + queue = _default_queue() + + kernel = self._resolve(queue) + kargs = [_marshal(a) for a in args] + n_expected = kernel.num_args + if n_expected != len(kargs): + raise ValueError( + f"kernel {self.name!r} expects {n_expected} arguments, " + f"got {len(kargs)}") + gS, lS = _cuda_ranges_to_sycl(grid, block) + # Blocking submit: keeps USM argument buffers alive for the duration + # of the launch without a separate lifetime-tracking mechanism, and + # matches the ordering the CUDA call sites assume. + queue.submit(kernel, kargs, gS, lS) + return None + + +class RawKernel(_KernelBase): + """Drop-in replacement for `cupy.RawKernel` on the SYCL backend. + + Signature mirrors CuPy's; `backend`/`translate_cucomplex`/`jitify` and + friends are accepted and ignored so call sites need no edits. + """ + + def __init__(self, code, name, options=(), backend="nvrtc", + translate_cucomplex=False, enable_cooperative_groups=False, + jitify=False, **kwargs): + self.code = code + self.name = name + self.options = tuple(options) + self._kernels = {} # device filter string -> SyclKernel + + def _resolve(self, queue): + key = queue.sycl_device.filter_string + krn = self._kernels.get(key) + if krn is None: + bundle = _get_bundle(self.code, self.options, queue) + if not bundle.has_sycl_kernel(self.name): + raise RuntimeError( + f"compiled bundle has no kernel named {self.name!r}; the " + "kernel must be declared extern \"C\"") + krn = bundle.get_sycl_kernel(self.name) + self._kernels[key] = krn + return krn + + # CuPy attribute-compatibility surface + @property + def kernel(self): + return self._resolve(_default_queue()) + + @property + def max_threads_per_block(self): + return self._resolve(_default_queue()).work_group_size + + @property + def num_regs(self): + return 0 + + @property + def shared_size_bytes(self): + return 0 + + @property + def local_size_bytes(self): + return self._resolve(_default_queue()).private_mem_size + + @property + def attributes(self): + k = self._resolve(_default_queue()) + return { + "max_threads_per_block": k.work_group_size, + "local_size_bytes": k.private_mem_size, + "preferred_work_group_size_multiple": + k.preferred_work_group_size_multiple, + } + + +class _ModuleKernel(_KernelBase): + def __init__(self, module, name): + self._module = module + self.name = name + + def _resolve(self, queue): + key = queue.sycl_device.filter_string + krn = self._module._kernels.get((key, self.name)) + if krn is None: + bundle = _get_bundle(self._module.code, self._module.options, queue) + if not bundle.has_sycl_kernel(self.name): + raise RuntimeError( + f"compiled bundle has no kernel named {self.name!r}") + krn = bundle.get_sycl_kernel(self.name) + self._module._kernels[(key, self.name)] = krn + return krn + + +class RawModule: + """Drop-in replacement for `cupy.RawModule` on the SYCL backend. + + One source string may define several `extern "C" __global__` kernels; each + is fetched with :meth:`get_function`, as in CuPy. + """ + + def __init__(self, code=None, path=None, options=(), backend="nvrtc", + translate_cucomplex=False, enable_cooperative_groups=False, + name_expressions=None, jitify=False, **kwargs): + if code is None: + if path is None: + raise TypeError("RawModule requires either code= or path=") + with open(path) as f: + code = f.read() + self.code = code + self.options = tuple(options) + self._kernels = {} + self._functions = {} + + def get_function(self, name): + fn = self._functions.get(name) + if fn is None: + fn = _ModuleKernel(self, name) + self._functions[name] = fn + return fn diff --git a/gpu4pyscf/cupyx/__init__.py b/gpu4pyscf/cupyx/__init__.py new file mode 100644 index 000000000..48895b428 --- /dev/null +++ b/gpu4pyscf/cupyx/__init__.py @@ -0,0 +1,27 @@ +# cupyx/__init__.py +# Fake cupyx package root + +import numpy as _np +import cupy as _cupy + + +def empty_pinned(shape, dtype=_np.float64, order='C'): + '''Equivalent of cupyx.empty_pinned: allocate an uninitialized host + ndarray backed by pinned (page-locked) memory, for fast host<->device + transfers. Backed by cupy.cuda.alloc_pinned_memory, the same pinned + allocator already used elsewhere in the shim (see + gpu4pyscf/lib/cupy_helper.py:pin_memory and cupy/cuda.py + :alloc_pinned_memory). + ''' + shape = tuple(int(s) for s in shape) if isinstance(shape, (tuple, list)) else (int(shape),) + dtype = _np.dtype(dtype) + nbytes = int(_np.prod(shape)) * dtype.itemsize if shape else dtype.itemsize + mem = _cupy.cuda.alloc_pinned_memory(nbytes) + return _np.ndarray(shape, dtype=dtype, buffer=mem, order=order) + + +def zeros_pinned(shape, dtype=_np.float64, order='C'): + '''Equivalent of cupyx.zeros_pinned: like empty_pinned but zero-filled.''' + out = empty_pinned(shape, dtype=dtype, order=order) + out.fill(0) + return out diff --git a/gpu4pyscf/cupyx/scipy/__init__.py b/gpu4pyscf/cupyx/scipy/__init__.py new file mode 100644 index 000000000..342279272 --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/__init__.py @@ -0,0 +1,7 @@ +# cupyx/scipy/__init__.py + +from . import fft # expose cupyx.scipy.fft +from . import sparse # expose cupyx.scipy.sparse +from . import special # expose cupyx.scipy.special + +__all__ = ["fft", "sparse", "special"] diff --git a/gpu4pyscf/cupyx/scipy/cupyx_linalg.py b/gpu4pyscf/cupyx/scipy/cupyx_linalg.py new file mode 100644 index 000000000..cdb28fb5e --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/cupyx_linalg.py @@ -0,0 +1,101 @@ +import numpy + +import cupy +from cupy.cuda import cublas +from cupy.cuda import device +from cupy.linalg import _util +from cupyx.scipy.linalg import _uarray + + +@_uarray.implements('solve_triangular') +def solve_triangular(a, b, trans=0, lower=False, unit_diagonal=False, + overwrite_b=False, check_finite=False): + """Solve the equation a x = b for x, assuming a is a triangular matrix. + + Args: + a (cupy.ndarray): The matrix with dimension ``(M, M)``. + b (cupy.ndarray): The matrix with dimension ``(M,)`` or + ``(M, N)``. + lower (bool): Use only data contained in the lower triangle of ``a``. + Default is to use upper triangle. + trans (0, 1, 2, 'N', 'T' or 'C'): Type of system to solve: + + - *'0'* or *'N'* -- :math:`a x = b` + - *'1'* or *'T'* -- :math:`a^T x = b` + - *'2'* or *'C'* -- :math:`a^H x = b` + + unit_diagonal (bool): If ``True``, diagonal elements of ``a`` are + assumed to be 1 and will not be referenced. + overwrite_b (bool): Allow overwriting data in b (may enhance + performance) + check_finite (bool): Whether to check that the input matrices contain + only finite numbers. Disabling may give a performance gain, but may + result in problems (crashes, non-termination) if the inputs do + contain infinities or NaNs. + + Returns: + cupy.ndarray: + The matrix with dimension ``(M,)`` or ``(M, N)``. + + .. seealso:: :func:`scipy.linalg.solve_triangular` + """ + + _util._assert_cupy_array(a, b) + + if len(a.shape) != 2 or a.shape[0] != a.shape[1]: + raise ValueError('expected square matrix') + if len(a) != len(b): + raise ValueError('incompatible dimensions') + + # Cast to float32 or float64 + if a.dtype.char in 'fd': + dtype = a.dtype + else: + dtype = numpy.promote_types(a.dtype.char, 'f') + + a = cupy.array(a, dtype=dtype, order='F', copy=False) + b = cupy.array(b, dtype=dtype, order='F', copy=(not overwrite_b)) + + if check_finite: + if a.dtype.kind == 'f' and not cupy.isfinite(a).all(): + raise ValueError( + 'array must not contain infs or NaNs') + if b.dtype.kind == 'f' and not cupy.isfinite(b).all(): + raise ValueError( + 'array must not contain infs or NaNs') + + m, n = (b.size, 1) if b.ndim == 1 else b.shape + cublas_handle = device.get_cublas_handle() + + if dtype == 'f': + trsm = cublas.strsm + elif dtype == 'd': + trsm = cublas.dtrsm + elif dtype == 'F': + trsm = cublas.ctrsm + else: # dtype == 'D' + trsm = cublas.ztrsm + one = numpy.array(1, dtype=dtype) + + if lower: + uplo = cublas.CUBLAS_FILL_MODE_LOWER + else: + uplo = cublas.CUBLAS_FILL_MODE_UPPER + + if trans == 'N': + trans = cublas.CUBLAS_OP_N + elif trans == 'T': + trans = cublas.CUBLAS_OP_T + elif trans == 'C': + trans = cublas.CUBLAS_OP_C + + if unit_diagonal: + diag = cublas.CUBLAS_DIAG_UNIT + else: + diag = cublas.CUBLAS_DIAG_NON_UNIT + + trsm( + cublas_handle, cublas.CUBLAS_SIDE_LEFT, uplo, + trans, diag, + m, n, one.ctypes.data, a.data.ptr, m, b.data.ptr, m) + return b diff --git a/gpu4pyscf/cupyx/scipy/fft/__init__.py b/gpu4pyscf/cupyx/scipy/fft/__init__.py new file mode 100644 index 000000000..cbd3c950d --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/fft/__init__.py @@ -0,0 +1,25 @@ +# cupyx/scipy/fft/__init__.py +""" +Minimal shim mapping cupyx.scipy.fft -> dpnp.fft +Exposes only: fftn, ifftn, fftfreq +""" + +import dpnp as _xp +_xp_fft = _xp.fft + + +def fftn(a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None): + """CuPy/SciPy-compatible fftn; extra args are accepted but ignored.""" + return _xp_fft.fftn(a, s=s, axes=axes, norm=norm) + + +def ifftn(a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None): + """CuPy/SciPy-compatible ifftn; extra args are accepted but ignored.""" + return _xp_fft.ifftn(a, s=s, axes=axes, norm=norm) + + +def fftfreq(n, d=1.0): + """CuPy/SciPy-compatible fftfreq.""" + return _xp_fft.fftfreq(n, d=d) + +__all__ = ["fftn", "ifftn", "fftfreq"] diff --git a/gpu4pyscf/cupyx/scipy/linalg.py b/gpu4pyscf/cupyx/scipy/linalg.py new file mode 100644 index 000000000..14ab67645 --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/linalg.py @@ -0,0 +1,325 @@ +# Copyright 2021-2024 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import cmath +import math +import dpnp +import ctypes + +# Load the oneMKL helper shared library from gpu4pyscf/lib/ +_lib_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'lib')) +libonemkl = ctypes.CDLL(os.path.join(_lib_dir, 'libonemkl_helper.so')) + +libonemkl.onemkl_trsm.argtypes = [ + ctypes.c_void_p, # A + ctypes.c_void_p, # B + ctypes.c_int, # m + ctypes.c_int, # n + ctypes.c_int, # lda + ctypes.c_int, # ldb + ctypes.c_int, # lower + ctypes.c_int, # trans + ctypes.c_int # unit_diagonal +] +libonemkl.onemkl_trsm.restype = None + +libonemkl.onemkl_strsm.argtypes = [ + ctypes.c_void_p, # A + ctypes.c_void_p, # B + ctypes.c_int, # m + ctypes.c_int, # n + ctypes.c_int, # lda + ctypes.c_int, # ldb + ctypes.c_int, # lower + ctypes.c_int, # trans + ctypes.c_int # unit_diagonal +] +libonemkl.onemkl_strsm.restype = None + +########################################################################################################### + +def solve_triangular(a, b, trans=0, lower=False, unit_diagonal=False, + overwrite_b=False, check_finite=False): + # print("inputs from a in linalg.py: ", a) + # print("inputs from b in linalg.py: ", b) + """ + Solve the equation a x = b for x, assuming a is a triangular matrix using dpnp + oneMKL. + + Args: + a (dpnp.ndarray): The matrix with dimension (M, M). + b (dpnp.ndarray): The matrix with dimension (M,) or (M, N). + lower (bool): Use lower triangle if True, otherwise upper. + trans (0, 1, 2, 'N', 'T', 'C'): Type of system to solve: + - 0 or 'N' -- a x = b + - 1 or 'T' -- a^T x = b + - 2 or 'C' -- a^H x = b + unit_diagonal (bool): If True, assumes diagonal elements are all 1. + overwrite_b (bool): Allow overwriting data in b (may enhance performance). + check_finite (bool): Whether to check for NaNs or Infs. + + Returns: + dpnp.ndarray: Solution x with same shape as b. + """ + + # Check shapes + if a.ndim != 2 or a.shape[0] != a.shape[1]: + raise ValueError("Matrix 'a' must be square.") + if a.shape[0] != b.shape[0]: + raise ValueError("Dimensions of 'a' and 'b' do not align.") + + # Handle trans parameter + trans_flag = 0 + if trans in [1, 'T']: + trans_flag = 1 + elif trans in [2, 'C']: + raise NotImplementedError("Hermitian transpose not supported") + + # Type promotion - cast to float32 or float64 + if a.dtype.char in 'fdFD': + dtype = a.dtype + else: + dtype = dpnp.promote_types(a.dtype.char, 'f') + + # FIX: Remove copy=False to allow dpnp to copy when necessary + # If conversion to F-order or dtype change is needed, dpnp will copy automatically + a = dpnp.array(a, dtype=dtype, order='F') + + # For b, handle overwrite_b properly + # If overwrite_b=True and no conversion needed, don't copy + # Otherwise, copy as needed + if overwrite_b: + # Try to avoid copy, but allow it if necessary + b = dpnp.asarray(b, dtype=dtype) + # Convert to F-order if needed (may copy) + if not b.flags['F_CONTIGUOUS']: + b = dpnp.asfortranarray(b) + else: + # Always make a copy + b = dpnp.array(b, dtype=dtype, order='F', copy=True) + + if check_finite: + if a.dtype.kind == 'f' and not dpnp.isfinite(a).all(): + raise ValueError('A array must not contain infs or NaNs') + if b.dtype.kind == 'f' and not dpnp.isfinite(b).all(): + raise ValueError('B array must not contain infs or NaNs') + + m, n = (b.size, 1) if b.ndim == 1 else b.shape + + if dtype.char == 'f': + trsm_func = libonemkl.onemkl_strsm + else: + trsm_func = libonemkl.onemkl_trsm + + trsm_func(ctypes.cast(a.data.ptr, ctypes.c_void_p), + ctypes.cast(b.data.ptr, ctypes.c_void_p), + ctypes.c_int(m), ctypes.c_int(n), + ctypes.c_int(m), ctypes.c_int(m), + ctypes.c_int(lower), ctypes.c_int(trans_flag), + ctypes.c_int(unit_diagonal)) + return b + +########################################################################################################### + + +def block_diag(*arrs): + """Create a block diagonal matrix from provided arrays. + + Given the inputs ``A``, ``B``, and ``C``, the output will have these + arrays arranged on the diagonal:: + + [A, 0, 0] + [0, B, 0] + [0, 0, C] + + Args: + A, B, C, ... (cupy.ndarray): Input arrays. A 1-D array of length ``n`` + is treated as a 2-D array with shape ``(1,n)``. + + Returns: + (cupy.ndarray): Array with ``A``, ``B``, ``C``, ... on the diagonal. + Output has the same dtype as ``A``. + + .. seealso:: :func:`scipy.linalg.block_diag` + """ + if not arrs: + return dpnp.empty((1, 0)) + + # --- NEW: unwrap gpu4pyscf wrappers like DPNPArrayWithTag --- + def _unwrap_dpnp_like(a): + base = getattr(a, "array", None) + if isinstance(base, dpnp.ndarray): + return base + return a + + arrs = tuple(_unwrap_dpnp_like(a) for a in arrs) + # --- END NEW --- + + # Convert to 2D and check + if len(arrs) == 1: + arrs = (dpnp.atleast_2d(*arrs),) + else: + arrs = dpnp.atleast_2d(*arrs) + if any(a.ndim != 2 for a in arrs): + bad = [k for k in range(len(arrs)) if arrs[k].ndim != 2] + raise ValueError('arguments in the following positions have dimension ' + 'greater than 2: {}'.format(bad)) + + shapes = tuple(a.shape for a in arrs) + shape = tuple(sum(x) for x in zip(*shapes)) + out = dpnp.zeros(shape, dtype=dpnp.result_type(*arrs)) + r, c = 0, 0 + for arr in arrs: + rr, cc = arr.shape + out[r:r + rr, c:c + cc] = arr + r += rr + c += cc + return out + +########################################################################################################### + +def lu(a, permute_l=False, overwrite_a=False, check_finite=True, + p_indices=False): + return dpnp.scipy.linalg.lu( + a, + permute_l=permute_l, + overwrite_a=overwrite_a, + check_finite=check_finite, + p_indices=p_indices, + ) + + +def lu_factor(a, overwrite_a=False, check_finite=True): + return dpnp.scipy.linalg.lu_factor( + a, + overwrite_a=overwrite_a, + check_finite=check_finite, + ) + + +def lu_solve(lu_and_piv, b, trans=0, overwrite_b=False, check_finite=True): + return dpnp.scipy.linalg.lu_solve( + lu_and_piv, + b, + trans=trans, + overwrite_b=overwrite_b, + check_finite=check_finite, + ) + +########################################################################################################### + +# Source: https://github.com/cupy/cupy/blob/main/cupyx/scipy/linalg/_matfuncs.py#L45 + +th13 = 5.37 + +b = [64764752532480000., + 32382376266240000., + 7771770303897600., + 1187353796428800., + 129060195264000., + 10559470521600., + 670442572800., + 33522128640., + 1323241920., + 40840800., + 960960., + 16380., + 182., + 1.,] + +def expm(a): + """Compute the matrix exponential. + + Parameters + ---------- + a : dpnp.ndarray, 2D + + Returns + ------- + matrix exponential of `a` + + Notes + ----- + Uses (a simplified) version of Algorithm 2.3 of [1]_: + a [13 / 13] Pade approximant with scaling and squaring. + + Simplifications: + + * we always use a [13/13] approximate + * no matrix balancing + + References + ---------- + .. [1] N. Higham, SIAM J. MATRIX ANAL. APPL. Vol. 26(4), p. 1179 (2005) + https://doi.org/10.1137/04061101X + + """ + if a.size == 0: + return dpnp.zeros((0, 0), dtype=a.dtype) + + n = a.shape[0] + + # follow scipy.linalg.expm dtype handling + a_dtype = a.dtype if dpnp.issubdtype( + a.dtype, dpnp.inexact) else dpnp.float64 + + # try reducing the norm + mu = dpnp.diag(a).sum() / n + A = a - dpnp.eye(n, dtype=a_dtype) * mu + + # scale factor + nrmA = dpnp.linalg.norm(A, ord=1).item() + + scale = nrmA > th13 + if scale: + s = int(math.ceil(math.log2(float(nrmA) / th13))) + 1 + else: + s = 1 + + A /= 2**s + + # compute [13/13] Pade approximant + A2 = A @ A + A4 = A2 @ A2 + A6 = A2 @ A4 + + E = dpnp.eye(A.shape[0], dtype=a_dtype) + bb = dpnp.asarray(b, dtype=a_dtype) + + u1, u2, v1, v2 = _expm_inner(E, A, A2, A4, A6, bb) + u = A @ (A6 @ u1 + u2) + v = A6 @ v1 + v2 + + r13 = dpnp.linalg.solve(-u + v, u + v) + + # squaring + x = r13 + for _ in range(s): + x = x @ x + + # undo preprocessing + emu = cmath.exp(mu) if dpnp.issubdtype( + mu.dtype, dpnp.complexfloating) else math.exp(mu) + x *= emu + + return x + +def _expm_inner(E, A, A2, A4, A6, b): + u1 = b[13]*A6 + b[11]*A4 + b[9]*A2 + u2 = b[7]*A6 + b[5]*A4 + b[3]*A2 + b[1]*E + + v1 = b[12]*A6 + b[10]*A4 + b[8]*A + v2 = b[6]*A6 + b[4]*A4 + b[2]*A2 + b[0]*E + return u1, u2, v1, v2 +########################################################################################################### diff --git a/gpu4pyscf/cupyx/scipy/sparse/__init__.py b/gpu4pyscf/cupyx/scipy/sparse/__init__.py new file mode 100644 index 000000000..f443fe0de --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/sparse/__init__.py @@ -0,0 +1,4 @@ +# cupyx/scipy/sparse/__init__.py +# Fake cupyx.scipy.sparse — forwards to dpnp.scipy.sparse +from . import linalg +__all__ = ["linalg"] diff --git a/gpu4pyscf/cupyx/scipy/sparse/linalg.py b/gpu4pyscf/cupyx/scipy/sparse/linalg.py new file mode 100644 index 000000000..0950c363d --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/sparse/linalg.py @@ -0,0 +1,30 @@ +# gpu4pyscf/gpu4pyscf/cupyx/scipy/sparse/linalg.py +# +# Shim: re-export dpnp equivalents under the cupyx namespace, with a +# tol -> rtol translation for minres since cupyx uses `tol` but dpnp +# (matching modern scipy) uses `rtol`. + +from dpnp.scipy.sparse.linalg import ( + LinearOperator, + cg, + gmres, + minres as _dpnp_minres, +) + + + +def minres(A, b, x0=None, *, shift=0.0, tol=1e-5, maxiter=None, + M=None, callback=None, check=False): + """cupyx-style minres. Translates tol -> rtol for dpnp.""" + return _dpnp_minres( + A, b, x0, + rtol=tol, + shift=shift, + maxiter=maxiter, + M=M, + callback=callback, + check=check, + ) + + +__all__ = ["LinearOperator", "minres", "cg", "gmres"] diff --git a/gpu4pyscf/cupyx/scipy/special/__init__.py b/gpu4pyscf/cupyx/scipy/special/__init__.py new file mode 100644 index 000000000..051759edc --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/special/__init__.py @@ -0,0 +1,5 @@ +# cupyx/scipy/special/__init__.py +# Shim: re-export dpnp.scipy.special under the cupyx namespace. +from dpnp.scipy import special as _dp_special +from dpnp.scipy.special import * # noqa: F401,F403 +__all__ = [n for n in dir(_dp_special) if not n.startswith("_")] diff --git a/gpu4pyscf/cupyx/scipy/special/linalg.py b/gpu4pyscf/cupyx/scipy/special/linalg.py new file mode 100644 index 000000000..a3b1f42ff --- /dev/null +++ b/gpu4pyscf/cupyx/scipy/special/linalg.py @@ -0,0 +1,10 @@ +# cupyx/scipy/sparse/linalg.py +# Shim: re-export dpnp equivalents under the cupyx namespace. +from dpnp.scipy.sparse.linalg import ( + LinearOperator, + minres, + cg, + gmres, +) + +__all__ = ["LinearOperator", "minres", "cg", "gmres"] diff --git a/gpu4pyscf/df/tests/test_df_hessian.py b/gpu4pyscf/df/tests/test_df_hessian.py index a5a365cf7..aad4fc2da 100644 --- a/gpu4pyscf/df/tests/test_df_hessian.py +++ b/gpu4pyscf/df/tests/test_df_hessian.py @@ -358,7 +358,11 @@ def test_unstable_j2c(self): mo_occ = mf.mo_occ test_hessian_round2 = hobj.partial_hess_elec(mo_energy, mo_coeff, mo_occ) - assert np.max(np.abs(test_hessian_round1 - test_hessian_round2)) < 2e-7 + # np.abs() is a ufunc and dpnp sets __array_ufunc__ = None, so NumPy + # refuses a device operand (CuPy implements the protocol, so this used + # to work on CUDA only). Stay in the array's own namespace and pull a + # single scalar at the end, per gpu4pyscf#810. + assert abs(test_hessian_round1 - test_hessian_round2).max().item() < 2e-7 if __name__ == "__main__": print("Full Tests for DF Hessian") diff --git a/gpu4pyscf/df/tests/test_df_rks_grad.py b/gpu4pyscf/df/tests/test_df_rks_grad.py index c69a0b63a..9a4b26d2b 100644 --- a/gpu4pyscf/df/tests/test_df_rks_grad.py +++ b/gpu4pyscf/df/tests/test_df_rks_grad.py @@ -94,8 +94,8 @@ def _check_grad(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-6): grad_fd = np.array(grad_fd).reshape(-1,3) print('finite difference gradient:') print(grad_fd) - print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd)) - assert(cupy.linalg.norm(g_analy - grad_fd) < tol) + print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd)) + assert(np.linalg.norm(g_analy - grad_fd) < tol) def _vs_cpu(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-9): mf = rks.RKS(mol, xc=xc).density_fit(auxbasis=auxbasis0) diff --git a/gpu4pyscf/df/tests/test_df_uhf.py b/gpu4pyscf/df/tests/test_df_uhf.py index cdb3dd2c1..7032f31a7 100644 --- a/gpu4pyscf/df/tests/test_df_uhf.py +++ b/gpu4pyscf/df/tests/test_df_uhf.py @@ -81,8 +81,8 @@ def _check_grad(mol, tol=1e-5, disp=None): grad_fd = np.array(grad_fd).reshape(-1,3) print('finite difference gradient:') print(grad_fd) - print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd)) - assert(cupy.linalg.norm(g_analy - grad_fd) < tol) + print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd)) + assert(np.linalg.norm(g_analy - grad_fd) < tol) class KnownValues(unittest.TestCase): ''' diff --git a/gpu4pyscf/dft/libxc.py b/gpu4pyscf/dft/libxc.py index e0a8ee5ef..4682e204b 100644 --- a/gpu4pyscf/dft/libxc.py +++ b/gpu4pyscf/dft/libxc.py @@ -146,6 +146,10 @@ def _check_arrays(current_arrays, fields, sizes, factor, required): class _xcfun(ctypes.Structure): pass +# Set from the device library below when it advertises a limit; None means +# "no limit reported", i.e. the device path handles every order it advertises. +MAX_DERIV_ORDER = None + if _libxc is not None: _xc_func_p = ctypes.POINTER(xc_func_type) _libxc.xc_func_alloc.restype = _xc_func_p @@ -155,6 +159,16 @@ class _xcfun(ctypes.Structure): _libxc.xc_functional_get_name.argtypes = (ctypes.c_int, ) _libxc.xc_functional_get_name.restype = ctypes.c_char_p + # Highest derivative order the device library implements. The CUDA libxc + # build handles every order it advertises, so it does not export this + # symbol; the SYCL ExchCXX shim stops at fxc and does. Anything beyond it + # is routed to PySCF's CPU libxc by XCfun.on_gpu below. + try: + _libxc.xc_device_max_deriv_order.restype = ctypes.c_int + MAX_DERIV_ORDER = _libxc.xc_device_max_deriv_order() + except AttributeError: + MAX_DERIV_ORDER = None + nfunc = _libxc.xc_number_of_functionals() XC_IDS = np.zeros(nfunc, dtype=np.int32) _libxc.xc_available_functional_numbers(XC_IDS.ctypes) diff --git a/gpu4pyscf/dft/numint.py b/gpu4pyscf/dft/numint.py index e9c39aa34..6b768a768 100644 --- a/gpu4pyscf/dft/numint.py +++ b/gpu4pyscf/dft/numint.py @@ -1807,8 +1807,13 @@ def eval_xc_eff(ni, xc_code, rho, deriv=1, omega=None, xctype=None, out[m] = cupy.zeros([2,1]*m + [ngrids]) return out - # Fall back to the libxc library provided by PySCF, evaluate xc on CPUs - if not all(x.on_gpu for x, w in xcfuns): + # Fall back to the libxc library provided by PySCF, evaluate xc on CPUs. + # Either the functional has no device implementation, or the requested + # derivative order is beyond what the device library supports (the SYCL + # ExchCXX shim stops at fxc; CUDA libxc reports no limit). + max_deriv = libxc.MAX_DERIV_ORDER + if (not all(x.on_gpu for x, w in xcfuns) + or (max_deriv is not None and deriv > max_deriv)): ni_cpu = ni.to_cpu() ret = ni_cpu.eval_xc_eff(xc_code, rho.get(), deriv, xctype=xctype) for i in range(deriv+1): diff --git a/gpu4pyscf/dft/tests/test_numint.py b/gpu4pyscf/dft/tests/test_numint.py index de51ba2ef..f2a0f77f1 100644 --- a/gpu4pyscf/dft/tests/test_numint.py +++ b/gpu4pyscf/dft/tests/test_numint.py @@ -272,7 +272,8 @@ def test_sparse_index(self): i1 = min(i0+numint.MIN_BLK_SIZE, ngrids) ref = numint._sparse_index( opt._sorted_mol, grids.coords[i0:i1], opt.l_ctr_offsets, ao_loc, opt) - assert all(np.array_equal(r, x) for r, x in zip(ref[1:], dat[i][1:])) + assert all(r.shape == x.shape and bool((r == x).all()) + for r, x in zip(ref[1:], dat[i][1:])) def test_scale_ao(self): ao = cupy.random.rand(1, 3, 256) diff --git a/gpu4pyscf/geomopt/ase_solver.py b/gpu4pyscf/geomopt/ase_solver.py index 9ce693660..4a4d4be2d 100644 --- a/gpu4pyscf/geomopt/ase_solver.py +++ b/gpu4pyscf/geomopt/ase_solver.py @@ -20,6 +20,7 @@ from ase.optimize import BFGS from ase.filters import UnitCellFilter, StrainFilter +from ase.constraints import FixCom from pyscf import lib from pyscf.lib import logger from pyscf.pbc import gto @@ -48,6 +49,9 @@ def kernel(method, target=None, logfile=None, fmax=0.05, max_steps=100, Maximum number of optimization steps. restart : bool Whether to restart from a previous optimization state. + + For PBC calculations, the current ``cell.mesh`` is fixed when the + optimizer starts and reused throughout the optimization. ''' assert not restart if hasattr(method, 'cell'): @@ -59,7 +63,13 @@ def kernel(method, target=None, logfile=None, fmax=0.05, max_steps=100, is_pbc = isinstance(cell, gto.Cell) atoms = pyscf_to_ase_atoms(cell) - atoms.calc = PySCF(method=method) + calculator = PySCF(method=method) + atoms.calc = calculator + + if is_pbc and target in (None, 'atoms', 'cell'): + # This constraint cannot be used to lattice optimization + # because it fixes the atomic positions. + atoms.set_constraint(FixCom()) if target is None: if is_pbc: @@ -79,6 +89,8 @@ def kernel(method, target=None, logfile=None, fmax=0.05, max_steps=100, atoms = atoms.atoms if is_pbc: cell = cell.set_geom_(atoms.get_positions(), unit='Ang', a=atoms.cell, inplace=False) + method._geomopt_mesh = calculator.mesh + cell._geomopt_mesh = calculator.mesh else: cell = cell.set_geom_(atoms.get_positions(), unit='Ang', inplace=False) diff --git a/gpu4pyscf/geomopt/tests/test_pbc_geomopt_ase.py b/gpu4pyscf/geomopt/tests/test_pbc_geomopt_ase.py index 78a86eb72..15b42bfb4 100644 --- a/gpu4pyscf/geomopt/tests/test_pbc_geomopt_ase.py +++ b/gpu4pyscf/geomopt/tests/test_pbc_geomopt_ase.py @@ -23,10 +23,17 @@ from pyscf.data.nist import BOHR, HARTREE2EV if ase is not None: + from ase.constraints import FixCom from pyscf.pbc.tools.pyscf_ase import pyscf_to_ase_atoms class _FakeGradients: + def __init__(self, cell): + self.cell = cell + + def kernel(self): + return np.arange(1, self.cell.natm * 3 + 1).reshape(-1, 3) + def get_stress(self): return np.eye(3) @@ -35,10 +42,11 @@ class _FakeScanner: converged = True def __call__(self, cell): + self.cell = cell return 0. def Gradients(self): - return _FakeGradients() + return _FakeGradients(self.cell) class _FakeMethod(lib.StreamObject): @@ -65,6 +73,75 @@ def test_ase_stress_units(): calculator.results['stress'], np.eye(3) * HARTREE2EV / BOHR**3) +@pytest.mark.skipif(ase is None, reason='ASE not available') +def test_pbc_optimizer_freezes_automatic_mesh(): + from gpu4pyscf.geomopt import ase_solver + + cell = pyscf.M( + atom='He 0 0 0', a=np.eye(3) * 4., unit='Angstrom', + basis='gth-szv', pseudo='gth-pade', precision=1e-8, verbose=0) + mesh = np.asarray(cell.mesh).copy() + assert cell._mesh_from_build + + method = _FakeMethod(cell) + _, optimized_cell = ase_solver.kernel(method, max_steps=0) + + assert not cell._mesh_from_build + np.testing.assert_array_equal(method._geomopt_mesh, mesh) + np.testing.assert_array_equal(cell.mesh, mesh) + np.testing.assert_array_equal(optimized_cell.mesh, mesh) + + strained_cell = optimized_cell.set_geom_( + optimized_cell.atom_coords(), + a=optimized_cell.lattice_vectors() * 1.01, + unit='Bohr', + inplace=False, + ) + np.testing.assert_array_equal(strained_cell.mesh, mesh) + + +@pytest.mark.skipif(ase is None, reason='ASE not available') +def test_pbc_optimizer_fixcom(monkeypatch): + from gpu4pyscf.geomopt import ase_solver + + optimized = [] + + class FakeBFGS: + def __init__(self, atoms, logfile=None): + optimized.append(atoms) + + def run(self, fmax, steps): + return True + + monkeypatch.setattr(ase_solver, 'BFGS', FakeBFGS) + + for target in (None, 'atoms', 'cell', 'lattice'): + cell = pyscf.M( + atom='He 0 0 0; He 1 1 1', + a=np.eye(3) * 4., + unit='Angstrom', + basis='gth-szv', + pseudo='gth-pade', + mesh=[15] * 3, + verbose=0, + ) + ase_solver.kernel( + _FakeMethod(cell), + target=target, + max_steps=0, + ) + + system = optimized[-1] + atoms = getattr(system, 'atoms', system) + has_fixcom = any( + isinstance(constraint, FixCom) + for constraint in atoms.constraints + ) + assert has_fixcom == (target != 'lattice') + + if target == 'atoms': + assert np.allclose(atoms.get_forces().sum(axis=0), 0.) + @pytest.mark.skipif(ase is None, reason='ASE not available') def test_ase_optimize_cell(): cell = pyscf.M( diff --git a/gpu4pyscf/grad/tests/test_tdrhf_grad.py b/gpu4pyscf/grad/tests/test_tdrhf_grad.py index 3e3834ea9..7516b3bcb 100644 --- a/gpu4pyscf/grad/tests/test_tdrhf_grad.py +++ b/gpu4pyscf/grad/tests/test_tdrhf_grad.py @@ -22,6 +22,7 @@ from gpu4pyscf.lib.cupy_helper import contract from gpu4pyscf.grad import rhf as rhf_grad from gpu4pyscf.grad.tdrhf import _jk_energies_per_atom +from gpu4pyscf.nac.finite_diff import diagonalize, diagonalize_tda atom = """ O 0.0000000000 0.0000000000 0.0000000000 @@ -31,42 +32,6 @@ bas0 = "cc-pvdz" -def diagonalize(a, b, nroots=5): - nocc, nvir = a.shape[:2] - nov = nocc * nvir - a = a.reshape(nov, nov) - b = b.reshape(nov, nov) - h = np.block([[a, b], - [-b.conj(), -a.conj()]]) - e, xy = np.linalg.eig(np.asarray(h)) - assert np.max(np.abs(e.imag)) < 1e-14 - assert np.max(np.abs(xy.imag)) < 1e-14 - e = e.real - xy = xy.real - sorted_indices = np.argsort(e) - - e_sorted = e[sorted_indices] - xy_sorted = xy[:, sorted_indices] - - e_sorted_final = e_sorted[e_sorted > 1e-3] - xy_sorted = xy_sorted[:, e_sorted > 1e-3] - return e_sorted_final[:nroots], xy_sorted[:, :nroots] - - -def diagonalize_tda(a, nroots=5): - nocc, nvir = a.shape[:2] - nov = nocc * nvir - a = a.reshape(nov, nov) - e, xy = np.linalg.eigh(np.asarray(a)) - sorted_indices = np.argsort(e) - - e_sorted = e[sorted_indices] - xy_sorted = xy[:, sorted_indices] - - e_sorted_final = e_sorted[e_sorted > 1e-3] - xy_sorted = xy_sorted[:, e_sorted > 1e-3] - return e_sorted_final[:nroots], xy_sorted[:, :nroots] - def cal_analytic_gradient(mol, td, tdgrad, nocc, nvir, tda, singlet=True): a, b = td.get_ab() diff --git a/gpu4pyscf/grad/tests/test_tduhf_grad.py b/gpu4pyscf/grad/tests/test_tduhf_grad.py index 6b12ee831..002a3ada4 100644 --- a/gpu4pyscf/grad/tests/test_tduhf_grad.py +++ b/gpu4pyscf/grad/tests/test_tduhf_grad.py @@ -21,6 +21,8 @@ import gpu4pyscf from gpu4pyscf import scf as gpu_scf from gpu4pyscf.lib.multi_gpu import num_devices +from gpu4pyscf.nac.finite_diff import diagonalize_u as diagonalize +from gpu4pyscf.nac.finite_diff import diagonalize_tda_u as diagonalize_tda atom = """ O 0.0000000000 0.0000000000 0.0000000000 @@ -30,51 +32,6 @@ bas0 = "cc-pvdz" -def diagonalize(a, b, nroots=5): - a_aa, a_ab, a_bb = a - b_aa, b_ab, b_bb = b - nocc_a, nvir_a, nocc_b, nvir_b = a_ab.shape - a_aa = a_aa.reshape((nocc_a * nvir_a, nocc_a * nvir_a)) - a_ab = a_ab.reshape((nocc_a * nvir_a, nocc_b * nvir_b)) - a_bb = a_bb.reshape((nocc_b * nvir_b, nocc_b * nvir_b)) - b_aa = b_aa.reshape((nocc_a * nvir_a, nocc_a * nvir_a)) - b_ab = b_ab.reshape((nocc_a * nvir_a, nocc_b * nvir_b)) - b_bb = b_bb.reshape((nocc_b * nvir_b, nocc_b * nvir_b)) - a = np.block([[a_aa, a_ab], [a_ab.T, a_bb]]) - b = np.block([[b_aa, b_ab], [b_ab.T, b_bb]]) - abba = np.asarray(np.block([[a, b], [-b.conj(), -a.conj()]])) - e, xy = np.linalg.eig(abba) - assert np.max(np.abs(e.imag)) < 1e-14 - assert np.max(np.abs(xy.imag)) < 1e-14 - e = e.real - xy = xy.real - sorted_indices = np.argsort(e) - - e_sorted = e[sorted_indices] - xy_sorted = xy[:, sorted_indices] - - e_sorted_final = e_sorted[e_sorted > 1e-3] - xy_sorted = xy_sorted[:, e_sorted > 1e-3] - return e_sorted_final[:nroots], xy_sorted[:, :nroots] - - -def diagonalize_tda(a, nroots=5): - a_aa, a_ab, a_bb = a - nocc_a, nvir_a, nocc_b, nvir_b = a_ab.shape - a_aa = a_aa.reshape((nocc_a * nvir_a, nocc_a * nvir_a)) - a_ab = a_ab.reshape((nocc_a * nvir_a, nocc_b * nvir_b)) - a_bb = a_bb.reshape((nocc_b * nvir_b, nocc_b * nvir_b)) - a = np.block([[a_aa, a_ab], [a_ab.T, a_bb]]) - e, xy = np.linalg.eigh(a) - sorted_indices = np.argsort(e) - - e_sorted = e[sorted_indices] - xy_sorted = xy[:, sorted_indices] - - e_sorted_final = e_sorted[e_sorted > 1e-3] - xy_sorted = xy_sorted[:, e_sorted > 1e-3] - return e_sorted_final[:nroots], xy_sorted[:, :nroots] - def cal_analytic_gradient(mol, td, tdgrad, nocc_a, nvir_a, nocc_b, nvir_b, tda): a, b = td.get_ab() diff --git a/gpu4pyscf/hessian/tests/test_uks_hessian_grid_response.py b/gpu4pyscf/hessian/tests/test_uks_hessian_grid_response.py index 198e75b4e..fb7584d93 100644 --- a/gpu4pyscf/hessian/tests/test_uks_hessian_grid_response.py +++ b/gpu4pyscf/hessian/tests/test_uks_hessian_grid_response.py @@ -644,7 +644,8 @@ def test_hessian_grid_response_one_atom(self): mf = mol.UKS(xc = "wB97M-V").density_fit(auxbasis = "def2-universal-jkfit").to_gpu() mf.grids.atom_grid = (10,14) mf.nlcgrids.atom_grid = (10,14) - mf.conv_tol = 1e-12 + mf.conv_tol = 5e-12 + mf.max_cycle = 70 mf.kernel() assert mf.converged diff --git a/gpu4pyscf/lib/CMakeLists.txt b/gpu4pyscf/lib/CMakeLists.txt index 43aaaba91..f632b66d8 100644 --- a/gpu4pyscf/lib/CMakeLists.txt +++ b/gpu4pyscf/lib/CMakeLists.txt @@ -13,20 +13,125 @@ # limitations under the License. cmake_minimum_required (VERSION 3.19 FATAL_ERROR) # 3.19 is required by cutlass -project (gpu4pyscf C CXX CUDA Fortran) +# OFF by default so that a plain `python setup.py build` produces the CUDA +# backend, as upstream does. An ON default sent the CUDA CI down the SYCL +# path, where it failed looking for MKL. setup.py passes -DUSE_SYCL=ON when +# GPU4PYSCF_BACKEND=sycl is set. +option(USE_SYCL "Using SYCL backend" OFF) + +if (USE_SYCL) + project (gpu4pyscf C CXX Fortran) +else() + project (gpu4pyscf C CXX CUDA Fortran) +endif() set(CMAKE_C_STANDARD "99") set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) +if (USE_SYCL) + set(DPCTL_CMAKE_MODULES_PATH "${PROJECT_SOURCE_DIR}/cmake") + list(APPEND CMAKE_MODULE_PATH "${DPCTL_CMAKE_MODULES_PATH}") + + add_compile_definitions(USE_SYCL=1) + + # Production builds AOT-compile device code for Intel PVC (spir64_gen + + # -device pvc). That cannot run anywhere without a PVC, so CPU-only CI + # (GitHub-hosted runners driving the OpenCL CPU device) passes + # -DGPU4PYSCF_SYCL_AOT=OFF to get a generic JIT (spir64) build instead. + # Default ON keeps production PVC builds byte-identical. + option(GPU4PYSCF_SYCL_AOT "AOT-compile SYCL device code for Intel PVC" ON) + + # Honor a user-supplied CMAKE_BUILD_TYPE (e.g. passed through + # CMAKE_CONFIGURE_ARGS="-DCMAKE_BUILD_TYPE=RelWithDebInfo"). When nothing is + # supplied, default to Release. + # + # With AOT (-fsycl-targets=spir64_gen) a -g injected by RelWithDebInfo/Debug + # into CMAKE_CXX_FLAGS_ is forwarded into the device path and reaches + # `ocloc -options "-g"`, which compiles device code in debug mode (disables + # optimization). To keep the device backend optimized while still giving host + # frames file:line for gdb-oneapi, we downgrade the C++ -g to + # -gline-tables-only below. Host-side full -g is re-added further down for the + # host-only (C/Fortran) translation units, which never carry SYCL device code. + if (NOT CMAKE_BUILD_TYPE OR CMAKE_BUILD_TYPE STREQUAL "") + set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE) + endif() + message(STATUS "SYCL build: CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}") + + # Downgrade any -g the selected build type injects into the C++ per-config + # flags to host-only line tables. This keeps ocloc device compilation + # optimized (no -options "-g") while preserving host file:line info. C and + # Fortran host TUs keep full -g via the generator expressions below. + foreach(_cfg RELWITHDEBINFO DEBUG) + if (CMAKE_CXX_FLAGS_${_cfg}) + string(REPLACE "-g" "-gline-tables-only" + CMAKE_CXX_FLAGS_${_cfg} "${CMAKE_CXX_FLAGS_${_cfg}}") + endif() + endforeach() + + add_compile_options( + -fsycl + -Wno-ignored-reference-qualifiers + -Wno-register + -fsycl-device-code-split=per_kernel + -fp-model=precise + -fno-system-debug + -fsycl-default-sub-group-size=32 + ) + if (GPU4PYSCF_SYCL_AOT) + add_compile_options( + -fsycl-targets=spir64_gen + "SHELL:-Xsycl-target-backend=spir64_gen \"-device pvc\"" + ) + endif() + + # Host-side debug info WITHOUT propagating -g to the SYCL device backend. + # Restrict -g to C and Fortran translation units (pure host code). SYCL + # device-bearing C++ TUs are intentionally left without -g so `ocloc` does + # not receive -options "-g" and the device build stays optimized. + add_compile_options( + "$<$:-g>" + "$<$:-g>" + ) + + add_link_options( + -fsycl + -flink-huge-device-code + -fsycl-max-parallel-link-jobs=4 + ) + if (GPU4PYSCF_SYCL_AOT) + add_link_options( + -ftarget-register-alloc-mode=pvc:auto + -fsycl-targets=spir64_gen + "SHELL:-Xsycl-target-backend=spir64_gen \"-device pvc\"" + ) + # ExchCXX validates this against a GPU-only allowlist; an empty value + # disables its AOT pass entirely (generic spir64 JIT), which is what the + # CPU-only CI needs. + set(GPU4PYSCF_EXCHCXX_SYCL_TARGET intel_gpu_pvc) + else() + set(GPU4PYSCF_EXCHCXX_SYCL_TARGET "") + endif() + + # This needs to be built first than anything else + add_subdirectory(gsycl) + # Global fallback so subdirectories that haven't yet gained + # target_link_libraries(... sycl_compat) still find + # and via gsycl/ shims. + # Remove once all subdirectory CMakeLists.txt are updated. + include_directories(${PROJECT_SOURCE_DIR}/gsycl) +endif() + # For better performance on A100, the option # -DCUDA_ARCHITECTURES="70-real;80-real" can be added to cmake command -if(DEFINED CUDA_ARCHITECTURES) - set(CMAKE_CUDA_ARCHITECTURES "${CUDA_ARCHITECTURES}") -else() - set(CMAKE_CUDA_ARCHITECTURES "70-real;80-real;90-real") +if (NOT USE_SYCL) + if(DEFINED CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES "${CUDA_ARCHITECTURES}") + else() + set(CMAKE_CUDA_ARCHITECTURES "70-real;80-real;90-real") + endif() + message("CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}") endif() -message("CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}") if (NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE RELWITHDEBINFO) @@ -104,21 +209,77 @@ set(CXX_LINK_TEMPLATE " - -DCMAKE_INSTALL_LIBDIR:PATH=lib - -DCMAKE_C_CREATE_SHARED_LIBRARY=${C_LINK_TEMPLATE} - -DENABLE_XHOST:STRING=${BUILD_MARCH_NATIVE} - CMAKE_CACHE_ARGS -DCMAKE_CUDA_ARCHITECTURES:STRING=${CMAKE_CUDA_ARCHITECTURES} - ) -endif() + + if(USE_SYCL) + set(_exchcxx_lib "${PROJECT_SOURCE_DIR}/deps/lib/libexchcxx.so") + set(_exchcxx_inc "${PROJECT_SOURCE_DIR}/deps/include") + + # Always define the IMPORTED target — both branches need it + file(MAKE_DIRECTORY "${_exchcxx_inc}") # must exist at configure time + file(MAKE_DIRECTORY "${PROJECT_SOURCE_DIR}/deps/lib") + + add_library(ExchCXX::ExchCXX SHARED IMPORTED GLOBAL) + set_target_properties(ExchCXX::ExchCXX PROPERTIES + IMPORTED_LOCATION "${_exchcxx_lib}" + INTERFACE_INCLUDE_DIRECTORIES "${_exchcxx_inc}" + INTERFACE_COMPILE_DEFINITIONS "EXCHCXX_HAS_CONFIG_H=1" + ) + + # Conditionally build or skip ExchCXX + if(NOT EXISTS "${_exchcxx_lib}") + message(STATUS "ExchCXX: not found -- building from source (~35 min)") + ExternalProject_Add(ExchCXX + GIT_REPOSITORY https://github.com/abagusetty/ExchCXX.git + GIT_TAG d7ed4af53ce0499da62d38e54e100713ee565e8f + UPDATE_DISCONNECTED ON + PREFIX ${PROJECT_BINARY_DIR}/deps + INSTALL_DIR ${PROJECT_SOURCE_DIR}/deps + CMAKE_ARGS + -DCMAKE_BUILD_TYPE=Release + -DBUILD_SHARED_LIBS=ON + # ExternalProject does not inherit the parent's compiler choice; without + # this the inner configure falls back to /usr/bin/c++ (no -fsycl). + -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER} + -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER} + -DEXCHCXX_ENABLE_SYCL=ON + -DEXCHCXX_SYCL_TARGET=${GPU4PYSCF_EXCHCXX_SYCL_TARGET} + -DEXCHCXX_ENABLE_TESTS=OFF + -DEXCHCXX_ENABLE_LIBXC=OFF + -DCMAKE_POSITION_INDEPENDENT_CODE=ON + -DCMAKE_INSTALL_PREFIX:PATH= + -DCMAKE_INSTALL_LIBDIR:PATH=lib + -DCMAKE_C_CREATE_SHARED_LIBRARY=${C_LINK_TEMPLATE} + -DCMAKE_CXX_CREATE_SHARED_LIBRARY=${CXX_LINK_TEMPLATE} + BUILD_BYPRODUCTS "${_exchcxx_lib}" + ) + else() + message(STATUS "ExchCXX: found ${_exchcxx_lib} -- skipping rebuild") + add_custom_target(ExchCXX) # no-op target, satisfies add_dependencies() + endif() + + # Wire IMPORTED target to whichever ExchCXX target was created above + add_dependencies(ExchCXX::ExchCXX ExchCXX) + else (USE_SYCL) + ExternalProject_Add(libxc + GIT_REPOSITORY https://github.com/wxj6000/libxc.git + GIT_TAG b225c254c063e1de835a4425115c9a6377478b32 + PATCH_COMMAND git apply --reject ${PROJECT_SOURCE_DIR}/libxc_prune.patch || true + PREFIX ${PROJECT_BINARY_DIR}/deps + INSTALL_DIR ${PROJECT_SOURCE_DIR}/deps + CMAKE_ARGS -DBUILD_SHARED_LIBS=ON -DENABLE_CUDA=ON + -DCMAKE_POLICY_VERSION_MINIMUM=3.5 + -DENABLE_FORTRAN=OFF -DDISABLE_KXC=OFF -DDISABLE_LXC=ON -DDISABLE_FHC=ON + -DCMAKE_INSTALL_PREFIX:PATH= + -DCMAKE_INSTALL_LIBDIR:PATH=lib + -DCMAKE_C_CREATE_SHARED_LIBRARY=${C_LINK_TEMPLATE} + -DENABLE_XHOST:STRING=${BUILD_MARCH_NATIVE} + CMAKE_CACHE_ARGS -DCMAKE_CUDA_ARCHITECTURES:STRING=${CMAKE_CUDA_ARCHITECTURES} + ) + endif(USE_SYCL) +endif(BUILD_LIBXC) diff --git a/gpu4pyscf/lib/cupy_helper/CMakeLists.txt b/gpu4pyscf/lib/cupy_helper/CMakeLists.txt index 1d7969fcf..dccb9338b 100644 --- a/gpu4pyscf/lib/cupy_helper/CMakeLists.txt +++ b/gpu4pyscf/lib/cupy_helper/CMakeLists.txt @@ -35,9 +35,18 @@ endif() add_library(cupy_helper SHARED ${cupy_helper_src}) if(BUILD_CUTLASS) -add_dependencies(cupy_helper cutlass) -target_include_directories(cupy_helper PRIVATE ${cutlass_SOURCE_DIR}/tools/util/include ${cutlass_SOURCE_DIR}/include) + add_dependencies(cupy_helper cutlass) + target_include_directories(cupy_helper PRIVATE ${cutlass_SOURCE_DIR}/tools/util/include ${cutlass_SOURCE_DIR}/include) endif() -set_target_properties(cupy_helper PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) -set_target_properties(cupy_helper PROPERTIES CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") +set_target_properties(cupy_helper PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + +if (USE_SYCL) + set_source_files_properties(${cupy_helper_src} PROPERTIES LANGUAGE CXX) + target_compile_options(cupy_helper PRIVATE -x c++ -nocudainc -nocudalib -DCUTLASS_ENABLE_SYCL=1) + target_link_libraries(cupy_helper PRIVATE sycl_compat gsycl) +else() + set_target_properties(cupy_helper PROPERTIES + CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") +endif() diff --git a/gpu4pyscf/lib/cupy_helper/add_sparse.cu b/gpu4pyscf/lib/cupy_helper/add_sparse.cu index 154ba852a..9bd316e1b 100644 --- a/gpu4pyscf/lib/cupy_helper/add_sparse.cu +++ b/gpu4pyscf/lib/cupy_helper/add_sparse.cu @@ -15,6 +15,7 @@ */ #include +#include "gsycl/gpu_compat.h" #define THREADS 32 #define BLOCK_DIM 32 @@ -22,8 +23,9 @@ __global__ void _add_sparse(double *a, double *b, int *indices, int n, int m, int count) { - int row = blockIdx.x * BLOCK_DIM + threadIdx.x; - int col = blockIdx.y * BLOCK_DIM + threadIdx.y; + setup_context(); + int row = blockIdx_x * BLOCK_DIM + threadIdx_x; + int col = blockIdx_y * BLOCK_DIM + threadIdx_y; if (row >= m || col >= m){ return; } @@ -38,9 +40,10 @@ extern "C" { __host__ int add_sparse(cudaStream_t stream, double *a, double *b, int *indices, int n, int m, int count){ int ntile = (m + THREADS - 1) / THREADS; - dim3 threads(THREADS, THREADS); - dim3 blocks(ntile, ntile); - _add_sparse<<>>(a, b, indices, n, m, count); + auto threads = make_block(THREADS, THREADS); + auto blocks = make_grid(ntile, ntile); + LAUNCH_KERNEL(_add_sparse, blocks, threads, 0, stream, + a, b, indices, n, m, count); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { return 1; diff --git a/gpu4pyscf/lib/cupy_helper/async_d2h_2d.cu b/gpu4pyscf/lib/cupy_helper/async_d2h_2d.cu index ea62eabc0..6e059eecf 100644 --- a/gpu4pyscf/lib/cupy_helper/async_d2h_2d.cu +++ b/gpu4pyscf/lib/cupy_helper/async_d2h_2d.cu @@ -29,6 +29,10 @@ int async_d2h_2d(cudaStream_t stream, double *dst, int dstride, const double *sr int width = rows * sizeof(double); int height = cols * sizeof(double); +#ifdef USE_SYCL + stream.ext_oneapi_memcpy2d(host_ptr, dpitch, device_ptr, spitch, + width, height); +#else // USE_SYCL cudaError_t err = cudaMemcpy2DAsync(host_ptr, dpitch, device_ptr, spitch, width, height, cudaMemcpyDeviceToHost); /* @@ -43,6 +47,7 @@ int async_d2h_2d(cudaStream_t stream, double *dst, int dstride, const double *sr return 1; } +#endif //USE_SYCL return 0; } } diff --git a/gpu4pyscf/lib/cupy_helper/block_diag.cu b/gpu4pyscf/lib/cupy_helper/block_diag.cu index eb53a7459..cf022bac1 100644 --- a/gpu4pyscf/lib/cupy_helper/block_diag.cu +++ b/gpu4pyscf/lib/cupy_helper/block_diag.cu @@ -21,8 +21,18 @@ __global__ static void _block_diag(double *out, int m, int n, double *diags, int ndiags, int *offsets, int *rows, int *cols) { - int r = blockIdx.x; - +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int blockIdx_x = item.get_group(1); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); +#else + int blockIdx_x = blockIdx.x; + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; +#endif + int r = blockIdx_x; + if (r >= ndiags){ return; } @@ -32,8 +42,8 @@ static void _block_diag(double *out, int m, int n, double *diags, int ndiags, in int row_offset = rows[r]; int col_offset = cols[r]; - for (int i = threadIdx.y; i < m0; i += THREADS){ - for (int j = threadIdx.x; j < n0; j += THREADS){ + for (int i = threadIdx_y; i < m0; i += THREADS){ + for (int j = threadIdx_x; j < n0; j += THREADS){ out[(i+row_offset)*n + (j+col_offset)] = diags[diag_offset + i*n0 + j]; } } @@ -42,6 +52,13 @@ static void _block_diag(double *out, int m, int n, double *diags, int ndiags, in extern "C" { int block_diag(cudaStream_t stream, double *out, int m, int n, double *diags, int ndiags, int *offsets, int *rows, int *cols) { +#ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(1, ndiags); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _block_diag(out, m, n, diags, ndiags, offsets, rows, cols); + }); +#else //USE_SYCL dim3 threads(THREADS, THREADS); dim3 blocks(ndiags); _block_diag<<>>(out, m, n, diags, ndiags, offsets, rows, cols); @@ -50,6 +67,7 @@ int block_diag(cudaStream_t stream, double *out, int m, int n, double *diags, in fprintf(stderr, "CUDA Error in block_diag: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/cupy_helper/cart2sph.cu b/gpu4pyscf/lib/cupy_helper/cart2sph.cu index ab5f79d0b..8cc7e5ec7 100644 --- a/gpu4pyscf/lib/cupy_helper/cart2sph.cu +++ b/gpu4pyscf/lib/cupy_helper/cart2sph.cu @@ -22,7 +22,12 @@ // (n,ncart,stride) -> (n,nsph,stride), count = n*stride __global__ static void _cart2sph_ang2(double *cart, double *sph, int stride, int count){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= count){ return; } @@ -46,7 +51,12 @@ static void _cart2sph_ang2(double *cart, double *sph, int stride, int count){ __global__ static void _cart2sph_ang3(double *cart, double *sph, int stride, int count){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= count){ return; } @@ -76,7 +86,12 @@ static void _cart2sph_ang3(double *cart, double *sph, int stride, int count){ __global__ static void _cart2sph_ang4(double *cart, double *sph, int stride, int count){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= count){ return; } @@ -113,7 +128,12 @@ static void _cart2sph_ang4(double *cart, double *sph, int stride, int count){ __global__ static void _cart2sph_ang5(double *cart, double *sph, int stride, int count){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= count){ return; } @@ -157,7 +177,12 @@ static void _cart2sph_ang5(double *cart, double *sph, int stride, int count){ __global__ static void _cart2sph_ang6(double *cart, double *sph, int stride, int count){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= count){ return; } @@ -210,7 +235,12 @@ static void _cart2sph_ang6(double *cart, double *sph, int stride, int count){ __global__ static void _cart2sph_ang7(double *cart, double *sph, int stride, int count){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= count){ return; } @@ -275,6 +305,23 @@ extern "C" { __host__ int cart2sph(cudaStream_t stream, double *cart_gto, double *sph_gto, int stride, int count, int ang) { +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks((count + THREADS - 1)/THREADS); + switch (ang) { + case 0: break; + case 1: break; + case 2: stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { _cart2sph_ang2 (cart_gto, sph_gto, stride, count); }); break; + case 3: stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { _cart2sph_ang3 (cart_gto, sph_gto, stride, count); }); break; + case 4: stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { _cart2sph_ang4 (cart_gto, sph_gto, stride, count); }); break; + case 5: stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { _cart2sph_ang5 (cart_gto, sph_gto, stride, count); }); break; + case 6: stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { _cart2sph_ang6 (cart_gto, sph_gto, stride, count); }); break; + case 7: stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { _cart2sph_ang7 (cart_gto, sph_gto, stride, count); }); break; + default: + fprintf(stderr, "Ang > 7 is not supported!\n"); + return 1; + } +#else // USE_SYCL dim3 threads(THREADS); dim3 blocks((count + THREADS - 1)/THREADS); switch (ang) { @@ -295,6 +342,7 @@ int cart2sph(cudaStream_t stream, double *cart_gto, double *sph_gto, int stride, if (err != cudaSuccess) { return 1; } +#endif // USE_SYCL return 0; } } diff --git a/gpu4pyscf/lib/cupy_helper/dist_matrix.cu b/gpu4pyscf/lib/cupy_helper/dist_matrix.cu index bbc8c5b19..6cc79ec6c 100644 --- a/gpu4pyscf/lib/cupy_helper/dist_matrix.cu +++ b/gpu4pyscf/lib/cupy_helper/dist_matrix.cu @@ -21,8 +21,14 @@ __global__ static void _calc_distances(double *dist, const double *x, const double *y, int m, int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int i = item.get_global_id(1); + int j = item.get_global_id(0); +#else int i = blockIdx.x * blockDim.x + threadIdx.x; int j = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= m || j >= n){ return; } @@ -38,6 +44,13 @@ int dist_matrix(cudaStream_t stream, double *dist, const double *x, const double { int ntilex = (m + THREADS - 1) / THREADS; int ntiley = (n + THREADS - 1) / THREADS; +#ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(ntiley, ntilex); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _calc_distances(dist, x, y, m, n); + }); +#else //USE_SYCL dim3 threads(THREADS, THREADS); dim3 blocks(ntilex, ntiley); _calc_distances<<>>(dist, x, y, m, n); @@ -45,6 +58,7 @@ int dist_matrix(cudaStream_t stream, double *dist, const double *x, const double if (err != cudaSuccess) { return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/cupy_helper/grouped_gemm.cu b/gpu4pyscf/lib/cupy_helper/grouped_gemm.cu index 7e753bf4b..a83ddd7a9 100644 --- a/gpu4pyscf/lib/cupy_helper/grouped_gemm.cu +++ b/gpu4pyscf/lib/cupy_helper/grouped_gemm.cu @@ -14,7 +14,6 @@ * limitations under the License. */ - #include #include #include @@ -45,6 +44,26 @@ static int get_device_compute_capability() { return properties.major * 10 + properties.minor; } +#ifdef USE_SYCL +// PVC 1550 +using cutlass_tensorop_d884gemm_grouped_64x128_16x3_tt_align1_base = + typename cutlass::gemm::kernel::DefaultGemmGrouped< + double, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 1, + double, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1, + double, cutlass::layout::RowMajor, + double, + cutlass::arch::OpClassTensorOp, + cutlass::arch::IntelXe, + cutlass::gemm::GemmShape<64, 128, 16>, + cutlass::gemm::GemmShape<32, 64, 16>, + cutlass::gemm::GemmShape<8, 8, 4>, + cutlass::epilogue::thread::LinearCombination, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, + 3, + cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly, + cutlass::arch::OpMultiplyAdd +>::GemmKernel; +#else // A100 using cutlass_tensorop_d884gemm_grouped_64x128_16x3_tt_align1_base = typename cutlass::gemm::kernel::DefaultGemmGrouped< @@ -83,6 +102,7 @@ using cutlass_simt_dgemm_grouped_64x128_8x2_tt_align1_base = cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly, cutlass::arch::OpMultiplyAdd >::GemmKernel; +#endif template cutlass::Status grouped_gemm_kernel_run(int problem_count, cutlass::gemm::GemmCoord* problem_sizes, diff --git a/gpu4pyscf/lib/cupy_helper/sparse_cderi.cu b/gpu4pyscf/lib/cupy_helper/sparse_cderi.cu index 0d7e2cf85..09cc711b0 100644 --- a/gpu4pyscf/lib/cupy_helper/sparse_cderi.cu +++ b/gpu4pyscf/lib/cupy_helper/sparse_cderi.cu @@ -35,8 +35,14 @@ typedef struct { __global__ void _unpack(CDERI_BLOCK block, int nao, int offset, double *out){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int ij = item.get_global_id(1); + int k = item.get_global_id(0); +#else int ij = blockIdx.x * blockDim.x + threadIdx.x; int k = blockIdx.y * blockDim.y + threadIdx.y; +#endif int nij = block.nij; int idx_aux = k + offset; @@ -95,6 +101,14 @@ int unpack_block(CDERI_BLOCK *block, int p1, int p2, int nao, double *buf){ int nij = block->nij; int blockx = (nij + THREADS - 1) / THREADS; int blocky = (p2 - p1 + THREADS - 1) / THREADS; +#ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(blocky, blockx); + CDERI_BLOCK dev_block = *block; + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _unpack(dev_block, nao, p1, buf); + }); +#else //USE_SYCL dim3 threads(THREADS, THREADS); dim3 blocks(blockx, blocky); @@ -104,6 +118,7 @@ int unpack_block(CDERI_BLOCK *block, int p1, int p2, int nao, double *buf){ if (err != cudaSuccess) { return 1; } +#endif return 0; } diff --git a/gpu4pyscf/lib/cupy_helper/take_last2d.cu b/gpu4pyscf/lib/cupy_helper/take_last2d.cu index 6cb77fc7e..4579273f2 100644 --- a/gpu4pyscf/lib/cupy_helper/take_last2d.cu +++ b/gpu4pyscf/lib/cupy_helper/take_last2d.cu @@ -22,9 +22,16 @@ __global__ static void _take_last2d(double *a, const double *b, int *indices, int na, int nb) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<3>(); + size_t i = item.get_group(0); + int j = item.get_global_id(2); + int k = item.get_global_id(1); +#else size_t i = blockIdx.z; int j = blockIdx.x * blockDim.x + threadIdx.x; int k = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (j >= na || k >= na) { return; } @@ -40,8 +47,14 @@ __global__ static void _takebak(double *out, double *a, int *indices, int count, int n_o, int n_a) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int i0 = item.get_group(0) * COUNT_BLOCK; + int j = item.get_global_id(1); +#else int i0 = blockIdx.y * COUNT_BLOCK; int j = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (j >= n_a) { return; } @@ -63,6 +76,13 @@ int take_last2d(cudaStream_t stream, double *a, const double *b, int *indices, { // reorder j and k in a[i,j,k] with indicies int ntile = (na + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<3> threads(1, THREADS, THREADS); + sycl::range<3> blocks(blk_size, ntile, ntile); + stream.parallel_for(sycl::nd_range<3>(blocks * threads, threads), [=](auto item) { + _take_last2d(a, b, indices, na, nb); + }); + #else dim3 threads(THREADS, THREADS); dim3 blocks(ntile, ntile, blk_size); _take_last2d<<>>(a, b, indices, na, nb); @@ -70,6 +90,7 @@ int take_last2d(cudaStream_t stream, double *a, const double *b, int *indices, if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -77,14 +98,23 @@ int takebak(cudaStream_t stream, double *out, double *a_h, int *indices, int count, int n_o, int n_a) { double *a_d; + int ntile = (n_a + THREADS*THREADS - 1) / (THREADS*THREADS); + int ncount = (count + COUNT_BLOCK - 1) / COUNT_BLOCK; + + #ifdef USE_SYCL + *(void **)&a_d = (double *)a_h; + sycl::range<2> threads(1, THREADS*THREADS); + sycl::range<2> blocks(ncount, ntile); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _takebak(out, a_d, indices, count, n_o, n_a); + }); + #else cudaError_t err; err = cudaHostGetDevicePointer(&a_d, a_h, 0); // zero-copy check if (err != cudaSuccess) { return 1; } - int ntile = (n_a + THREADS*THREADS - 1) / (THREADS*THREADS); - int ncount = (count + COUNT_BLOCK - 1) / COUNT_BLOCK; dim3 threads(THREADS*THREADS); dim3 blocks(ntile, ncount); _takebak<<>>(out, a_d, indices, count, n_o, n_a); @@ -92,6 +122,7 @@ int takebak(cudaStream_t stream, double *out, double *a_h, int *indices, if (err != cudaSuccess) { return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/cupy_helper/transpose.cu b/gpu4pyscf/lib/cupy_helper/transpose.cu index 7f0918fef..eaf290f67 100644 --- a/gpu4pyscf/lib/cupy_helper/transpose.cu +++ b/gpu4pyscf/lib/cupy_helper/transpose.cu @@ -22,17 +22,31 @@ static __global__ void _transpose_dsum(double *a, int n, int counts, int hermi) { - if(blockIdx.x > blockIdx.y){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + using tile_t = double[THREADS][THREADS]; + tile_t& block = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + __shared__ double block[THREADS][THREADS]; +#endif + if(blockIdx_x > blockIdx_y){ return; } - __shared__ double block[THREADS][THREADS]; - int blockx_off = blockIdx.x * BLOCK_DIM; - int blocky_off = blockIdx.y * BLOCK_DIM; - size_t x0 = blockx_off + threadIdx.x; - size_t y0 = blocky_off + threadIdx.y; - size_t x1 = blocky_off + threadIdx.x; - size_t y1 = blockx_off + threadIdx.y; + int blockx_off = blockIdx_x * BLOCK_DIM; + int blocky_off = blockIdx_y * BLOCK_DIM; + size_t x0 = blockx_off + threadIdx_x; + size_t y0 = blocky_off + threadIdx_y; + size_t x1 = blocky_off + threadIdx_x; + size_t y1 = blockx_off + threadIdx_y; size_t nn = n * n; size_t xy0 = y0 * n + x0; size_t xy1 = y1 * n + x1; @@ -40,26 +54,26 @@ void _transpose_dsum(double *a, int n, int counts, int hermi) for (int k = 0; k < counts; ++k) { double *pa = a + nn * k; if (x0 < n && y0 < n){ - block[threadIdx.y][threadIdx.x] = pa[xy0]; + block[threadIdx_y][threadIdx_x] = pa[xy0]; } __syncthreads(); if (x1 < n && y1 < n){ if (hermi == 1) { - block[threadIdx.x][threadIdx.y] += pa[xy1]; + block[threadIdx_x][threadIdx_y] += pa[xy1]; } else { - block[threadIdx.x][threadIdx.y] -= pa[xy1]; + block[threadIdx_x][threadIdx_y] -= pa[xy1]; } } __syncthreads(); if(x0 < n && y0 < n){ - pa[xy0] = block[threadIdx.y][threadIdx.x]; + pa[xy0] = block[threadIdx_y][threadIdx_x]; } if(x1 < n && y1 < n){ if (hermi == 1) { - pa[xy1] = block[threadIdx.x][threadIdx.y]; + pa[xy1] = block[threadIdx_x][threadIdx_y]; } else { - pa[xy1] = -block[threadIdx.x][threadIdx.y]; + pa[xy1] = -block[threadIdx_x][threadIdx_y]; } } __syncthreads(); @@ -69,18 +83,33 @@ void _transpose_dsum(double *a, int n, int counts, int hermi) static __global__ void _transpose_zsum(double *a, int n, int counts, int hermi) { - if(blockIdx.x > blockIdx.y){ - return; - } +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + using tile_t = double[THREADS][THREADS]; + tile_t& blockR = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + tile_t& blockI = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; __shared__ double blockR[THREADS][THREADS]; __shared__ double blockI[THREADS][THREADS]; +#endif + if(blockIdx_x > blockIdx_y){ + return; + } - int blockx_off = blockIdx.x * BLOCK_DIM; - int blocky_off = blockIdx.y * BLOCK_DIM; - size_t x0 = blockx_off + threadIdx.x; - size_t y0 = blocky_off + threadIdx.y; - size_t x1 = blocky_off + threadIdx.x; - size_t y1 = blockx_off + threadIdx.y; + int blockx_off = blockIdx_x * BLOCK_DIM; + int blocky_off = blockIdx_y * BLOCK_DIM; + size_t x0 = blockx_off + threadIdx_x; + size_t y0 = blocky_off + threadIdx_y; + size_t x1 = blocky_off + threadIdx_x; + size_t y1 = blockx_off + threadIdx_y; size_t nn = n * n * 2; size_t xy0 = (y0 * n + x0) * 2; size_t xy1 = (y1 * n + x1) * 2; @@ -88,32 +117,32 @@ void _transpose_zsum(double *a, int n, int counts, int hermi) for (int k = 0; k < counts; ++k) { double *pa = a + nn * k; if (x0 < n && y0 < n){ - blockR[threadIdx.y][threadIdx.x] = pa[xy0 ]; - blockI[threadIdx.y][threadIdx.x] = pa[xy0+1]; + blockR[threadIdx_y][threadIdx_x] = pa[xy0 ]; + blockI[threadIdx_y][threadIdx_x] = pa[xy0+1]; } __syncthreads(); if (x1 < n && y1 < n){ if (hermi == 1) { - blockR[threadIdx.x][threadIdx.y] += pa[xy1 ]; - blockI[threadIdx.x][threadIdx.y] -= pa[xy1+1]; + blockR[threadIdx_x][threadIdx_y] += pa[xy1 ]; + blockI[threadIdx_x][threadIdx_y] -= pa[xy1+1]; } else { - blockR[threadIdx.x][threadIdx.y] -= pa[xy1 ]; - blockI[threadIdx.x][threadIdx.y] += pa[xy1+1]; + blockR[threadIdx_x][threadIdx_y] -= pa[xy1 ]; + blockI[threadIdx_x][threadIdx_y] += pa[xy1+1]; } } __syncthreads(); if(x0 < n && y0 < n){ - pa[xy0 ] = blockR[threadIdx.y][threadIdx.x]; - pa[xy0+1] = blockI[threadIdx.y][threadIdx.x]; + pa[xy0 ] = blockR[threadIdx_y][threadIdx_x]; + pa[xy0+1] = blockI[threadIdx_y][threadIdx_x]; } if(x1 < n && y1 < n){ if (hermi == 1) { - pa[xy1 ] = blockR[threadIdx.x][threadIdx.y]; - pa[xy1+1] = -blockI[threadIdx.x][threadIdx.y]; + pa[xy1 ] = blockR[threadIdx_x][threadIdx_y]; + pa[xy1+1] = -blockI[threadIdx_x][threadIdx_y]; } else { - pa[xy1 ] = -blockR[threadIdx.x][threadIdx.y]; - pa[xy1+1] = blockI[threadIdx.x][threadIdx.y]; + pa[xy1 ] = -blockR[threadIdx_x][threadIdx_y]; + pa[xy1+1] = blockI[threadIdx_x][threadIdx_y]; } } __syncthreads(); @@ -124,9 +153,17 @@ extern "C" { int transpose_dsum(cudaStream_t stream, double *a, int n, int counts, int hermi) { int ntile = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(ntile, ntile); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _transpose_dsum(a, n, counts, hermi); + }); + #else dim3 threads(THREADS, THREADS); dim3 blocks(ntile, ntile); _transpose_dsum<<>>(a, n, counts, hermi); + #endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { return 1; @@ -137,9 +174,17 @@ int transpose_dsum(cudaStream_t stream, double *a, int n, int counts, int hermi) int transpose_zsum(cudaStream_t stream, double *a, int n, int counts, int hermi) { int ntile = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(ntile, ntile); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _transpose_zsum(a, n, counts, hermi); + }); + #else dim3 threads(THREADS, THREADS); dim3 blocks(ntile, ntile); _transpose_zsum<<>>(a, n, counts, hermi); + #endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { return 1; diff --git a/gpu4pyscf/lib/cupy_helper/unpack.cu b/gpu4pyscf/lib/cupy_helper/unpack.cu index f00b35eb4..565302402 100644 --- a/gpu4pyscf/lib/cupy_helper/unpack.cu +++ b/gpu4pyscf/lib/cupy_helper/unpack.cu @@ -30,8 +30,14 @@ __global__ static void _pack_tril(double *a_tril, double *a, size_t n, int counts) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + size_t j = item.get_global_id(1); + size_t i = item.get_global_id(0); +#else size_t j = blockIdx.x * blockDim.x + threadIdx.x; size_t i = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= n || j >= n || i < j) { return; @@ -47,8 +53,14 @@ void _pack_tril(double *a_tril, double *a, size_t n, int counts) __global__ static void _unpack_tril(double *eri_tril, double *eri, size_t nao, int counts) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + size_t j = item.get_global_id(1); + size_t i = item.get_global_id(0); +#else size_t j = blockIdx.x * blockDim.x + threadIdx.x; size_t i = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= nao || j >= nao || i < j) { return; } @@ -63,8 +75,14 @@ void _unpack_tril(double *eri_tril, double *eri, size_t nao, int counts) __global__ static void _dfill_triu(double *eri, size_t nao, int counts, int hermi) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int j = item.get_global_id(1); + int i = item.get_global_id(0); +#else int j = blockIdx.x * blockDim.x + threadIdx.x; int i = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= nao || j >= nao || i >= j) { return; } @@ -82,8 +100,14 @@ void _dfill_triu(double *eri, size_t nao, int counts, int hermi) __global__ static void _zfill_triu(double *eri, size_t nao, int counts, int hermi) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int j = item.get_global_id(1); + int i = item.get_global_id(0); +#else int j = blockIdx.x * blockDim.x + threadIdx.x; int i = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= nao || j >= nao || i >= j) { return; } @@ -107,9 +131,16 @@ void decompress_kernel(double *out, size_t out_stride, double *cderi, int *pair_idx, int npairs, int nao, size_t naux, int aux0, int aux1) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int threads = item.get_local_range(0); + int batch_id = item.get_group(0); +#else int thread_id = threadIdx.x; int threads = blockDim.x; int batch_id = blockIdx.x; +#endif int dcol = aux1 - aux0; int pair0 = batch_id * RBLKSIZE; int pair1 = min(pair0 + RBLKSIZE, npairs); @@ -135,9 +166,16 @@ void d_t_kernel(double *out, size_t out_stride, double *cderi, int *pair_idx, int npairs, int nao, int aux0, int aux1, int fill_triu) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int bx = item.get_group(1); + int by = item.get_group(0); + int thread_id = item.get_local_id(1); +#else int bx = blockIdx.x; int by = blockIdx.y; int thread_id = threadIdx.x; +#endif int threads = STRIDE * CBLKSIZE; int tx = thread_id % CBLKSIZE; int ty = thread_id / CBLKSIZE; @@ -147,7 +185,12 @@ void d_t_kernel(double *out, size_t out_stride, size_t Npairs = npairs; size_t Nao = nao; +#ifdef USE_SYCL + using buf_t = double[RBLKSIZE][CBLKSIZE+1]; + buf_t& buf = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ double buf[RBLKSIZE][CBLKSIZE+1]; +#endif if (pair_start+tx < npairs) { for (int k = ty; k < min(RBLKSIZE, daux-aux_start); k += STRIDE) { buf[k][tx] = cderi[(aux_start+k)*Npairs+pair_start+tx]; @@ -176,9 +219,16 @@ void z_d_t_kernel(double2 *out, size_t out_stride, double2 *cderi, int *pair_idx, int npairs, int nao, int aux0, int aux1) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int bx = item.get_group(1); + int by = item.get_group(0); + int thread_id = item.get_local_id(1); +#else int bx = blockIdx.x; int by = blockIdx.y; int thread_id = threadIdx.x; +#endif int threads = STRIDE * CBLKSIZE; int tx = thread_id % CBLKSIZE; int ty = thread_id / CBLKSIZE; @@ -188,7 +238,12 @@ void z_d_t_kernel(double2 *out, size_t out_stride, size_t Npairs = npairs; size_t Nao = nao; +#ifdef USE_SYCL + using zbuf_t = double2[RBLKSIZE][CBLKSIZE+1]; + zbuf_t& buf = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ double2 buf[RBLKSIZE][CBLKSIZE+1]; +#endif if (pair_start+tx < npairs) { for (int k = ty; k < min(RBLKSIZE, daux-aux_start); k += STRIDE) { buf[k][tx] = cderi[(aux_start+k)*Npairs+pair_start+tx]; @@ -213,6 +268,17 @@ extern "C" { int fill_triu(cudaStream_t stream, double *a, int n, int counts, int hermi, int dtype) { +#ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + int nx = (n + threads[1] - 1) / threads[1]; + int ny = (n + threads[0] - 1) / threads[0]; + sycl::range<2> blocks(ny, nx); + if (dtype == 1) { // float64 + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { _dfill_triu(a, n, counts, hermi); }); + } else { + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { _zfill_triu(a, n, counts, hermi); }); + } +#else dim3 threads(THREADS, THREADS); int nx = (n + threads.x - 1) / threads.x; int ny = (n + threads.y - 1) / threads.y; @@ -227,11 +293,19 @@ int fill_triu(cudaStream_t stream, double *a, int n, int counts, int hermi, fprintf(stderr, "fill_tril error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int pack_tril(cudaStream_t stream, double *a_tril, double *a, int n, int counts) { +#ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + int nx = (n + threads[1] - 1) / threads[1]; + int ny = (n + threads[0] - 1) / threads[0]; + sycl::range<2> blocks(ny, nx); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { _pack_tril(a_tril, a, n, counts); }); +#else dim3 threads(THREADS, THREADS); int nx = (n + threads.x - 1) / threads.x; int ny = (n + threads.y - 1) / threads.y; @@ -242,12 +316,21 @@ int pack_tril(cudaStream_t stream, double *a_tril, double *a, int n, int counts) fprintf(stderr, "pack_tril error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int unpack_tril(cudaStream_t stream, double *eri_tril, double *eri, int nao, int counts, int hermi) { +#ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + int nx = (nao + threads[1] - 1) / threads[1]; + int ny = (nao + threads[0] - 1) / threads[0]; + sycl::range<2> blocks(ny, nx); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { _unpack_tril(eri_tril, eri, nao, counts); }); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { _dfill_triu(eri, nao, counts, hermi); }); +#else dim3 threads(THREADS, THREADS); int nx = (nao + threads.x - 1) / threads.x; int ny = (nao + threads.y - 1) / threads.y; @@ -258,6 +341,7 @@ int unpack_tril(cudaStream_t stream, double *eri_tril, double *eri, if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -265,6 +349,13 @@ int decompress_and_fill(cudaStream_t stream, double *out, int out_stride, double *cderi, int *pair_idx, int npairs, int nao, int naux, int aux0, int aux1) { +#ifdef USE_SYCL + sycl::range<1> threads(512); + sycl::range<1> blocks((npairs+RBLKSIZE-1)/RBLKSIZE); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + decompress_kernel(out, out_stride, cderi, pair_idx, npairs, nao, naux, aux0, aux1); + }); +#else dim3 blocks((npairs+RBLKSIZE-1)/RBLKSIZE); decompress_kernel<<>>( out, out_stride, cderi, pair_idx, npairs, nao, naux, aux0, aux1); @@ -273,6 +364,7 @@ int decompress_and_fill(cudaStream_t stream, double *out, int out_stride, fprintf(stderr, "decompress_and_fill error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -280,6 +372,15 @@ int decompress_and_transpose(cudaStream_t stream, double *out, int out_stride, double *cderi, int *pair_idx, int npairs, int nao, int aux0, int aux1, int fill_triu, int on_host) { +#ifdef USE_SYCL + // Host USM allocations are directly device-accessible; no address mapping. + double *eri_gpu = cderi; + sycl::range<2> threads(1, CBLKSIZE * STRIDE); + sycl::range<2> blocks((aux1-aux0+RBLKSIZE-1)/RBLKSIZE, (npairs+CBLKSIZE-1)/CBLKSIZE); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + d_t_kernel(out, out_stride, eri_gpu, pair_idx, npairs, nao, aux0, aux1, fill_triu); + }); +#else double *eri_gpu = cderi; if (on_host) { cudaError_t err = cudaHostGetDevicePointer(&eri_gpu, cderi, 0); @@ -297,6 +398,7 @@ int decompress_and_transpose(cudaStream_t stream, double *out, int out_stride, fprintf(stderr, "decompress_and_transpose error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -304,6 +406,15 @@ int z_decompress_and_transpose(cudaStream_t stream, double2 *out, int out_stride double2 *cderi, int *pair_idx, int npairs, int nao, int aux0, int aux1, int fill_triu, int on_host) { +#ifdef USE_SYCL + // Host USM allocations are directly device-accessible; no address mapping. + double2 *eri_gpu = cderi; + sycl::range<2> threads(1, CBLKSIZE * STRIDE); + sycl::range<2> blocks((aux1-aux0+RBLKSIZE-1)/RBLKSIZE, (npairs+CBLKSIZE-1)/CBLKSIZE); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + z_d_t_kernel(out, out_stride, eri_gpu, pair_idx, npairs, nao, aux0, aux1); + }); +#else double2 *eri_gpu = cderi; if (on_host) { cudaError_t err = cudaHostGetDevicePointer(&eri_gpu, cderi, 0); @@ -321,6 +432,7 @@ int z_decompress_and_transpose(cudaStream_t stream, double2 *out, int out_stride fprintf(stderr, "decompress_and_transpose error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/dpnp_helper.py b/gpu4pyscf/lib/dpnp_helper.py new file mode 100644 index 000000000..ecc88a61e --- /dev/null +++ b/gpu4pyscf/lib/dpnp_helper.py @@ -0,0 +1,1591 @@ +# Copyright 2021-2024 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import functools +import inspect +import ctypes +import warnings +import numpy as np +import scipy.linalg +import cupy +import dpnp +import dpctl +import dpctl.memory as dpmem +from dpnp.dpnp_array import dpnp_array # low-level constructor +from pyscf import lib +from gpu4pyscf.lib import logger +from gpu4pyscf.lib.cutensor import contract +from gpu4pyscf.lib.onemkl_lapack import eigh as onemkl_eigh +#from gpu4pyscf.lib.onemkl_lapack import eigh as onemkl_eigh, cholesky as onemkl_cholesky +#from gpu4pyscf.lib.onemkl_lapack import cholesky #NOQA +from gpu4pyscf.lib.memcpy import copy_array, p2p_transfer #NOQA +from gpu4pyscf.lib import multi_gpu +from gpu4pyscf.lib.utils import load_library +from gpu4pyscf.lib.multi_gpu import lru_cache +from gpu4pyscf.__config__ import num_devices, _p2p_access + +LMAX_ON_GPU = 7 +DSOLVE_LINDEP = 1e-13 +MAX_EIGH_DIM = 23150 + +# Fall back to scipy.linalg.eigh for arrays larger than the onemkl/dpnp eigh +# limit (see MAX_EIGH_DIM). Referenced by cond() and eigh(). +SCIPY_EIGH_FOR_LARGE_ARRAYS = True + +# Threshold (in bytes) above which allocations bypass any pooled allocator. +# Kept for API parity with cupy_helper; SYCL/USM manages memory automatically. +MEMPOOL_THRESHOLD = 100000000 + +_kernel_registery = {} + +libdpnp_helper = load_library('libcupy_helper') +# Upstream modules import the library handle under its cupy_helper name. +libcupy_helper = libdpnp_helper + +def pin_memory(array): + mem = dpctl.memory.MemoryUSMHost(array.nbytes) + ret = np.frombuffer(mem, array.dtype, array.size).reshape(array.shape) + ret[...] = array + return ret + +def release_gpu_stack(): + pass + +def print_mem_info(): + total_mem = cupy.cuda.get_total_memory() + free_mem = cupy.cuda.get_free_memory() + used_mem = total_mem - free_mem + GB = 1024 * 1024 * 1024 + msg = (f'mem_avail: {free_mem/GB:.3f} GB, total_mem: {total_mem/GB:.3f} GB, ' + f'used_mem: {used_mem/GB:.3f} GB') + print(msg) + return msg + +def get_avail_mem(exclude_memory_pool=False): + return cupy.cuda.get_free_memory() + +def concatenate(array_list): + ''' Concatenate axis=0 only + ''' + if _p2p_access: + return dpnp.concatenate(array_list) + else: + #array_list_cpu = [a.get() for a in array_list] + n = sum([a.shape[0] for a in array_list]) + a0_shape = list(array_list[0].shape) + out_shape = tuple([n] + a0_shape[1:]) + out = dpnp.empty(out_shape) + p0 = p1 = 0 + for a in array_list: + p1 = p0 + a.shape[0] + #out[p0:p1].set(a) + copy_array(a, out[p0:p1]) + p0 = p1 + return out + +def broadcast_to_devices(): + ''' Broadcast dpnp ndarray to all the devices, return a list of dpnp ndarray + ''' + raise NotImplementedError + +def reduce_to_device(array_list, inplace=False): + return multi_gpu.array_reduce(array_list, inplace) + + # ''' Reduce a list of ndarray in different devices to device 0 + # TODO: reduce memory footprint, improve throughput + # ''' + # assert len(array_list) == num_devices + # if num_devices == 1: + # return array_list[0] + + # out_shape = array_list[0].shape + # for s in _streams: + # s.synchronize() + + # if inplace: + # result = array_list[0] + # else: + # result = array_list[0].copy() + + # # Transfer data chunk by chunk, reduce memory footprint, + # result = result.reshape(-1) + # for device_id, matrix in enumerate(array_list): + # if device_id == 0: + # continue + + # assert matrix.device.id == device_id + # matrix = matrix.reshape(-1) + # blksize = 1024*1024*1024 // matrix.itemsize # 1GB + # for p0, p1 in lib.prange(0,len(matrix), blksize): + # result[p0:p1] += copy_array(matrix[p0:p1]) + # #result[p0:p1] += cupy.asarray(matrix[p0:p1]) + # return result.reshape(out_shape) + +def device2host_2d(a_cpu, a_gpu, stream=None): + if stream is None: + stream = cupy.cuda.get_current_stream() + libdpnp_helper.async_d2h_2d( + ctypes.cast(stream.ptr, ctypes.c_void_p), + a_cpu.ctypes.data_as(ctypes.c_void_p), + ctypes.c_int(a_cpu.strides[0]), + ctypes.cast(a_gpu.data.ptr, ctypes.c_void_p), + ctypes.c_int(a_gpu.strides[0]), + ctypes.c_int(a_gpu.shape[0]), + ctypes.c_int(a_gpu.shape[1])) + +# Define dpnp array with tag using Python class wrapper +class CPArrayWithTag(dpnp.ndarray): + """A dpnp array that can carry pyscf's ad-hoc tags. + + dpnp's array methods build plain ``dpnp_array`` results instead of + honouring the subclass the way NumPy/CuPy do, so a tagged array silently + degrades to an untaggable ``dpnp_array`` on indexing, ``.T``, ``reshape`` + and friends. Call sites written against CuPy rely on those results still + being subclass instances so they can assign attributes onto them -- + ``df_jk._transpose_dm`` does ``dm_T.factor_l = ...`` and + ``nac.tdrhf_grad_nacv._dms_to_list`` does ``dm.factor_l = ...`` while + iterating. Re-view the derived array to restore that. + + Like CuPy, the tags themselves are *not* propagated -- only the ability to + hold them. Copying ``__dict__`` here would be wrong: a slice of a stacked + density matrix must not inherit the parent's ``factor_l``. + """ + + @staticmethod + def _retag(res): + if isinstance(res, dpnp.ndarray) and not isinstance(res, CPArrayWithTag): + return res.view(CPArrayWithTag) + return res + + def __getitem__(self, key): + return CPArrayWithTag._retag(dpnp.ndarray.__getitem__(self, key)) + + @property + def T(self): + return CPArrayWithTag._retag(dpnp.ndarray.T.fget(self)) + + def reshape(self, *args, **kwargs): + return CPArrayWithTag._retag( + dpnp.ndarray.reshape(self, *args, **kwargs)) + + def transpose(self, *args, **kwargs): + return CPArrayWithTag._retag( + dpnp.ndarray.transpose(self, *args, **kwargs)) + +#@functools.wraps(lib.tag_array) +def tag_array(a, **kwargs): + ''' + a should be cupy/numpy array or tuple of cupy/numpy array + + attach attributes to cupy ndarray for cupy array + attach attributes to numpy ndarray for numpy array + ''' + if isinstance(a, dpnp.ndarray) or isinstance(a[0], dpnp.ndarray): + t = dpnp.asarray(a).view(CPArrayWithTag) + if isinstance(a, CPArrayWithTag): + t.__dict__.update(a.__dict__) + else: + t = np.asarray(a).view(lib.NPArrayWithTag) + if isinstance(a, lib.NPArrayWithTag): + t.__dict__.update(a.__dict__) + t.__dict__.update(kwargs) + return t + +def asarray(a, **kwargs): + ''' + Similar to `dpnp.asarray`, but optimized for transferring NumPy arrays from host to device. + If the input object is an instance of `CPArrayWithTag`, this function will remove any + associated attributes from the tagged array during the transfer. + + Unlike `dpnp.asarray`, which may allocate a temporary buffer during array transfer, + this function eliminates that buffer for efficiency. + ''' + if isinstance(a, np.ndarray): + # Avoid temporary buffer allocation during host-to-device transfer. + # In DPNP/SYCL, we use usm_data.copy_from_host for a direct DMA transfer. + + allow_fast_transfer = kwargs.get('dtype', a.dtype) == a.dtype + # a must be C-contiguous or F-contiguous + if not a.flags.c_contiguous and not a.flags.f_contiguous: + allow_fast_transfer = False + + if allow_fast_transfer: + if a.size == 0: + # Empty array — just create an empty device array directly + order = 'F' if a.flags.f_contiguous and not a.flags.c_contiguous else 'C' + return dpnp.empty(a.shape, dtype=a.dtype, order=order) + # Preserve memory layout (C or F order) + order = 'F' if a.flags.f_contiguous and not a.flags.c_contiguous else 'C' + out = dpnp.empty(a.shape, dtype=a.dtype, order=order) + # Direct host-to-device copy via USM memory, no intermediate pinned buffer + out.get_array().usm_data.copy_from_host(a.ravel(order=order).view(np.uint8)) + if kwargs.get('blocking', False): + dpnp.get_sycl_queue().wait() # SYCL sync, not CUDA + return out + + elif isinstance(a, CPArrayWithTag): + a = a.view(dpnp.ndarray) + + return dpnp.asarray(a, **kwargs) + +ensure_numpy = dpnp.asnumpy + +def to_dpnp(a): + '''Convert numpy → dpnp (handles nested structures)''' + if isinstance(a, lib.NPArrayWithTag): + attrs = {k: to_dpnp(v) for k, v in a.__dict__.items()} + return tag_array(dpnp.asarray(a), **attrs) + if isinstance(a, np.ndarray): + return dpnp.asarray(a) + if isinstance(a, (tuple, list)): + return type(a)(to_dpnp(x) for x in a) + if isinstance(a, dict): + return {k: to_dpnp(v) for k, v in a.items()} + return a + # '''Converts a numpy (and subclass) object to a dpnp object''' + # if isinstance(a, lib.NPArrayWithTag): + # attrs = {k: to_dpnp(v) for k, v in a.__dict__.items()} + # return tag_array(cupy.asarray(a), **attrs) + # if isinstance(a, np.ndarray): + # return cupy.asarray(a) + # return a + +######################################################################## +# This section guards the return_cupy_array() section when a pyscf.cpu method +# is passed with DPNP arrays. It explicitly copies the array to numpy.ndarrat type +# Similar to cupy but it does implictly (hiding the transfer) + +def _to_numpy(a): + '''Convert GPU → NumPy (handles nested structures)''' + if isinstance(a, dpnp.ndarray): + return dpnp.asnumpy(a) + if hasattr(a, 'asnumpy'): + return a.asnumpy() + if isinstance(a, (tuple, list)): + return type(a)(_to_numpy(x) for x in a) + if isinstance(a, dict): + return {k: _to_numpy(v) for k, v in a.items()} + return a + +def _is_cpu_function(fn): + '''Detect if function is from CPU PySCF (pyscf.scf.*)''' + fn_module = fn.__module__ or '' + if 'pyscf' in fn_module and 'gpu4pyscf' not in fn_module: + return True + if 'cpu' in (fn.__name__ or '').lower(): + return True + return False + +class _GPUMethodProxy: + """ + Proxy that wraps an mf object so that any method call + automatically converts numpy inputs back to dpnp. + """ + def __init__(self, mf): + object.__setattr__(self, '_mf', mf) + + def __getattr__(self, name): + attr = getattr(object.__getattribute__(self, '_mf'), name) + # Only wrap genuine bound methods / functions. Callable *objects* + # (e.g. gto.Mole, which defines __call__) must be returned as-is, + # otherwise `self.mol` becomes a wrapper function and attribute + # access like `mol.inertia_moment()` fails. + if inspect.isroutine(attr): + @functools.wraps(attr) + def wrapper(*args, **kwargs): + # Convert numpy arrays back to dpnp before calling GPU method + args = tuple(to_dpnp(a) for a in args) + kwargs = {k: to_dpnp(v) for k, v in kwargs.items()} + return attr(*args, **kwargs) + return wrapper + return attr + + def __setattr__(self, name, value): + setattr(object.__getattribute__(self, '_mf'), name, value) + +def return_cupy_array(fn): + '''Convert inputs for CPU functions, wrapping mf to auto-convert on callbacks''' + is_cpu = _is_cpu_function(fn) + + @functools.wraps(fn) + def filter_ret(*args, **kwargs): + if is_cpu and args: + # Wrap mf (first arg) so GPU method calls auto-convert numpy→dpnp + mf_proxy = _GPUMethodProxy(args[0]) + args = (mf_proxy,) + tuple(_to_numpy(a) for a in args[1:]) + kwargs = {k: _to_numpy(v) for k, v in kwargs.items()} + + ret = fn(*args, **kwargs) + + if isinstance(ret, tuple): + return tuple(to_dpnp(x) for x in ret) + return to_dpnp(ret) + return filter_ret + +# ## How this works: +# ``` +# Decorator: +# 1. Detects CPU function (pyscf.scf.uhf.get_occ) +# 2. Wraps mf with _GPUMethodProxy +# 3. Converts mo_energy, mo_coeff to numpy +# 4. Calls CPU function + +# CPU get_occ: +# - Works with numpy arrays ✓ +# - numpy.zeros_like(mo_energy) works ✓ +# - Calls mf.spin_square(numpy_arrays) +# ↓ +# _GPUMethodProxy intercepts: +# - Converts numpy → dpnp +# - Calls real GPU spin_square(dpnp_arrays) ✓ + +######################################################################## + +def pack_tril(a, stream=None): + ndim = a.ndim + assert ndim in (2, 3) + if ndim == 2: + a = a[None] + + counts, n = a.shape[:2] + if a.dtype != np.float64 or not a.flags.c_contiguous: + idx = dpnp.arange(n) + mask = idx[:,None] >= idx + a_tril = a[:,mask] + else: + if stream is None: + stream = cupy.cuda.get_current_stream() + a_tril = dpnp.empty((counts, n*(n+1)//2), dtype=np.float64) + err = libdpnp_helper.pack_tril( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(a_tril.data.ptr, ctypes.c_void_p), + ctypes.cast(a.data.ptr, ctypes.c_void_p), + ctypes.c_int(n), ctypes.c_int(counts)) + if err != 0: + raise RuntimeError('pack_tril kernel failed') + + if ndim == 2: + a_tril = a_tril[0] + return a_tril + +def unpack_tril(cderi_tril, out=None, stream=None, hermi=1): + assert cderi_tril.flags.c_contiguous + assert hermi in (1, 2) + ndim = cderi_tril.ndim + assert ndim in (1, 2) + if ndim == 1: + cderi_tril = cderi_tril[None] + count = cderi_tril.shape[0] + nao = int((2*cderi_tril.shape[1])**.5) + out = ndarray((count,nao,nao), dtype=cderi_tril.dtype, buffer=out) + + if cderi_tril.dtype != np.float64: + idx = dpnp.arange(nao) + mask = idx[:,None] >= idx + cderiT = out.transpose(0,2,1) + if hermi == 1: + cderiT[:,mask] = cderi_tril.conj() + else: + raise NotImplementedError + out [:,mask] = cderi_tril + return out + + if stream is None: + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.unpack_tril( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(cderi_tril.data.ptr, ctypes.c_void_p), + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.c_int(nao), + ctypes.c_int(count), + ctypes.c_int(hermi)) + if err != 0: + raise RuntimeError('failed in unpack_tril kernel') + if ndim == 1: + out = out[0] + return out + +def unpack_sparse(cderi_sparse, row, col, p0, p1, nao, out=None, stream=None): + warnings.warn( + 'unpack_sparse is deprecated. Use fill_symmetric instead', + DeprecationWarning, stacklevel=2) + return fill_symmetric(cderi_sparse, row*nao+col, nao, p0, p1, out, stream) + +def fill_symmetric(a, pair_addresses, nao, p0=0, p1=None, out=None, stream=None): + ''' + Performs: + + i, j = divmod(pair_addresses, nao) + out[j,i,:] = out[i,j,:] = a[:,p0:p1] + ''' + assert a.ndim == 2 + if p1 is None: + p1 = a.shape[1] + + if out is None: + out = dpnp.zeros([nao,nao,p1-p0], dtype=a.dtype) + else: + assert out.ndim == 3 + + if stream is None: + stream = cupy.cuda.get_current_stream() + + pair_addresses = dpnp.asarray(pair_addresses, dtype=np.int32) + out_stride = out.strides[-2] // out.itemsize + if a.strides[-1] == 8: # a is in row major + a_stride = a.strides[-2] // a.itemsize + err = libdpnp_helper.decompress_and_fill( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.c_int(out_stride), + ctypes.cast(a.data.ptr, ctypes.c_void_p), + ctypes.cast(pair_addresses.data.ptr, ctypes.c_void_p), + ctypes.c_int(len(pair_addresses)), + ctypes.c_int(nao), + ctypes.c_int(a_stride), + ctypes.c_int(p0), ctypes.c_int(p1)) + if err != 0: + raise RuntimeError('decompress_and_fill failed') + else: # a is in column major + err = libdpnp_helper.decompress_and_transpose( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.c_int(out_stride), ctypes.cast(a.data.ptr, ctypes.c_void_p), + ctypes.cast(pair_addresses.data.ptr, ctypes.c_void_p), + ctypes.c_int(len(pair_addresses)), + ctypes.c_int(nao), + ctypes.c_int(p0), ctypes.c_int(p1), + ctypes.c_int(1), ctypes.c_int(0)) + if err != 0: + raise RuntimeError('decompress_and_transpose failed') + return out + +def add_sparse(a, b, indices): + ''' + a[:,...,:np.ix_(indices, indices)] += b + ''' + assert a.device == b.device + assert a.flags.c_contiguous + assert b.flags.c_contiguous + if len(indices) == 0: return a + indices = dpnp.asarray(indices, dtype=np.int32) + n = a.shape[-1] + m = b.shape[-1] + if a.ndim > 2: + count = np.prod(a.shape[:-2]) + elif a.ndim == 2: + count = 1 + else: + raise RuntimeError('add_sparse only supports 2d or 3d tensor') + + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.add_sparse( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(a.data.ptr, ctypes.c_void_p), + ctypes.cast(b.data.ptr, ctypes.c_void_p), + ctypes.cast(indices.data.ptr, ctypes.c_void_p), + ctypes.c_int(n), + ctypes.c_int(m), + ctypes.c_int(count) + ) + if err != 0: + raise RuntimeError('failed in sparse_add2d') + return a + +def dist_matrix(x, y, out=None): + '''np.linalg.norm(x[:,None,:] - y[None,:,:], axis=2)''' + x = dpnp.asarray(x, dtype=np.float64) + y = dpnp.asarray(y, dtype=np.float64) + assert x.flags.c_contiguous + assert y.flags.c_contiguous + + m = x.shape[0] + n = y.shape[0] + if out is None: + out = dpnp.empty([m,n]) + + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.dist_matrix( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.cast(x.data.ptr, ctypes.c_void_p), + ctypes.cast(y.data.ptr, ctypes.c_void_p), + ctypes.c_int(m), + ctypes.c_int(n) + ) + if err != 0: + raise RuntimeError('failed in calculating distance matrix') + return out + +@multi_gpu.lru_cache(1) +def _initialize_c2s_data(): + from gpu4pyscf.gto import mole + c2s_l = [mole.cart2sph_by_l(l) for l in range(LMAX_ON_GPU)] + c2s_data = dpnp.concatenate([x.ravel() for x in c2s_l]) + c2s_offset = np.cumsum([0] + [x.shape[0]*x.shape[1] for x in c2s_l]) + return c2s_l, c2s_data, c2s_offset + +def block_c2s_diag(angular, counts): + ''' + Diagonal blocked cartesian to spherical transformation + Args: + angular (list): angular momentum type, e.g. [0,1,2,3] + counts (list): count of each angular momentum + ''' + c2s_l, c2s_data, c2s_offset = _initialize_c2s_data() + + nshells = np.sum(counts) + rows = [np.array([0], dtype='int32')] + cols = [np.array([0], dtype='int32')] + offsets = [] + for l, count in zip(angular, counts): + r, c = c2s_l[l].shape + rows.append(rows[-1][-1] + np.arange(1,count+1, dtype='int32') * r) + cols.append(cols[-1][-1] + np.arange(1,count+1, dtype='int32') * c) + offsets += [c2s_offset[l]] * count + rows = dpnp.asarray(np.hstack(rows)) + cols = dpnp.asarray(np.hstack(cols)) + + ncart, nsph = int(rows[-1]), int(cols[-1]) + cart2sph = dpnp.zeros([ncart, nsph]) + offsets = dpnp.asarray(offsets, dtype='int32') + + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.block_diag( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(cart2sph.data.ptr, ctypes.c_void_p), + ctypes.c_int(ncart), + ctypes.c_int(nsph), + ctypes.cast(c2s_data.data.ptr, ctypes.c_void_p), + ctypes.c_int(nshells), + ctypes.cast(offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(rows.data.ptr, ctypes.c_void_p), + ctypes.cast(cols.data.ptr, ctypes.c_void_p), + ) + if err != 0: + raise RuntimeError('failed in block_diag kernel') + return cart2sph + +def block_diag(blocks, out=None): + ''' + each block size is up to 16x16 + ''' + rows = np.cumsum(np.asarray([0] + [x.shape[0] for x in blocks])) + cols = np.cumsum(np.asarray([0] + [x.shape[1] for x in blocks])) + offsets = np.cumsum(np.asarray([0] + [x.shape[0]*x.shape[1] for x in blocks])) + + m, n = rows[-1], cols[-1] + if out is None: out = dpnp.zeros([m, n]) + rows = dpnp.asarray(rows, dtype='int32') + cols = dpnp.asarray(cols, dtype='int32') + offsets = dpnp.asarray(offsets, dtype='int32') + data = dpnp.concatenate([x.ravel() for x in blocks]) + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.block_diag( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.c_int(m), + ctypes.c_int(n), + ctypes.cast(data.data.ptr, ctypes.c_void_p), + ctypes.c_int(len(blocks)), + ctypes.cast(offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(rows.data.ptr, ctypes.c_void_p), + ctypes.cast(cols.data.ptr, ctypes.c_void_p), + ) + if err != 0: + raise RuntimeError('failed in block_diag kernel') + return out + +def take_last2d(a, indices, out=None): + ''' + Reorder the last 2 dimensions as a[..., indices[:,None], indices] + ''' + assert a.flags.c_contiguous + assert a.shape[-1] == a.shape[-2] + nao = a.shape[-1] + nidx = len(indices) + if a.ndim == 2: + count = 1 + else: + count = np.prod(a.shape[:-2]) + out = ndarray((count, nidx, nidx), buffer=out) + indices_int32 = dpnp.asarray(indices, dtype='int32') + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.take_last2d( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.cast(a.data.ptr, ctypes.c_void_p), + ctypes.cast(indices_int32.data.ptr, ctypes.c_void_p), + ctypes.c_int(count), + ctypes.c_int(nidx), + ctypes.c_int(nao) + ) + if err != 0: + raise RuntimeError('failed in take_last2d kernel') + if a.ndim == 2: + out = out.reshape(nidx,nidx) + return out + +def takebak(out, a, indices, axis=-1): + '''(experimental) + Take elements from a NumPy array along an axis and write to CuPy array. + out[..., indices] = a + ''' + assert axis == -1 + assert isinstance(a, np.ndarray) + assert isinstance(out, dpnp.ndarray) + assert out.ndim == a.ndim + assert a.shape[-1] == len(indices) + if a.ndim == 1: + count = 1 + else: + assert out.shape[:-1] == a.shape[:-1] + count = np.prod(a.shape[:-1]) + n_a = a.shape[-1] + n_o = out.shape[-1] + indices_int32 = dpnp.asarray(indices, dtype=dpnp.int32) + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.takebak( + ctypes.c_void_p(stream.ptr), + ctypes.c_void_p(out.data.ptr), a.ctypes, + ctypes.c_void_p(indices_int32.data.ptr), + ctypes.c_int(count), ctypes.c_int(n_o), ctypes.c_int(n_a) + ) + if err != 0: # Not the mapped host memory + out[...,indices] = dpnp.asarray(a) + return out + +def transpose_sum(a, stream=None, inplace=True, hermi=1): + ''' + perform + a + a.transpose(0,2,1) for hermi=1 or + a - a.transpose(0,2,1) hermi=2 + inplace + ''' + if not inplace: + a = dpnp.copy(a, order='C') + ndim = a.ndim + assert hermi == 1 or hermi == 2 + assert isinstance(a, dpnp.ndarray) + assert a.flags.c_contiguous + assert ndim == 2 or ndim == 3 + if ndim == 2: + a = a[None] + count, m, n = a.shape + assert m == n + out = a + stream = cupy.cuda.get_current_stream() + if a.dtype == np.float64: + fn = libdpnp_helper.transpose_dsum + else: + fn = libdpnp_helper.transpose_zsum + err = fn(ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(a.data.ptr, ctypes.c_void_p), + ctypes.c_int(n), ctypes.c_int(count), ctypes.c_int(hermi)) + if err != 0: + raise RuntimeError('failed in transpose_sum kernel') + if ndim == 2: + out = out[0] + return out + +def hermi_triu(mat, hermi=1, inplace=True, stream=None): + ''' + Use the elements of the lower triangular part to fill the upper triangular part. + See also pyscf.lib.hermi_triu + + hermi=1 performs symmetric; hermi=2 performs anti-symmetric + ''' + assert hermi in (1, 2) + if inplace: + assert mat.flags.c_contiguous + else: + mat = mat.copy('C') + + if mat.ndim == 2: + n = mat.shape[0] + counts = 1 + elif mat.ndim == 3: + counts, n = mat.shape[:2] + else: + raise ValueError(f'dimension not supported {mat.ndim}') + + if mat.dtype == np.float64: + dtype = 1 + elif mat.dtype == np.complex128: + dtype = 2 + else: + raise ValueError(f'{mat.ndim} type not supported') + + if stream is None: + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.fill_triu( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(mat.data.ptr, ctypes.c_void_p), + ctypes.c_int(n), ctypes.c_int(counts), ctypes.c_int(hermi), + ctypes.c_int(dtype)) + if err != 0: + raise RuntimeError('hermi_triu kernel failed') + return mat + +def cart2sph_cutensor(t, axis=0, ang=1, out=None): + ''' + transform 'axis' of a tensor from cartesian basis into spherical basis with cutensor + ''' + from gpu4pyscf.gto import mole + if(ang <= 1): + if(out is not None): out[:] = t + return t + size = list(t.shape) + c2s = mole.cart2sph_by_l(ang) + if(not t.flags['C_CONTIGUOUS']): t = dpnp.asarray(t, order='C') + li_size = c2s.shape + nli = size[axis] // li_size[0] + i0 = max(1, np.prod(size[:axis])) + i3 = max(1, np.prod(size[axis+1:])) + out_shape = size[:axis] + [nli*li_size[1]] + size[axis+1:] + + t_cart = t.reshape([i0*nli, li_size[0], i3]) + if(out is not None): + out = out.reshape([i0*nli, li_size[1], i3]) + t_sph = contract('min,ip->mpn', t_cart, c2s, out=out) + return t_sph.reshape(out_shape) + +def cart2sph(t, axis=0, ang=1, out=None, stream=None): + ''' + transform 'axis' of a tensor from cartesian basis into spherical basis + ''' + from gpu4pyscf.gto import mole + if(ang <= 1): + if(out is not None): out[:] = t + return t + size = list(t.shape) + c2s = mole.cart2sph_by_l(ang) + if(not t.flags['C_CONTIGUOUS']): t = dpnp.asarray(t, order='C') + li_size = c2s.shape + nli = size[axis] // li_size[0] + i0 = max(1, np.prod(size[:axis])) + i3 = max(1, np.prod(size[axis+1:])) + out_shape = size[:axis] + [nli*li_size[1]] + size[axis+1:] + + t_cart = t.reshape([i0*nli, li_size[0], i3]) + if(out is not None): + out = out.reshape([i0*nli, li_size[1], i3]) + else: + out = dpnp.empty(out_shape) + count = i0*nli*i3 + if stream is None: + stream = cupy.cuda.get_current_stream() + err = libdpnp_helper.cart2sph( + ctypes.cast(stream.ptr, ctypes.c_void_p), + ctypes.cast(t_cart.data.ptr, ctypes.c_void_p), + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.c_int(i3), + ctypes.c_int(count), + ctypes.c_int(ang) + ) + if err != 0: + raise RuntimeError('failed in cart2sph kernel') + return out.reshape(out_shape) + +# a copy with modification from +# https://github.com/pyscf/pyscf/blob/9219058ac0a1bcdd8058166cad0fb9127b82e9bf/pyscf/lib/linalg_helper.py#L1536 +def krylov(aop, b, x0=None, tol=1e-10, max_cycle=30, dot=dpnp.dot, + lindep=DSOLVE_LINDEP, callback=None, hermi=False, + verbose=logger.WARN): + r'''Krylov subspace method to solve (1+a) x = b. Ref: + J. A. Pople et al, Int. J. Quantum. Chem. Symp. 13, 225 (1979). + Args: + aop : function(x) => array_like_x + aop(x) to mimic the matrix vector multiplication :math:`\sum_{j}a_{ij} x_j`. + The argument is a 1D array. The returned value is a 1D array. + b : a vector or a list of vectors + Kwargs: + x0 : 1D array + Initial guess + tol : float + Tolerance to terminate the operation aop(x). + max_cycle : int + max number of iterations. + lindep : float + Linear dependency threshold. The function is terminated when the + smallest eigenvalue of the metric of the trial vectors is lower + than this threshold. + dot : function(x, y) => scalar + Inner product + callback : function(envs_dict) => None + callback function takes one dict as the argument which is + generated by the builtin function :func:`locals`, so that the + callback function can access all local variables in the current + environment. + Returns: + x : ndarray like b + ''' + if isinstance(aop, dpnp.ndarray) and aop.ndim == 2: + return dpnp.linalg.solve(aop+dpnp.eye(aop.shape[0]), b) + + if isinstance(verbose, logger.Logger): + log = verbose + else: + log = logger.Logger(sys.stdout, verbose) + + if not (isinstance(b, dpnp.ndarray) and b.ndim == 1): + b = dpnp.asarray(b) + + if x0 is None: + x1 = b + else: + b = b - (x0 + aop(x0)) + x1 = b + if x1.ndim == 1: + x1 = x1.reshape(1, x1.size) + nroots, ndim = x1.shape + x1, rmat = _stable_qr(x1, dpnp.dot, lindep=lindep) + if len(x1) == 0: + return dpnp.zeros_like(b) + + x1 *= rmat.diagonal()[:,None] + + innerprod = [rmat[i,i].real ** 2 for i in range(x1.shape[0])] + max_innerprod = max(innerprod) + + if max_innerprod < lindep or max_innerprod < tol**2: + if x0 is None: + return dpnp.zeros_like(b) + else: + return x0 + + xs = [] + ax = [] + + max_cycle = min(max_cycle, ndim) + for cycle in range(max_cycle): + axt = aop(x1) + if axt.ndim == 1: + axt = axt.reshape(1,ndim) + xs.extend(x1) + ax.extend(axt) + if callable(callback): + callback(cycle, xs, ax) + x1 = axt.copy() + for i in range(len(xs)): + xsi = dpnp.asarray(xs[i]) + w = dpnp.dot(x1, xsi.conj()) / innerprod[i] + x1 -= xsi * dpnp.expand_dims(w,-1) + axt = xsi = None + x1, rmat = _stable_qr(x1, dpnp.dot, lindep=lindep) + x1 *= rmat.diagonal()[:,None] + innerprod1 = rmat.diagonal().real ** 2 + max_innerprod = max(innerprod1, default=0.) + + log.info(f'krylov cycle {cycle}, r = {max_innerprod**.5:.3e}, {x1.shape[0]} equations') + if max_innerprod < lindep or max_innerprod < tol**2: + break + mask = (innerprod1 > lindep) & (innerprod1 > tol**2) + x1 = x1[mask] + innerprod.extend(innerprod1[mask]) + if max_innerprod > 1e10: + raise RuntimeError('Krylov subspace iterations diverge') + + else: + raise RuntimeError('Krylov solver failed to converge') + + log.info(f'krylov space size {len(xs)}') + xs = dpnp.asarray(xs) + ax = dpnp.asarray(ax) + nd = xs.shape[0] + + h = dpnp.dot(xs, ax.T) + + # Add the contribution of I in (1+a) + h += dpnp.diag(dpnp.asarray(innerprod[:nd])) + g = dpnp.zeros((nd,nroots), dtype=x1.dtype) + + if b.ndim == 1: + g[0] = innerprod[0] + else: + # Restore the first nroots vectors, which are array b or b-(1+a)x0 + for i in range(min(nd, nroots)): + xsi = dpnp.asarray(xs[i]) + g[i] = dpnp.dot(xsi.conj(), b.T) + + c = dpnp.linalg.solve(h, g) + x = _gen_x0(c, dpnp.asarray(xs)) + if b.ndim == 1: + x = x[0] + + if x0 is not None: + x += x0 + return x + +def _qr(xs, dot, lindep=1e-14): + '''QR decomposition for a list of vectors (for linearly independent vectors only). + xs = (r.T).dot(qs) + ''' + nvec = len(xs) + dtype = xs[0].dtype + qs = dpnp.empty((nvec,xs[0].size), dtype=dtype) + rmat = dpnp.eye(nvec, order='F', dtype=dtype) + + nv = 0 + for i in range(nvec): + xi = dpnp.array(xs[i], copy=True) + prod = dot(qs[:nv].conj(), xi) + xi -= dpnp.dot(qs[:nv].T, prod) + + innerprod = dot(xi.conj(), xi).real + norm = innerprod**0.5 + if innerprod > lindep: + rmat[:,nv] -= dpnp.dot(rmat[:,:nv], prod) + qs[nv] = xi/norm + rmat[:nv+1,nv] /= norm + nv += 1 + return qs[:nv], dpnp.linalg.inv(rmat[:nv,:nv]) + +def _stable_qr(xs, dot, lindep=1e-14): + '''QR decomposition for a list of vectors (for linearly independent vectors only). + using the modified Gram-Schmidt process + ''' + nvec = len(xs) + dtype = xs[0].dtype + Q = dpnp.empty((nvec,xs[0].size), dtype=dtype) + R = dpnp.zeros((nvec,nvec), dtype=dtype) + V = xs.copy() + nv = 0 + for i in range(nvec): + norm = dpnp.linalg.norm(V[i]) + if norm**2 > lindep: + R[nv,nv] = norm + Q[nv] = V[i] / norm + R[nv, i+1:] = dot(Q[nv], V[i+1:].T) + V[i+1:] -= dpnp.outer(R[nv, i+1:], Q[nv]) + nv += 1 + return Q[:nv], R[:nv,:nv] + +def _gen_x0(v, xs): + ndim = v.ndim + if ndim == 1: + v = v[:,None] + space, nroots = v.shape + x0 = dpnp.einsum('c,x->cx', v[space-1], dpnp.asarray(xs[space-1])) + for i in reversed(range(space-1)): + xsi = dpnp.asarray(xs[i]) + x0 += dpnp.expand_dims(v[i],-1) * xsi + if ndim == 1: + x0 = x0[0] + return x0 + +def empty_mapped(shape, dtype=float, order='C'): + '''(experimental) + Returns a new, uninitialized NumPy array with the given shape and dtype. + + This is a convenience function which is just :func:`numpy.empty`, + except that the underlying buffer is a pinned and mapped memory. + This array can be used as the buffer of zero-copy memory. + ''' + nbytes = np.prod(shape) * np.dtype(dtype).itemsize + q = dpctl.SyclQueue() # or _create_from_ptr(int(libgpu.sycl_get_queue_ptr())) + mem = dpmem.MemoryUSMShared(nbytes, queue=q) # use MemoryUSMHost(...) if you don't need device mapping + # mem = cupy.cuda.PinnedMemoryPointer( + # cupy.cuda.PinnedMemory(nbytes, cupy.cuda.runtime.hostAllocMapped), 0) + out = np.ndarray(shape, dtype=dtype, buffer=mem, order=order) + return out + +def ndarray(shape, dtype=np.float64, buffer=None): + ''' + Construct a dpnp ndarray object using the NumPy ndarray API. + + dpnp counterpart of cupy_helper.ndarray. + + Args: + shape : tuple or int + Shape of the array to allocate. + + Kwargs: + dtype : Numpy data type. + + buffer : dpnp array + If buffer is specified, the array is a view over its memory. + Otherwise a new allocation is made. + ''' + # Normalise shape the way numpy.empty() does. Callers pass plain ints, + # tuples/lists, numpy integer scalars, and occasionally a device array + # produced by shape arithmetic. dpnp's __int__ rejects anything with + # ndim != 0, so `int(shape)` alone raises + # TypeError: only 0-dimensional arrays can be converted to Python scalars + # for a 1-element 1-D array, which numpy accepts as the sequence [n]. + if hasattr(shape, 'ndim') and not isinstance(shape, (list, tuple)): + # 0-d array is a scalar length; anything higher is a sequence of them. + shape = int(shape) if shape.ndim == 0 else tuple(int(s) for s in shape) + elif isinstance(shape, (list, tuple)): + shape = tuple(int(s) for s in shape) + else: + shape = int(shape) + + if buffer is None: + return dpnp.empty(shape, dtype=dtype) + + # cupy_helper builds the view from the raw pointer (memptr=buffer.data), + # so the buffer's own shape and dtype are irrelevant -- only its extent + # matters. dpnp validates the request against the buffer object instead + # and raises if the declared shape does not fit its element count, so + # reinterpret through a flat byte view to get the same semantics. + out_nbytes = int(np.prod(shape)) * np.dtype(dtype).itemsize + assert buffer.nbytes >= out_nbytes, ( + 'buffer of %d bytes is too small for the requested %d bytes' + % (buffer.nbytes, out_nbytes)) + return dpnp.ndarray(shape, dtype, buffer=buffer) + +def pinv(a, lindep=1e-10): + '''psudo-inverse with eigh, to be consistent with pyscf + ''' + a = dpnp.asarray(a) + w, v = dpnp.linalg.eigh(a) + mask = w > lindep + v1 = v[:,mask] + j2c = dpnp.dot(v1/w[mask], v1.conj().T) + return j2c + +def cond(a, sympos=False, verbose=logger.WARN): + """ + Calculate the condition number of a matrix. + + Parameters: + a (cupy.ndarray): The input matrix. + sympos : Whether the input matrix is symmetric and positive definite. + + Returns: + float: The condition number of the matrix. + """ + if isinstance(verbose, logger.Logger): + log = verbose + else: + log = logger.Logger(sys.stdout, verbose) + + if a.shape[0] > MAX_EIGH_DIM: + if not SCIPY_EIGH_FOR_LARGE_ARRAYS: + raise RuntimeError( + f'Array size exceeds the maximum size {MAX_EIGH_DIM}.') + a = a.get() + if sympos: + s = scipy.linalg.eigvalsh(a) + if s[0] > 0: + return s[-1] / s[0] + else: + log.warn(f'In condition number calculation, matrix is assumed to be positive definite, but it is not (minimal eigenvalue = {s[0]:e})') + _, s, _ = scipy.linalg.svd(a) + cond_number = s[0] / s[-1] + return cond_number + + else: + if sympos: + s = dpnp.linalg.eigvalsh(a) + if s[0] > 0: + return s[-1] / s[0] + else: + log.warn(f'In condition number calculation, matrix is assumed to be positive definite, but it is not (minimal eigenvalue = {s[0]:e})') + _, s, _ = dpnp.linalg.svd(a) + cond_number = s[0] / s[-1] + return cond_number + +def grouped_dot(As, Bs, Cs=None): + ''' + As: dpnp 2D array list. + Bs: dpnp 2D array list. + Cs: dpnp 2D array list. + einsum('ik,jk->ij', A, B, C) C=A@B.T + ''' + assert len(As) > 0 + assert len(As) == len(Bs) + assert As[0].flags.c_contiguous + assert Bs[0].flags.c_contiguous + groups = len(As) + + if Cs is None: + Cs = [] + for a, b in zip(As, Bs): + Cs.append(dpnp.empty((a.shape[0], b.shape[0]))) + + # Pure DPNP implementation using matmul with transpose + # C = A @ B.T (einsum 'ik,jk->ij') + for i in range(groups): + # B.T: transpose B so that (N, K) -> (K, N) + # Result: (M, K) @ (K, N) -> (M, N) + Cs[i][...] = dpnp.matmul(As[i], Bs[i].T) + + return Cs + +# def grouped_dot(As, Bs, Cs=None): +# ''' +# todo: layout of cutlass kernel +# As: dpnp 2D array list. +# Bs: dpnp 2D array list. +# Cs: dpnp 2D array list. +# einsum('ik,jk->ij', A, B, C) C=A@B.T +# ''' +# assert len(As) > 0 +# assert len(As) == len(Bs) +# assert As[0].flags.c_contiguous +# assert Bs[0].flags.c_contiguous +# groups = len(As) +# Ms, Ns, Ks = [], [], [] +# for a, b in zip(As, Bs): +# Ms.append(a.shape[0]) +# Ns.append(b.shape[0]) +# Ks.append(a.shape[1]) + +# if Cs is None: +# Cs = [] +# for i in range(groups): +# Cs.append(cupy.empty((Ms[i], Ns[i]))) + +# As_ptr, Bs_ptr, Cs_ptr = [], [], [] +# for a, b, c in zip(As, Bs, Cs): +# As_ptr.append(a.data.ptr) +# Bs_ptr.append(b.data.ptr) +# Cs_ptr.append(c.data.ptr) + +# As_ptr = np.array(As_ptr) +# Bs_ptr = np.array(Bs_ptr) +# Cs_ptr = np.array(Cs_ptr) + +# Ms = np.array(Ms) +# Ns = np.array(Ns) +# Ks = np.array(Ks) +# total_size = 68 * groups +# ''' +# 68 is the result of +# sizeof(cutlass::gemm::GemmCoord) + +# sizeof(typename DeviceKernel::ElementA*) + +# sizeof(typename DeviceKernel::ElementB*) + +# sizeof(typename DeviceKernel::ElementC*) + +# sizeof(typename DeviceKernel::ElementC*) + +# sizeof(int64_t) + sizeof(int64_t) + sizeof(int64_t) +# ''' +# padding = 8 - (total_size % 8) +# total_size += padding +# cutlass_space = cupy.empty(total_size, dtype=cupy.uint8) + +# stream = cupy.cuda.get_current_stream() +# err = libdpnp_helper.grouped_dot( +# ctypes.cast(stream.ptr, ctypes.c_void_p), +# ctypes.cast(Cs_ptr.ctypes.data, ctypes.c_void_p), +# ctypes.cast(As_ptr.ctypes.data, ctypes.c_void_p), +# ctypes.cast(Bs_ptr.ctypes.data, ctypes.c_void_p), +# ctypes.cast(Ms.ctypes.data, ctypes.c_void_p), +# ctypes.cast(Ns.ctypes.data, ctypes.c_void_p), +# ctypes.cast(Ks.ctypes.data, ctypes.c_void_p), +# ctypes.cast(cutlass_space.data.ptr, ctypes.c_void_p), +# ctypes.c_int(groups) +# ) +# if err != 0: +# raise RuntimeError('failed in grouped_gemm kernel') +# return Cs + +def grouped_gemm(As, Bs, Cs=None): + ''' + As: dpnp 2D array list. + Bs: dpnp 2D array list. + Cs: dpnp 2D array list. + assuming (X, 64).T @ (X, Y) + einsum('ki,kj->ij', A, B, C) C=A.T@B + Compare with grouped_dot, this function handles the case M < 128 + + Pure DPNP implementation: the CUTLASS grouped-GEMM kernel is not built + under SYCL (BUILD_CUTLASS defaults OFF), so this loops over dpnp.matmul. + ''' + assert len(As) > 0 + assert len(As) == len(Bs) + assert As[0].flags.c_contiguous + assert Bs[0].flags.c_contiguous + groups = len(As) + Ms, Ns, Ks = [], [], [] + for a, b in zip(As, Bs): + Ms.append(a.shape[1]) + Ns.append(b.shape[1]) + Ks.append(a.shape[0]) + + if Cs is None: + Cs = [] + for i in range(groups): + Cs.append(dpnp.empty((Ms[i], Ns[i]))) + + # Pure DPNP implementation using matmul with transpose + # C = A.T @ B (einsum 'ki,kj->ij') + for i in range(groups): + # A.T: transpose A so that (K, M) -> (M, K) + # Result: (M, K) @ (K, N) -> (M, N) + Cs[i][...] = dpnp.matmul(As[i].T, Bs[i]) + + return Cs + +# def condense(opname, a, loc_x, loc_y=None): +# """ +# dpnp version of condense() without any explicit SYCL kernel. + +# Parameters +# ---------- +# opname : {'sum', 'max', 'min', 'abssum', 'absmax', 'norm'} +# a : np.ndarray or dpnp.ndarray, float64, ndim >= 2 +# loc_x : 1D array-like of ints (partition on axis -2) +# loc_y : 1D array-like of ints (partition on axis -1), defaults to loc_x + +# Returns +# ------- +# out : dpnp.ndarray (or numpy if you convert back) +# Shape (len(loc_x)-1, len(loc_y)-1) (possibly transposed for Fortran input). +# """ +# assert opname in ("sum", "max", "min", "abssum", "absmax", "norm") +# assert a.dtype == np.float64 +# assert a.ndim >= 2 +# if loc_y is None: +# loc_y = loc_x +# do_transpose = False +# loc_x = np.asarray(loc_x, dtype=np.int64) +# loc_y = np.asarray(loc_y, dtype=np.int64) + +# if a.ndim == 2: +# if a.flags.f_contiguous: +# a = dpnp.transpose(a) +# loc_x, loc_y = loc_y, loc_x +# do_transpose = True +# a = a[None] +# else: +# nx, ny = a.shape[-2:] +# a = a.reshape(-1, nx, ny) + +# counts, nx, ny = a.shape +# assert loc_x[-1] == nx +# assert loc_y[-1] == ny + +# # Move to device +# a_dev = dp.asarray(a) # shape (counts, nx, ny) +# loc_x_dev = loc_x # indices are small, fine on host +# loc_y_dev = loc_y + +# nloc_x = loc_x_dev.size - 1 +# nloc_y = loc_y_dev.size - 1 + +# out = dp.zeros((nloc_x, nloc_y), dtype=dp.float64) + +# for i in range(nloc_x): +# i0, i1 = loc_x_dev[i], loc_x_dev[i + 1] +# for j in range(nloc_y): +# j0, j1 = loc_y_dev[j], loc_y_dev[j + 1] + +# # Slice all counts, block in x,y -> shape (counts, i1-i0, j1-j0) +# block = a_dev[:, i0:i1, j0:j1] + +# if opname == "sum": +# val = dp.sum(block) # over all axes +# elif opname == "max": +# val = dp.max(block) +# elif opname == "min": +# val = dp.min(block) +# elif opname == "abssum": +# val = dp.sum(dp.abs(block)) +# elif opname == "absmax": +# val = dp.max(dp.abs(block)) +# elif opname == "norm": +# # sqrt of sum of squares over all elements +# val = dp.sqrt(dp.sum(block * block)) + +# out[i, j] = val + +# if do_transpose: +# out = dpnp.transpose(out) +# return out + +def absmax(a): + '''abs(a).max() while limiting temporary memory use. The optimization is + only valid for real-valued arrays. + ''' + if a.dtype == np.complex128 or a.nbytes < MEMPOOL_THRESHOLD: + return abs(a).max() + return max(a.max(), -a.min()) + +def condense(opname, a, loc_x, loc_y=None): + """ + DPNP/SYCL port of condense(): reduce over the last two dims in windows. + Reduces across counts and the i/j window just like the CUDA kernel. + + """ + assert opname in ('sum', 'max', 'min', 'abssum', 'absmax', 'norm') + assert a.dtype == np.float64 + assert a.ndim >= 2 + if loc_y is None: + loc_y = loc_x + do_transpose = False + if a.ndim == 2: + # Match CUDA path: if input is F-contig, transpose and swap locators + if a.flags.f_contiguous: + a = dpnp.transpose(a) + loc_x, loc_y = loc_y, loc_x + do_transpose = True + a = a[None, ...] # shape -> (counts=1, nx, ny) + else: + nx, ny = int(a.shape[-2]), int(a.shape[-1]) + a = a.reshape(-1, nx, ny) # (counts, nx, ny) + + # Work with host-side integer indices; windows stay on device + a = dpnp.asarray(a, order='C') + loc_x = np.asarray(loc_x, dtype=np.int32) + loc_y = np.asarray(loc_y, dtype=np.int32) + nloc_x = loc_x.size - 1 + nloc_y = loc_y.size - 1 + counts, nx, ny = a.shape + assert loc_x[-1] == nx + assert loc_y[-1] == ny + + out = dpnp.zeros((nloc_x, nloc_y), dtype=a.dtype) + + # Helper for a single window reduction + def _reduce_window(win): + if opname == 'sum': + return dpnp.sum(win) + elif opname == 'max': + return dpnp.max(win) + elif opname == 'min': + return dpnp.min(win) + elif opname == 'abssum': + return dpnp.sum(dpnp.abs(win)) + elif opname == 'absmax': + return dpnp.max(dpnp.abs(win)) + elif opname == 'norm': + return dpnp.sqrt(dpnp.sum(win * win)) + else: + raise ValueError(opname) + + # Host loops over blocks; device does heavy reductions per window + for i in range(nloc_x): + i0, i1 = int(loc_x[i]), int(loc_x[i+1]) + for j in range(nloc_y): + j0, j1 = int(loc_y[j]), int(loc_y[j+1]) + win = a[:, i0:i1, j0:j1] # (counts, i1-i0, j1-j0) on device + out[i, j] = _reduce_window(win) # device reduction + + if do_transpose: + out = dpnp.transpose(out) + + return out + +def sandwich_dot(a, c, out=None): + '''Performs c.T.dot(a).dot(c)''' + a = dpnp.asarray(a) + c = dpnp.asarray(c) + a_ndim = a.ndim + if a_ndim == 2: + a = a[None] + counts = a.shape[0] + m = c.shape[1] + dtype = dpnp.result_type(a, c) + out = dpnp.empty((counts, m, m), dtype=dtype) + tmp = None + for i in range(counts): + tmp = dpnp.dot(c.conj().T, a[i], out=tmp) + dpnp.dot(tmp, c, out=out[i]) + if a_ndim == 2: + out = out[0] + return out + +def set_conditional_mempool_malloc(n_bytes_threshold=MEMPOOL_THRESHOLD): + """No-op: SYCL/USM manages memory automatically. + + In CuPy, this sets conditional memory pool allocation based on size. + With DPNP/SYCL USM, memory management is handled by the runtime. + """ + pass +# def set_conditional_mempool_malloc(n_bytes_threshold=100000000): +# ''' +# Customize CuPy memory allocator. + +# For large memory allocations (>100MB by default), the custom allocator bypasses +# the CuPy memory pool, directly calling the CUDA malloc API. The large memory +# chunks will be released back to the system when the associated object is +# destroyed. Only small memory blocks are allocated from the CuPy memory pool. + +# Execute the following command to restore the default CuPy malloc +# cupy.cuda.set_allocator(cupy.get_default_memory_pool().malloc) +# ''' +# cuda_malloc = cupy.cuda.memory._malloc +# default_mempool_malloc = cupy.get_default_memory_pool().malloc +# def malloc(size): +# if size >= n_bytes_threshold: +# return cuda_malloc(size) +# return default_mempool_malloc(size) +# cupy.cuda.set_allocator(malloc) + +def batched_vec3_norm2(batched_vec3): + ''' + einsum('gx,gx->g', vec3, vec3) for the (N,3)-array vec3 + + Accepts either C-order (N,3) or F-order (3,N) layout, mirroring the CuPy + implementation. All work stays on the device; no host transfers. + ''' + assert type(batched_vec3) is dpnp.ndarray + assert batched_vec3.dtype == dpnp.float64 + assert batched_vec3.ndim == 2 + assert batched_vec3.shape[0] == 3 or batched_vec3.shape[1] == 3 + assert batched_vec3.flags.c_contiguous + + order = "c" if batched_vec3.shape[1] == 3 else "f" + + n = batched_vec3.shape[0] if order == "c" else batched_vec3.shape[1] + assert n != 3, "Ambiguous array order, cannot determine if the array is C or Fortran order from the shape" + assert n * 3 < np.iinfo(np.int32).max + + if order == "c": + return batched_vec_norm2(batched_vec3) + else: + return batched_vec_norm2(batched_vec3.T) + +def batched_vec_norm2(vec, out=None): + ''' + einsum('gx,gx->g', vec, vec) + + `vec` is expected to be a device (dpnp) array; dpnp.asarray is a no-op for + device arrays, so no host<->device transfer occurs. Both C- and F-contiguous + inputs are supported (callers pass transposed views, e.g. nabla_rho_i.T). + ''' + vec = dpnp.asarray(vec) + assert vec.dtype == dpnp.float64 + assert vec.ndim == 2 + n, x = vec.shape + out = ndarray(n, np.float64, out) + dpnp.einsum("ij,ij->i", vec, vec, out=out) + return out + +def vec_dot(vec1, vec2): + ''' + einsum('g,g->', vec1, vec2) + + dpnp counterpart of cupy_helper.vec_dot. Both inputs are expected to be + device (dpnp) arrays; no host transfers. + ''' + vec1 = dpnp.asarray(vec1) + vec2 = dpnp.asarray(vec2) + assert vec1.dtype == dpnp.float64 + assert vec2.dtype == dpnp.float64 + assert vec1.shape == vec2.shape + return dpnp.einsum("i,i->", vec1.ravel(), vec2.ravel()) + +def batched_vec_dot(vec1, vec2, out=None): + ''' + einsum('gx,gx->g', vec1, vec2) + + Both inputs are expected to be device (dpnp) arrays; no host transfers. + ''' + vec1 = dpnp.asarray(vec1) + vec2 = dpnp.asarray(vec2) + assert vec1.dtype == dpnp.float64 + assert vec2.dtype == dpnp.float64 + assert vec1.ndim == 2 + assert vec1.shape == vec2.shape + n, x = vec1.shape + out = ndarray(n, np.float64, out) + dpnp.einsum("ij,ij->i", vec1, vec2, out=out) + return out + +class LinAlgError(RuntimeError): + pass + +def cholesky(a, /, *, upper=False): + ''' + dpnp counterpart of cupy_helper.cholesky (cusolver.cholesky). + + dpnp.linalg.LinAlgError derives from ValueError, while the CUDA path's + cusolver.LinAlgError derives from RuntimeError. Several call sites + (df.py, df_jk.py, grad/rhf.py) rely on a failed Cholesky raising + RuntimeError to fall back to an eigendecomposition, so re-raise here + as this module's RuntimeError-derived LinAlgError. + ''' + try: + return dpnp.linalg.cholesky(a, upper=upper) + except dpnp.linalg.LinAlgError as e: + raise LinAlgError(str(e)) from e + +def eigh(a, b=None, overwrite=False): + ''' + Solve a standard or generalized eigenvalue problem for a complex + Hermitian or real symmetric matrix. + + Note: both a and b matrices are overwritten when overwrite is specified. + ''' + if b is None: + if a.shape[0] > 32600: + if not SCIPY_EIGH_FOR_LARGE_ARRAYS: + raise RuntimeError('Array is too large for DPNP eigh.') + a = a.get() + e, c = scipy.linalg.eigh(a, overwrite_a=True) + e = asarray(e) + c = asarray(c) + return e, c + return dpnp.linalg.eigh(a) + + if a.shape[0] > MAX_EIGH_DIM: + if not SCIPY_EIGH_FOR_LARGE_ARRAYS: + raise RuntimeError( + f'Array size exceeds the maximum size {MAX_EIGH_DIM}.') + a = a.get() + b = b.get() + e, c = scipy.linalg.eigh(a, b, overwrite_a=True) + e = asarray(e) + c = asarray(c) + return e, c + + return onemkl_eigh(a, b, overwrite) + +def stack_with_padding(arrays): + ''' + Stack orbital coefficients, padding zeros to smaller arrays + ''' + if not arrays: + raise ValueError("arrays must be a non-empty sequence") + + max_nmo = max(a.shape[1] for a in arrays) + nao = arrays[0].shape[0] + dtype = np.result_type(*arrays) + out = dpnp.empty((len(arrays), nao, max_nmo), dtype=dtype) + + for k, a in enumerate(arrays): + nmo = a.shape[1] + out[k,:,:nmo] = a + if nmo < max_nmo: + out[k,:,nmo:] = 0 + return out + +def empty_aligned(shape, dtype, alignment=128): + ''' + Allocate an array with a memory alignment. + + Args: + shape : tuple or int + Shape of the array to allocate. + + Kwargs: + dtype : Numpy data type. + + alignment : int + Byte alignment for the underlying device memory pointer. + 128 bytes is optimal for coalesced global memory access on most CUDA + architectures. + + ''' + dtype = np.dtype(dtype) + size = int(np.prod(shape)) + nbytes = size * dtype.itemsize + alignment + buf = cupy.empty(nbytes, dtype=np.uint8) + offset = (alignment - buf.data.ptr % alignment) % alignment + return ndarray(shape, dtype, buf[offset:]) diff --git a/gpu4pyscf/lib/ecp/CMakeLists.txt b/gpu4pyscf/lib/ecp/CMakeLists.txt index 8bd20c852..bde5b3266 100644 --- a/gpu4pyscf/lib/ecp/CMakeLists.txt +++ b/gpu4pyscf/lib/ecp/CMakeLists.txt @@ -12,36 +12,44 @@ # See the License for the specific language governing permissions and # limitations under the License. -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128") +set(GPU_SRCS + nr_ecp_driver.cu +) -# If compiling with CUDA<13.1 and arch>100, ECP will silently provide incorrect result. -set(HAS_BLACKWELL_OR_NEWER OFF) -foreach(arch IN LISTS CMAKE_CUDA_ARCHITECTURES) - string(REGEX MATCH "([0-9]+)" arch_number ${arch}) - if("${arch_number}" STREQUAL "") - message(FATAL_ERROR "Cannot extract architecture id from CMAKE_CUDA_ARCHITECTURES term ${arch}.") - endif() - if(arch_number GREATER_EQUAL 100) - set(HAS_BLACKWELL_OR_NEWER ON) - break() - endif() -endforeach() +add_library(gecp SHARED ${GPU_SRCS}) -set(CUDA_VERSION_LESS_THAN_13_1 OFF) -if (CMAKE_CUDA_COMPILER_VERSION VERSION_LESS "13.1") - set(CUDA_VERSION_LESS_THAN_13_1 ON) -endif() +if (USE_SYCL) + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + target_compile_options(gecp PRIVATE -x c++ -nocudainc -nocudalib) + target_link_libraries(gecp PRIVATE sycl_compat gsycl) +else() + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128") -if(HAS_BLACKWELL_OR_NEWER AND CUDA_VERSION_LESS_THAN_13_1) - message(WARNING "Blackwell or newer GPU detected with sm ${arch_number} >= 100, and a CUDA version ${CMAKE_CUDA_COMPILER_VERSION} < 13.1 is detected. " - "To get around with a severe bug in nvcc, we turn off compiler optimization. The performance will be damaged.") - set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --Ofast-compile max") -endif() -# End of hacking CUDA<13.1 and arch>100 + # If compiling with CUDA<13.1 and arch>100, ECP will silently provide incorrect result. + set(HAS_BLACKWELL_OR_NEWER OFF) + foreach(arch IN LISTS CMAKE_CUDA_ARCHITECTURES) + string(REGEX MATCH "([0-9]+)" arch_number ${arch}) + if("${arch_number}" STREQUAL "") + message(FATAL_ERROR "Cannot extract architecture id from CMAKE_CUDA_ARCHITECTURES term ${arch}.") + endif() + if(arch_number GREATER_EQUAL 100) + set(HAS_BLACKWELL_OR_NEWER ON) + break() + endif() + endforeach() -add_library(gecp SHARED - nr_ecp_driver.cu -) + set(CUDA_VERSION_LESS_THAN_13_1 OFF) + if (CMAKE_CUDA_COMPILER_VERSION VERSION_LESS "13.1") + set(CUDA_VERSION_LESS_THAN_13_1 ON) + endif() + + if(HAS_BLACKWELL_OR_NEWER AND CUDA_VERSION_LESS_THAN_13_1) + message(WARNING "Blackwell or newer GPU detected with sm ${arch_number} >= 100, and a CUDA version ${CMAKE_CUDA_COMPILER_VERSION} < 13.1 is detected. " + "To get around with a severe bug in nvcc, we turn off compiler optimization. The performance will be damaged.") + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --Ofast-compile max") + endif() + # End of hacking CUDA<13.1 and arch>100 +endif() set_target_properties(gecp PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) diff --git a/gpu4pyscf/lib/ecp/common.cu b/gpu4pyscf/lib/ecp/common.cu index e05b9f1a2..ba0d87917 100644 --- a/gpu4pyscf/lib/ecp/common.cu +++ b/gpu4pyscf/lib/ecp/common.cu @@ -44,6 +44,12 @@ Cartesian<(l+1)*(l+2)/2> ang_nuc_l(double rx, double ry, double rz){ __device__ double rad_part(const int ish, const int *ecpbas, const double *env){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + #else + const int threadIdx_x = threadIdx.x; + #endif const int npk = ecpbas[ish*BAS_SLOTS+NPRIM_OF]; const int r_order = ecpbas[ish*BAS_SLOTS+RADI_POWER]; const int exp_ptr = ecpbas[ish*BAS_SLOTS+PTR_EXP]; @@ -51,8 +57,8 @@ double rad_part(const int ish, const int *ecpbas, const double *env){ double u1 = 0.0; double r = 0.0; - if (threadIdx.x < NGAUSS){ - r = r128[threadIdx.x]; + if (threadIdx_x < NGAUSS){ + r = r128[threadIdx_x]; } for (int kp = 0; kp < npk; kp++){ const double ak = env[exp_ptr+kp]; @@ -60,8 +66,8 @@ double rad_part(const int ish, const int *ecpbas, const double *env){ u1 += ck * exp(-ak * r * r); } double w = 0.0; - if (threadIdx.x < NGAUSS){ - w = w128[threadIdx.x]; + if (threadIdx_x < NGAUSS){ + w = w128[threadIdx_x]; } return u1 * pow(r, r_order) * w; } @@ -118,8 +124,14 @@ void cache_fac(double *fx, double *ri){ __device__ void block_reduce(double val, double *d_out) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + double (&sdata)[THREADS] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + const unsigned int tid = item.get_local_id(0); +#else // USE_SYCL __shared__ double sdata[THREADS]; const unsigned int tid = threadIdx.x; +#endif sdata[tid] = val; __syncthreads(); @@ -150,7 +162,15 @@ void block_reduce(double val, double *d_out) { __device__ __forceinline__ void set_shared_memory(double *smem, const int size) { - for (int i = threadIdx.x; i < size; i += blockDim.x) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif + for (int i = threadIdx_x; i < size; i += blockDim_x) { smem[i] = 0.0; } __syncthreads(); @@ -158,6 +178,14 @@ void set_shared_memory(double *smem, const int size) { __device__ void _li_up(double *out, double *buf, const int li, const int lj){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif const int nfj = (lj+1) * (lj+2) / 2; const int nfi = (li+1) * (li+2) / 2; const int nfi0 = li * (li+1) / 2; @@ -165,7 +193,7 @@ void _li_up(double *out, double *buf, const int li, const int lj){ double *outy = outx + nfi*nfj; double *outz = outy + nfi*nfj; const double fac = 1.0 / _ecp_fac[li-1]; - for (int ij = threadIdx.x; ij < nfi0*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi0*nfj; ij+=blockDim_x){ const int i = ij % nfi0; const int j = ij / nfi0; const double yfac = fac * (_cart_pow_y[i] + 1); @@ -180,6 +208,14 @@ void _li_up(double *out, double *buf, const int li, const int lj){ __device__ void _li_up_and_write(double *out, double *buf, const int li, const int lj, const int nao){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif const int nfi0 = li * (li+1) / 2; const int nfj = (lj+1) * (lj+2) / 2; double *outxx = out ; @@ -192,7 +228,7 @@ void _li_up_and_write(double *out, double *buf, const int li, const int lj, cons double *outzy = out + 7*nao*nao; double *outzz = out + 8*nao*nao; const double fac = 1.0 / _ecp_fac[li-1]; - for (int ij = threadIdx.x; ij < nfi0*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi0*nfj; ij+=blockDim_x){ const int i = ij % nfi0; const int j = ij / nfi0; const double yfac = fac * (_cart_pow_y[i] + 1); @@ -217,6 +253,14 @@ void _li_up_and_write(double *out, double *buf, const int li, const int lj, cons __device__ void _li_down(double *out, double *buf, const int li, const int lj){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif const int nfi = (li+1) * (li+2) / 2; const int nfj = (lj+1) * (lj+2) / 2; const int nfi1= (li+2) * (li+3) / 2; @@ -225,7 +269,7 @@ void _li_down(double *out, double *buf, const int li, const int lj){ double *outz = outy + nfi*nfj; const double fac = _ecp_fac[li]; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij % nfi; const int j = ij / nfi; atomicAdd(outx + j*nfi+i, fac * buf[j*nfi1+i]); @@ -236,6 +280,14 @@ void _li_down(double *out, double *buf, const int li, const int lj){ __device__ void _li_down_and_write(double *out, double *buf, const int li, const int lj, const int nao){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif const int nfi = (li+1) * (li+2) / 2; const int nfj = (lj+1) * (lj+2) / 2; const int nfi1= (li+2) * (li+3) / 2; @@ -250,7 +302,7 @@ void _li_down_and_write(double *out, double *buf, const int li, const int lj, co double *outzz = out + 8*nao*nao; const double fac = _ecp_fac[li]; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij % nfi; const int j = ij / nfi; const int i_addr[3] = {i, _y_addr[i], _z_addr[i]}; @@ -272,6 +324,14 @@ void _li_down_and_write(double *out, double *buf, const int li, const int lj, co __device__ void _lj_up_and_write(double *out, double *buf, const int li, const int lj, const int nao){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif const int nfi = (li+1)*(li+2)/2; const int nfj0 = lj * (lj+1) / 2; double *outxx = out; @@ -284,7 +344,7 @@ void _lj_up_and_write(double *out, double *buf, const int li, const int lj, cons double *outzy = out + 7*nao*nao; double *outzz = out + 8*nao*nao; const double fac = 1.0 / _ecp_fac[lj-1]; - for (int ij = threadIdx.x; ij < nfi*nfj0; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj0; ij+=blockDim_x){ const int i = ij % nfi; const int j = ij / nfi; const double yfac = fac * (_cart_pow_y[j] + 1); @@ -308,6 +368,14 @@ void _lj_up_and_write(double *out, double *buf, const int li, const int lj, cons __device__ void _lj_down_and_write(double *out, double *buf, const int li, const int lj, const int nao){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif const int nfi = (li+1) * (li+2) / 2; const int nfj = (lj+1) * (lj+2) / 2; const int nfj1 = (lj+2) * (lj+3) / 2; @@ -321,7 +389,7 @@ void _lj_down_and_write(double *out, double *buf, const int li, const int lj, co double *outzy = out + 7*nao*nao; double *outzz = out + 8*nao*nao; const double fac = _ecp_fac[lj]; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij % nfi; const int j = ij / nfi; const int j_addr[3] = {j, _y_addr[j], _z_addr[j]}; @@ -459,4 +527,4 @@ void _li_down_down(double *out, double *buf, const int li, const int lj){ } __syncthreads(); } -*/ \ No newline at end of file +*/ diff --git a/gpu4pyscf/lib/ecp/ecp.h b/gpu4pyscf/lib/ecp/ecp.h index 5b67e5b02..c2887894b 100644 --- a/gpu4pyscf/lib/ecp/ecp.h +++ b/gpu4pyscf/lib/ecp/ecp.h @@ -1,3 +1,7 @@ +#ifdef USE_SYCL +#include +#endif + #define EXPCUTOFF 39 // 1e-17 #define CUTOFF 460 // ~ 1e200 #define CART_CUM (165) diff --git a/gpu4pyscf/lib/ecp/ecp_type1.cu b/gpu4pyscf/lib/ecp/ecp_type1.cu index 488e83b1a..6e6417f69 100644 --- a/gpu4pyscf/lib/ecp/ecp_type1.cu +++ b/gpu4pyscf/lib/ecp/ecp_type1.cu @@ -17,11 +17,19 @@ __device__ void type1_rad_part(double* __restrict__ rad_all, const int LIJ, double k, double aij, double ur) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif const double kaij = k / (2*aij); const double fac = kaij * kaij * aij; double r = 0.0; - if (threadIdx.x < NGAUSS){ - r = r128[threadIdx.x]; + if (threadIdx_x < NGAUSS){ + r = r128[threadIdx_x]; } double tmp = r - kaij; tmp = fac - aij*tmp*tmp; @@ -38,7 +46,7 @@ void type1_rad_part(double* __restrict__ rad_all, const int LIJ, double k, doubl _ine(bval, LIJ, k*r); } - for (int i = threadIdx.x; i < LIJ1*LIJ1; i+=blockDim.x){ + for (int i = threadIdx_x; i < LIJ1*LIJ1; i+=blockDim_x){ rad_all[i] = 0.0; } __syncthreads(); @@ -93,6 +101,14 @@ double type1_ang_nuc_l(const int i, const int j, const int k, double *unitr){ __device__ void type1_rad_ang(double *rad_ang, const int LIJ, double *r, double *rad_all, const double fac) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif double unitr[3]; if (r[0]*r[0] + r[1]*r[1] + r[2]*r[2] < 1e-16){ unitr[0] = 0; @@ -107,7 +123,7 @@ void type1_rad_ang(double *rad_ang, const int LIJ, double *r, double *rad_all, c // loop over i+j+k<=LIJ // TODO: find a closed form? - for (int n = threadIdx.x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim_x){ const int i = n/(LIJ+1)/(LIJ+1); const int j = n/(LIJ+1)%(LIJ+1); const int k = n%(LIJ+1); @@ -127,7 +143,7 @@ void type1_rad_ang(double *rad_ang, const int LIJ, double *r, double *rad_all, c //atomicAdd(rad_ang + i*(LIJ+1)*(LIJ+1) + j*(LIJ+1) + k, fac*s); } - for (int n = threadIdx.x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim_x){ const int i = n/(LIJ+1)/(LIJ+1); const int j = n/(LIJ+1)%(LIJ+1); const int k = n%(LIJ+1); @@ -150,6 +166,14 @@ void type1_rad_ang(double *rad_ang, const int LIJ, double *r, double *rad_all, c template __device__ void type1_rad_ang(double *rad_ang, double *r, double *rad_all, const double fac) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif double unitr[3]; if (r[0]*r[0] + r[1]*r[1] + r[2]*r[2] < 1e-16){ unitr[0] = 0; @@ -164,7 +188,7 @@ void type1_rad_ang(double *rad_ang, double *r, double *rad_all, const double fac // loop over i+j+k<=LIJ // TODO: find a closed form? - for (int n = threadIdx.x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim_x){ const int i = n/(LIJ+1)/(LIJ+1); const int j = n/(LIJ+1)%(LIJ+1); const int k = n%(LIJ+1); @@ -184,7 +208,7 @@ void type1_rad_ang(double *rad_ang, double *r, double *rad_all, const double fac //atomicAdd(rad_ang + i*(LIJ+1)*(LIJ+1) + j*(LIJ+1) + k, fac*s); } - for (int n = threadIdx.x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LIJ+1)*(LIJ+1)*(LIJ+1); n+=blockDim_x){ const int i = n/(LIJ+1)/(LIJ+1); const int j = n/(LIJ+1)%(LIJ+1); const int k = n%(LIJ+1); @@ -211,7 +235,21 @@ void type1_cart(double *gctr, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { + constexpr int LIJ1 = LI+LJ+1; + +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + auto thread_block = item.get_group(); + const int task_id = thread_block.get_group_id(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + double (&rad_ang)[LIJ1*LIJ1*LIJ1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + __shared__ double rad_ang[LIJ1*LIJ1*LIJ1]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -247,8 +285,6 @@ void type1_cart(double *gctr, ur += rad_part(kbas, ecpbas, env); } - constexpr int LIJ1 = LI+LJ+1; - __shared__ double rad_ang[LIJ1*LIJ1*LIJ1]; set_shared_memory(rad_ang, LIJ1*LIJ1*LIJ1); const double fac = 16.0 * M_PI * M_PI * _common_fac[LI] * _common_fac[LJ]; @@ -261,7 +297,11 @@ void type1_cart(double *gctr, const double k = 2.0 * norm3d(rij[0], rij[1], rij[2]); const double aij = ai[ip] + aj[jp]; + #ifdef USE_SYCL + double (&rad_all)[LIJ1*LIJ1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else __shared__ double rad_all[LIJ1*LIJ1]; + #endif type1_rad_part(rad_all, LI+LJ, k, aij, ur); __syncthreads(); @@ -280,7 +320,7 @@ void type1_cart(double *gctr, double fj[3*nfj]; cache_fac(fj, rcb); - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int mi = ij%nfi; const int mj = ij/nfi; @@ -331,9 +371,22 @@ void type1_cart(double *gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { +#ifdef USE_SYCL + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); +#else const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + extern __shared__ double smem[]; +#endif if (task_id >= ntasks){ return; } @@ -342,8 +395,6 @@ void type1_cart(double *gctr, const int jsh = tasks[task_id + ntasks]; const int ksh = tasks[task_id + 2*ntasks]; - extern __shared__ double smem[]; - const int npi = bas[NPRIM_OF+ish*BAS_SLOTS]; const int npj = bas[NPRIM_OF+jsh*BAS_SLOTS]; const double *ai = env + bas[PTR_EXP+ish*BAS_SLOTS]; @@ -403,7 +454,7 @@ void type1_cart(double *gctr, } const int nfi = (LI+1) * (LI+2) / 2; const int nfj = (LJ+1) * (LJ+2) / 2; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ double fi[3*NF_MAX]; cache_fac(fi, LI, rca); double fj[3*NF_MAX]; @@ -447,7 +498,7 @@ void type1_cart(double *gctr, const int ioff = ao_loc[ish]; const int joff = ao_loc[jsh]; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double tmp = reg_gctr[ij/THREADS]; @@ -458,4 +509,3 @@ void type1_cart(double *gctr, } return; } - diff --git a/gpu4pyscf/lib/ecp/ecp_type1_ip.cu b/gpu4pyscf/lib/ecp/ecp_type1_ip.cu index 846b8f498..86d19bc23 100644 --- a/gpu4pyscf/lib/ecp/ecp_type1_ip.cu +++ b/gpu4pyscf/lib/ecp/ecp_type1_ip.cu @@ -20,6 +20,23 @@ void type1_cart_unrolled_kernel(double *gctr, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { + constexpr int LIJ1 = LI+LJ+1; + constexpr int LIJ3 = LIJ1*LIJ1*LIJ1; + +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + + auto thread_block = item.get_group(); + double (&rad_ang)[LIJ3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else // USE_SYCL + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + + __shared__ double rad_ang[LIJ3]; +#endif // USE_SYCL + const int npi = bas[NPRIM_OF+ish*BAS_SLOTS]; const int npj = bas[NPRIM_OF+jsh*BAS_SLOTS]; const double *ai = env + bas[PTR_EXP+ish*BAS_SLOTS]; @@ -47,10 +64,6 @@ void type1_cart_unrolled_kernel(double *gctr, ur += rad_part(kbas, ecpbas, env); } - constexpr int LIJ1 = LI+LJ+1; - constexpr int LIJ3 = LIJ1*LIJ1*LIJ1; - - __shared__ double rad_ang[LIJ3]; set_shared_memory(rad_ang, LIJ3); const double fac = 16.0 * M_PI * M_PI * _common_fac[LI] * _common_fac[LJ]; @@ -65,7 +78,11 @@ void type1_cart_unrolled_kernel(double *gctr, const double k = 2.0 * norm3d(rij[0], rij[1], rij[2]); const double aij = ai_prim + aj_prim; + #ifdef USE_SYCL + double (&rad_all)[LIJ1*LIJ1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else __shared__ double rad_all[LIJ1*LIJ1]; + #endif type1_rad_part(rad_all, LI+LJ, k, aij, ur); __syncthreads(); @@ -87,7 +104,7 @@ void type1_cart_unrolled_kernel(double *gctr, constexpr int nfi = (LI+1) * (LI+2) / 2; constexpr int nfj = (LJ+1) * (LJ+2) / 2; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int mi = ij%nfi; const int mj = ij/nfi; @@ -123,14 +140,26 @@ void type1_cart_unrolled_kernel(double *gctr, } } +// `smem` is passed in rather than declared here as `extern __shared__`: +// under SYCL the dynamic local memory comes from a sycl::local_accessor +// created by the enclosing submit(), which cannot be reached from a +// device function. The CUDA callers pass their own `extern __shared__` +// block, so behaviour is unchanged. template __device__ -void type1_cart_kernel(double *gctr, +void type1_cart_kernel(double *smem, double *gctr, const int LI, const int LJ, const int ish, const int jsh, const int ksh, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { - extern __shared__ double smem[]; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); +#else // USE_SYCL + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; +#endif // USE_SYCL const int npi = bas[NPRIM_OF+ish*BAS_SLOTS]; const int npj = bas[NPRIM_OF+jsh*BAS_SLOTS]; @@ -163,7 +192,7 @@ void type1_cart_kernel(double *gctr, const int LIJ3 = LIJ1*LIJ1*LIJ1; double *rad_ang = smem; - for (int i = threadIdx.x; i < LIJ3; i+=blockDim.x) { + for (int i = threadIdx_x; i < LIJ3; i+=blockDim_x) { rad_ang[i] = 0; } __syncthreads(); @@ -200,7 +229,7 @@ void type1_cart_kernel(double *gctr, const int nfi = (LI+1) * (LI+2) / 2; const int nfj = (LJ+1) * (LJ+2) / 2; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int mi = ij%nfi; const int mj = ij/nfi; @@ -244,7 +273,24 @@ void type1_cart_ip1(double *gctr, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { + constexpr int nfi = (LI+1) * (LI+2) / 2; + constexpr int nfj = (LJ+1) * (LJ+2) / 2; + constexpr int nfi1 = (LI+2)*(LI+3)/2; + +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + double (&gctr_smem)[nfi*nfj*3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&buf)[nfi1*nfj] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + __shared__ double gctr_smem[nfi*nfj*3]; + __shared__ double buf[nfi1*nfj]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -257,17 +303,11 @@ void type1_cart_ip1(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += 3*ecp_id*nao*nao + ioff*nao + joff; - constexpr int nfi = (LI+1) * (LI+2) / 2; - constexpr int nfj = (LJ+1) * (LJ+2) / 2; - __shared__ double gctr_smem[nfi*nfj*3]; - for (int ij = threadIdx.x; ij < nfi*nfj*3; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj*3; ij+=blockDim_x){ gctr_smem[ij] = 0.0; } __syncthreads(); - constexpr int nfi1 = (LI+2)*(LI+3)/2; - __shared__ double buf[nfi1*nfj]; - type1_cart_unrolled_kernel<1,0,LI+1,LJ>( buf, ish, jsh, ksh, ecpbas, ecploc, @@ -286,7 +326,7 @@ void type1_cart_ip1(double *gctr, __syncthreads(); } - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double *gx = gctr; @@ -305,9 +345,29 @@ void type1_cart_ip1_general(double *gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { + constexpr int nfi_max = (AO_LMAX+2)*(AO_LMAX+3)/2; + constexpr int nfj_max = (AO_LMAX+1)*(AO_LMAX+2)/2; + +#ifdef USE_SYCL + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + double (&gctr_smem)[NF_MAX*NF_MAX*3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&buf)[nfi_max*nfj_max] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + __shared__ double gctr_smem[NF_MAX*NF_MAX*3]; + __shared__ double buf[nfi_max*nfj_max]; + extern __shared__ double smem[]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -320,17 +380,12 @@ void type1_cart_ip1_general(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += 3*ecp_id*nao*nao + ioff*nao + joff; - __shared__ double gctr_smem[NF_MAX*NF_MAX*3]; - for (int ij = threadIdx.x; ij < NF_MAX*NF_MAX*3; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < NF_MAX*NF_MAX*3; ij+=blockDim_x){ gctr_smem[ij] = 0.0; } __syncthreads(); - constexpr int nfi_max = (AO_LMAX+2)*(AO_LMAX+3)/2; - constexpr int nfj_max = (AO_LMAX+1)*(AO_LMAX+2)/2; - __shared__ double buf[nfi_max*nfj_max]; - - type1_cart_kernel<1,0>( + type1_cart_kernel<1,0>(smem, buf, LI+1, LJ, ish, jsh, ksh, ecpbas, ecploc, @@ -340,7 +395,7 @@ void type1_cart_ip1_general(double *gctr, __syncthreads(); if (LI > 0){ - type1_cart_kernel<0,0>( + type1_cart_kernel<0,0>(smem, buf, LI-1, LJ, ish, jsh, ksh, ecpbas, ecploc, @@ -352,7 +407,7 @@ void type1_cart_ip1_general(double *gctr, const int nfi = (LI+1) * (LI+2) / 2; const int nfj = (LJ+1) * (LJ+2) / 2; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double *gx = gctr; diff --git a/gpu4pyscf/lib/ecp/ecp_type1_ipip.cu b/gpu4pyscf/lib/ecp/ecp_type1_ipip.cu index 3717e4824..3a05b7c1d 100644 --- a/gpu4pyscf/lib/ecp/ecp_type1_ipip.cu +++ b/gpu4pyscf/lib/ecp/ecp_type1_ipip.cu @@ -20,9 +20,30 @@ void type1_cart_ipipv(double *gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { + constexpr int nfi2_max = (AO_LMAX+3)*(AO_LMAX+4)/2; + constexpr int nfj_max = (AO_LMAX+1)*(AO_LMAX+2)/2; + constexpr int nfi1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; + +#ifdef USE_SYCL + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + double (&buf1)[nfi2_max*nfj_max] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&buf)[3*nfi1_max*nfj_max] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + __shared__ double buf1[nfi2_max*nfj_max]; + __shared__ double buf[3*nfi1_max*nfj_max]; + extern __shared__ double smem[]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -36,26 +57,21 @@ void type1_cart_ipipv(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += ioff*nao + joff + 9*ecp_id*nao*nao; - constexpr int nfi2_max = (AO_LMAX+3)*(AO_LMAX+4)/2; - constexpr int nfj_max = (AO_LMAX+1)*(AO_LMAX+2)/2; - __shared__ double buf1[nfi2_max*nfj_max]; - type1_cart_kernel<2,0>( + type1_cart_kernel<2,0>(smem, buf1, LI+2, LJ, ish, jsh, ksh, ecpbas, ecploc, atm, bas, env); __syncthreads(); - constexpr int nfi1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; - __shared__ double buf[3*nfi1_max*nfj_max]; - for (int i = threadIdx.x; i < 3*nfi1_max*nfj_max; i+=blockDim.x){ + for (int i = threadIdx_x; i < 3*nfi1_max*nfj_max; i+=blockDim_x){ buf[i] = 0.0; } __syncthreads(); _li_down(buf, buf1, LI+1, LJ); __syncthreads(); - type1_cart_kernel<1,0>( + type1_cart_kernel<1,0>(smem, buf1, LI, LJ, ish, jsh, ksh, ecpbas, ecploc, @@ -66,13 +82,13 @@ void type1_cart_ipipv(double *gctr, _li_down_and_write(gctr, buf, LI, LJ, nao); __syncthreads(); if (LI > 0){ - for (int i = threadIdx.x; i < 3*nfi1_max*nfj_max; i+=blockDim.x){ + for (int i = threadIdx_x; i < 3*nfi1_max*nfj_max; i+=blockDim_x){ buf[i] = 0.0; } __syncthreads(); _li_down(buf, buf1, LI-1, LJ); if (LI > 1){ - type1_cart_kernel<0,0>( + type1_cart_kernel<0,0>(smem, buf1, LI-2, LJ, ish, jsh, ksh, ecpbas, ecploc, @@ -93,9 +109,31 @@ void type1_cart_ipvip(double *gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { + + constexpr int nfi1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; + constexpr int nfj1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; + constexpr int nfi_max = (AO_LMAX+1)*(AO_LMAX+2)/2; + +#ifdef USE_SYCL + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + double (&buf1)[nfi1_max*nfj1_max] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&buf)[3*nfi_max*nfj1_max] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + __shared__ double buf1[nfi1_max*nfj1_max]; + __shared__ double buf[3*nfi_max*nfj1_max]; + extern __shared__ double smem[]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -109,26 +147,21 @@ void type1_cart_ipvip(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += ioff*nao + joff + 9*ecp_id*nao*nao; - constexpr int nfi1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; - constexpr int nfj1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; - __shared__ double buf1[nfi1_max*nfj1_max]; - type1_cart_kernel<1,1>( + type1_cart_kernel<1,1>(smem, buf1, LI+1, LJ+1, ish, jsh, ksh, ecpbas, ecploc, atm, bas, env); __syncthreads(); - constexpr int nfi_max = (AO_LMAX+1)*(AO_LMAX+2)/2; - __shared__ double buf[3*nfi_max*nfj1_max]; - for (int i = threadIdx.x; i < 3*nfi_max*nfj1_max; i+=blockDim.x){ + for (int i = threadIdx_x; i < 3*nfi_max*nfj1_max; i+=blockDim_x){ buf[i] = 0.0; } __syncthreads(); _li_down(buf, buf1, LI, LJ+1); __syncthreads(); if (LI > 0){ - type1_cart_kernel<0,1>( + type1_cart_kernel<0,1>(smem, buf1, LI-1, LJ+1, ish, jsh, ksh, ecpbas, ecploc, @@ -141,11 +174,11 @@ void type1_cart_ipvip(double *gctr, __syncthreads(); if (LJ > 0){ - for (int i = threadIdx.x; i < 3*nfi_max*nfj1_max; i+=blockDim.x){ + for (int i = threadIdx_x; i < 3*nfi_max*nfj1_max; i+=blockDim_x){ buf[i] = 0.0; } __syncthreads(); - type1_cart_kernel<1,0>( + type1_cart_kernel<1,0>(smem, buf1, LI+1, LJ-1, ish, jsh, ksh, ecpbas, ecploc, @@ -154,7 +187,7 @@ void type1_cart_ipvip(double *gctr, _li_down(buf, buf1, LI, LJ-1); __syncthreads(); if (LI > 0){ - type1_cart_kernel<0,0>( + type1_cart_kernel<0,0>(smem, buf1, LI-1, LJ-1, ish, jsh, ksh, ecpbas, ecploc, diff --git a/gpu4pyscf/lib/ecp/ecp_type2.cu b/gpu4pyscf/lib/ecp/ecp_type2.cu index 89c7c3098..dfa0190f3 100644 --- a/gpu4pyscf/lib/ecp/ecp_type2.cu +++ b/gpu4pyscf/lib/ecp/ecp_type2.cu @@ -17,9 +17,16 @@ template __device__ void type2_facs_rad(double* facs, const int LIC, const int np, const double rca, const double *ci, const double *ai){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + #else + const int threadIdx_x = threadIdx.x; + #endif + double root = 0.0; - if (threadIdx.x < NGAUSS){ - root = r128[threadIdx.x]; + if (threadIdx_x < NGAUSS){ + root = r128[threadIdx_x]; } const double r = root - rca; const double r2 = r*r; @@ -52,6 +59,14 @@ void type2_facs_rad(double* facs, const int LIC, const int np, const double rca, __device__ void type2_facs_omega(double* __restrict__ omega, const int LI, const int LC, double *r){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif double unitr[3]; if (r[0]*r[0] + r[1]*r[1] + r[2]*r[2] < 1e-16){ unitr[0] = 0; @@ -66,7 +81,7 @@ void type2_facs_omega(double* __restrict__ omega, const int LI, const int LC, do // LC + (i+j+k) + (LI + LC) needs to be even // When i+j+k + LC is even - for (int n = threadIdx.x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim_x){ const int i = n/(LI+1)/(LI+1); const int j = n/(LI+1)%(LI+1); const int k = n%(LI+1); @@ -90,7 +105,7 @@ void type2_facs_omega(double* __restrict__ omega, const int LI, const int LC, do } // When i+j+k + LC is odd - for (int n = threadIdx.x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim_x){ const int i = n/(LI+1)/(LI+1); const int j = n/(LI+1)%(LI+1); const int k = n%(LI+1); @@ -115,6 +130,14 @@ void type2_facs_omega(double* __restrict__ omega, const int LI, const int LC, do template __device__ void type2_facs_omega(double* __restrict__ omega, double *r){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif double unitr[3]; if (r[0]*r[0] + r[1]*r[1] + r[2]*r[2] < 1e-16){ unitr[0] = 0; @@ -129,7 +152,7 @@ void type2_facs_omega(double* __restrict__ omega, double *r){ // LC + (i+j+k) + (LI + LC) needs to be even // When i+j+k + LC is even - for (int n = threadIdx.x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim_x){ const int i = n/(LI+1)/(LI+1); const int j = n/(LI+1)%(LI+1); const int k = n%(LI+1); @@ -153,7 +176,7 @@ void type2_facs_omega(double* __restrict__ omega, double *r){ } // When i+j+k + LC is odd - for (int n = threadIdx.x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim.x){ + for (int n = threadIdx_x; n < (LI+1)*(LI+1)*(LI+1); n+=blockDim_x){ const int i = n/(LI+1)/(LI+1); const int j = n/(LI+1)%(LI+1); const int k = n%(LI+1); @@ -178,6 +201,15 @@ void type2_facs_omega(double* __restrict__ omega, double *r){ __device__ void type2_ang(double * __restrict__ facs, const int LI, const int LC, double *rca, double *omega){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif + const int LI1 = LI+1; const int nfi = LI1*(LI1+1)/2; const int LCC1 = (2*LC+1); @@ -189,7 +221,7 @@ void type2_ang(double * __restrict__ facs, const int LI, const int LC, double *r cache_fac(fi, LI, rca); // i,j,k,ijkmn->(i+j+k)pmn - for (int pmn = threadIdx.x; pmn < nfi*LIC1; pmn+=blockDim.x){ + for (int pmn = threadIdx_x; pmn < nfi*LIC1; pmn+=blockDim_x){ const int m = pmn/nfi; const int p = pmn%nfi; @@ -228,6 +260,14 @@ void type2_ang(double * __restrict__ facs, const int LI, const int LC, double *r template __device__ void type2_ang(double * __restrict__ facs, double *rca, double *omega){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + #endif constexpr int LI1 = LI+1; constexpr int nfi = LI1*(LI1+1)/2; constexpr int LCC1 = (2*LC+1); @@ -239,7 +279,7 @@ void type2_ang(double * __restrict__ facs, double *rca, double *omega){ cache_fac
  • (fi, rca); // i,j,k,ijkmn->(i+j+k)pmn - for (int pmn = threadIdx.x; pmn < nfi*LIC1; pmn+=blockDim.x){ + for (int pmn = threadIdx_x; pmn < nfi*LIC1; pmn+=blockDim_x){ const int m = pmn/nfi; const int p = pmn%nfi; @@ -283,7 +323,43 @@ void type2_cart(double * __restrict__ gctr, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { + constexpr int LI1 = LI+1; + constexpr int LJ1 = LJ+1; + constexpr int LIC1 = LI+LC+1; + constexpr int LJC1 = LJ+LC+1; + constexpr int LCC1 = (2*LC+1); + + constexpr int BLKI = (LIC1+1)/2 * LCC1; + constexpr int BLKJ = (LJC1+1)/2 * LCC1; + + constexpr int nfi = (LI+1) * (LI+2) / 2; + constexpr int nfj = (LJ+1) * (LJ+2) / 2; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + + auto thread_block = item.get_group(); + const int task_id = thread_block.get_group_id(0); + using tile_t1 = double[LI1*(LI1+1)*(LI1+2)/6 * BLKI]; + tile_t1& omegai = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + using tile_t2 = double[LJ1*(LJ1+1)*(LJ1+2)/6 * BLKJ]; + tile_t2& omegaj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + using tile_t3 = double[(LI+LJ+1) * LIC1 * LJC1]; + tile_t3& rad_all = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&angi)[LI1*nfi*LIC1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&angj)[LJ1*nfj*LJC1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + + __shared__ double omegai[LI1*(LI1+1)*(LI1+2)/6 * BLKI]; // up to 12600 Bytes + __shared__ double omegaj[LJ1*(LJ1+1)*(LJ1+2)/6 * BLKJ]; + __shared__ double rad_all[(LI+LJ+1) * LIC1 * LJC1]; + __shared__ double angi[LI1*nfi*LIC1]; + __shared__ double angj[LJ1*nfj*LJC1]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -306,18 +382,6 @@ void type2_cart(double * __restrict__ gctr, rcb[1] = rc[1] - rj[1]; rcb[2] = rc[2] - rj[2]; - constexpr int LI1 = LI+1; - constexpr int LJ1 = LJ+1; - constexpr int LIC1 = LI+LC+1; - constexpr int LJC1 = LJ+LC+1; - constexpr int LCC1 = (2*LC+1); - - constexpr int BLKI = (LIC1+1)/2 * LCC1; - constexpr int BLKJ = (LJC1+1)/2 * LCC1; - - __shared__ double omegai[LI1*(LI1+1)*(LI1+2)/6 * BLKI]; - __shared__ double omegaj[LJ1*(LJ1+1)*(LJ1+2)/6 * BLKJ]; - type2_facs_omega(omegai, rca); type2_facs_omega(omegaj, rcb); __syncthreads(); @@ -337,7 +401,6 @@ void type2_cart(double * __restrict__ gctr, double radj[LJC1]; type2_facs_rad<0>(radj, LJ+LC, npj, dcb, cj, aj); - __shared__ double rad_all[(LI+LJ+1) * LIC1 * LJC1]; set_shared_memory(rad_all, (LI+LJ+1)*LIC1*LJC1); double ur = 0.0; @@ -347,8 +410,8 @@ void type2_cart(double * __restrict__ gctr, } double root = 0.0; - if (threadIdx.x < NGAUSS){ - root = r128[threadIdx.x]; + if (threadIdx_x < NGAUSS){ + root = r128[threadIdx_x]; } for (int p = 0; p <= LI+LJ; p++){ double *prad = rad_all + p*LIC1*LJC1; @@ -360,12 +423,6 @@ void type2_cart(double * __restrict__ gctr, } __syncthreads(); - constexpr int nfi = (LI+1) * (LI+2) / 2; - constexpr int nfj = (LJ+1) * (LJ+2) / 2; - - __shared__ double angi[LI1*nfi*LIC1]; - __shared__ double angj[LJ1*nfj*LJC1]; - const double fac = 16.0 * M_PI * M_PI * _common_fac[LI] * _common_fac[LJ]; constexpr int nreg = (nfi*nfj + THREADS - 1)/THREADS; @@ -379,7 +436,7 @@ void type2_cart(double * __restrict__ gctr, type2_ang(angi, rca, omegai+m); type2_ang(angj, rcb, omegaj+m); __syncthreads(); - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double s = 0; @@ -404,7 +461,7 @@ void type2_cart(double * __restrict__ gctr, const int ioff = ao_loc[ish]; const int joff = ao_loc[jsh]; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double tmp = reg_gctr[ij/THREADS]; @@ -422,9 +479,22 @@ void type2_cart(double * __restrict__ gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { + #ifdef USE_SYCL + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + #else const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + extern __shared__ double smem[]; + #endif if (task_id >= ntasks){ return; } @@ -433,8 +503,6 @@ void type2_cart(double * __restrict__ gctr, const int jsh = tasks[task_id + ntasks]; const int ksh = tasks[task_id + 2*ntasks]; - extern __shared__ double smem[]; - const double *ri = env + atm[PTR_COORD+bas[ATOM_OF+ish*BAS_SLOTS]*ATM_SLOTS]; const double *rj = env + atm[PTR_COORD+bas[ATOM_OF+jsh*BAS_SLOTS]*ATM_SLOTS]; @@ -478,8 +546,8 @@ void type2_cart(double * __restrict__ gctr, type2_facs_rad<0>(radj, LJ+LC, npj, dcb, cj, aj); double root = 0.0; - if (threadIdx.x < NGAUSS){ - root = r128[threadIdx.x]; + if (threadIdx_x < NGAUSS){ + root = r128[threadIdx_x]; } double* rad_all = smem; set_shared_memory(rad_all, (LI+LJ+1)*(LI+LC+1)*(LJ+LC+1)); @@ -511,7 +579,7 @@ void type2_cart(double * __restrict__ gctr, type2_ang(angj, LJ, LC, rcb, omegaj+m); __syncthreads(); - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double s = 0; @@ -540,7 +608,7 @@ void type2_cart(double * __restrict__ gctr, const int joff = ao_loc[jsh]; double *gctr_ij = gctr + ioff + joff*nao; double *gctr_ji = gctr + joff + ioff*nao; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double tmp = reg_gctr[ij/THREADS]; diff --git a/gpu4pyscf/lib/ecp/ecp_type2_ip.cu b/gpu4pyscf/lib/ecp/ecp_type2_ip.cu index f0d3c5eea..868809d91 100644 --- a/gpu4pyscf/lib/ecp/ecp_type2_ip.cu +++ b/gpu4pyscf/lib/ecp/ecp_type2_ip.cu @@ -14,13 +14,46 @@ * limitations under the License. */ - template __device__ void type2_cart_unrolled_kernel(double *gctr, const int ish, const int jsh, const int ksh, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { + constexpr int LI1 = LI+1; + constexpr int LJ1 = LJ+1; + constexpr int LIC1 = LI+LC+1; + constexpr int LJC1 = LJ+LC+1; + constexpr int LCC1 = (2*LC+1); + + constexpr int BLKI = (LIC1+1)/2 * LCC1; + constexpr int BLKJ = (LJC1+1)/2 * LCC1; + + constexpr int nfi = (LI+1) * (LI+2) / 2; + constexpr int nfj = (LJ+1) * (LJ+2) / 2; + +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + + auto thread_block = item.get_group(); + double (&omegai)[LI1*(LI1+1)*(LI1+2)/6 * BLKI] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&omegaj)[LJ1*(LJ1+1)*(LJ1+2)/6 * BLKJ] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rad_all)[(LI+LJ+1)*LIC1*LJC1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&angi)[LI1*nfi*LIC1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&angj)[LJ1*nfj*LJC1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else // USE_SYCL + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + + __shared__ double omegai[LI1*(LI1+1)*(LI1+2)/6 * BLKI]; + __shared__ double omegaj[LJ1*(LJ1+1)*(LJ1+2)/6 * BLKJ]; + __shared__ double rad_all[(LI+LJ+1)*LIC1*LJC1]; + __shared__ double angi[LI1*nfi*LIC1]; + __shared__ double angj[LJ1*nfj*LJC1]; +#endif // USE_SYCL + const double *ri = env + atm[PTR_COORD+bas[ATOM_OF+ish*BAS_SLOTS]*ATM_SLOTS]; const double *rj = env + atm[PTR_COORD+bas[ATOM_OF+jsh*BAS_SLOTS]*ATM_SLOTS]; @@ -35,23 +68,10 @@ void type2_cart_unrolled_kernel(double *gctr, rcb[1] = rc[1] - rj[1]; rcb[2] = rc[2] - rj[2]; - constexpr int LI1 = LI+1; - constexpr int LJ1 = LJ+1; - constexpr int LIC1 = LI+LC+1; - constexpr int LJC1 = LJ+LC+1; - constexpr int LCC1 = (2*LC+1); - - constexpr int BLKI = (LIC1+1)/2 * LCC1; - constexpr int BLKJ = (LJC1+1)/2 * LCC1; - - __shared__ double omegai[LI1*(LI1+1)*(LI1+2)/6 * BLKI]; - __shared__ double omegaj[LJ1*(LJ1+1)*(LJ1+2)/6 * BLKJ]; - type2_facs_omega(omegai, rca); type2_facs_omega(omegaj, rcb); __syncthreads(); - __shared__ double rad_all[(LI+LJ+1)*LIC1*LJC1]; set_shared_memory(rad_all, (LI+LJ+1)*LIC1*LJC1); const int npi = bas[NPRIM_OF+ish*BAS_SLOTS]; @@ -80,20 +100,14 @@ void type2_cart_unrolled_kernel(double *gctr, for (int j = 0; j <= LJ+LC; j++){ block_reduce(radi[i]*radj[j]*ur_tmp, prad+i*(LJ+LC+1)+j); }} - const int ir = threadIdx.x; + const int ir = threadIdx_x; ur_tmp *= r128[ir]; } __syncthreads(); - constexpr int nfi = (LI+1) * (LI+2) / 2; - constexpr int nfj = (LJ+1) * (LJ+2) / 2; - - __shared__ double angi[LI1*nfi*LIC1]; - __shared__ double angj[LJ1*nfj*LJC1]; - const double fac = 16.0 * M_PI * M_PI * _common_fac[LI] * _common_fac[LJ]; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ gctr[ij] = 0.0; } @@ -102,7 +116,7 @@ void type2_cart_unrolled_kernel(double *gctr, type2_ang(angi, rca, omegai+m); type2_ang(angj, rcb, omegaj+m); __syncthreads(); - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double s = 0; @@ -128,14 +142,26 @@ void type2_cart_unrolled_kernel(double *gctr, } } +// `smem` is passed in rather than declared here as `extern __shared__`: +// under SYCL the dynamic local memory comes from a sycl::local_accessor +// created by the enclosing submit(), which cannot be reached from a +// device function. The CUDA callers pass their own `extern __shared__` +// block, so behaviour is unchanged. template __device__ -void type2_cart_kernel(double *gctr, +void type2_cart_kernel(double *smem, double *gctr, const int LI, const int LJ, const int LC, const int ish, const int jsh, const int ksh, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { - extern __shared__ double smem[]; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); +#else // USE_SYCL + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; +#endif // USE_SYCL const double *ri = env + atm[PTR_COORD+bas[ATOM_OF+ish*BAS_SLOTS]*ATM_SLOTS]; const double *rj = env + atm[PTR_COORD+bas[ATOM_OF+jsh*BAS_SLOTS]*ATM_SLOTS]; @@ -188,7 +214,7 @@ void type2_cart_kernel(double *gctr, for (int j = 0; j <= LJ+LC; j++){ block_reduce(radi[i]*radj[j]*ur_tmp, prad+i*(LJ+LC+1)+j); }} - const int ir = threadIdx.x; + const int ir = threadIdx_x; ur_tmp *= r128[ir]; } __syncthreads(); @@ -200,7 +226,7 @@ void type2_cart_kernel(double *gctr, const double fac = 16.0 * M_PI * M_PI * _common_fac[LI] * _common_fac[LJ]; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ gctr[ij] = 0.0; } __syncthreads(); @@ -210,7 +236,7 @@ void type2_cart_kernel(double *gctr, type2_ang(angi, LI, LC, rca, omegai+m); type2_ang(angj, LJ, LC, rcb, omegaj+m); __syncthreads(); - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double s = 0; @@ -246,7 +272,24 @@ void type2_cart_ip1(double *gctr, const int *ecpbas, const int *ecploc, const int *atm, const int *bas, const double *env) { + constexpr int nfi = (LI+1) * (LI+2) / 2; + constexpr int nfj = (LJ+1) * (LJ+2) / 2; + constexpr int nfi1 = (LI+2) * (LI+3)/2; + +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + double (&gctr_smem)[nfi*nfj*3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&buf)[nfi1*nfj] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + __shared__ double gctr_smem[nfi*nfj*3]; + __shared__ double buf[nfi1*nfj]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -259,16 +302,11 @@ void type2_cart_ip1(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += 3*ecp_id*nao*nao + ioff*nao + joff; - constexpr int nfi = (LI+1) * (LI+2) / 2; - constexpr int nfj = (LJ+1) * (LJ+2) / 2; - __shared__ double gctr_smem[nfi*nfj*3]; - for (int ij = threadIdx.x; ij < nfi*nfj*3; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj*3; ij+=blockDim_x){ gctr_smem[ij] = 0.0; } __syncthreads(); - constexpr int nfi1 = (LI+2) * (LI+3)/2; - __shared__ double buf[nfi1*nfj]; type2_cart_unrolled_kernel<1,0,LI+1,LJ,LC>( buf, ish, jsh, ksh, ecpbas, ecploc, @@ -284,7 +322,7 @@ void type2_cart_ip1(double *gctr, __syncthreads(); } - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double *gx = gctr; @@ -304,9 +342,29 @@ void type2_cart_ip1_general(double *gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { + constexpr int NFI_MAX = (AO_LMAX+2)*(AO_LMAX+3)/2; + constexpr int NFJ_MAX = (AO_LMAX+1)*(AO_LMAX+2)/2; + +#ifdef USE_SYCL + const int task_id = item.get_group(0); + const int threadIdx_x = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + double (&gctr_smem)[NF_MAX*NF_MAX*3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&buf)[NFI_MAX*NFJ_MAX] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; + __shared__ double gctr_smem[NF_MAX*NF_MAX*3]; + __shared__ double buf[NFI_MAX*NFJ_MAX]; + extern __shared__ double smem[]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -319,16 +377,12 @@ void type2_cart_ip1_general(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += 3*ecp_id*nao*nao + ioff*nao + joff; - __shared__ double gctr_smem[NF_MAX*NF_MAX*3]; - for (int ij = threadIdx.x; ij < NF_MAX*NF_MAX*3; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < NF_MAX*NF_MAX*3; ij+=blockDim_x){ gctr_smem[ij] = 0.0; } __syncthreads(); - constexpr int NFI_MAX = (AO_LMAX+2)*(AO_LMAX+3)/2; - constexpr int NFJ_MAX = (AO_LMAX+1)*(AO_LMAX+2)/2; - __shared__ double buf[NFI_MAX*NFJ_MAX]; - type2_cart_kernel<1,0>( + type2_cart_kernel<1,0>(smem, buf, LI+1, LJ, LC, ish, jsh, ksh, ecpbas, ecploc, @@ -336,7 +390,7 @@ void type2_cart_ip1_general(double *gctr, _li_down(gctr_smem, buf, LI, LJ); __syncthreads(); if (LI > 0){ - type2_cart_kernel<0,0>( + type2_cart_kernel<0,0>(smem, buf, LI-1, LJ, LC, ish, jsh, ksh, ecpbas, ecploc, @@ -347,7 +401,7 @@ void type2_cart_ip1_general(double *gctr, const int nfi = (LI+1) * (LI+2) / 2; const int nfj = (LJ+1) * (LJ+2) / 2; - for (int ij = threadIdx.x; ij < nfi*nfj; ij+=blockDim.x){ + for (int ij = threadIdx_x; ij < nfi*nfj; ij+=blockDim_x){ const int i = ij%nfi; const int j = ij/nfi; double *gx = gctr; diff --git a/gpu4pyscf/lib/ecp/ecp_type2_ipip.cu b/gpu4pyscf/lib/ecp/ecp_type2_ipip.cu index 1164786e7..8a9fc56ad 100644 --- a/gpu4pyscf/lib/ecp/ecp_type2_ipip.cu +++ b/gpu4pyscf/lib/ecp/ecp_type2_ipip.cu @@ -21,9 +21,23 @@ void type2_cart_ipipv(double *gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { + constexpr int nfi2_max = (AO_LMAX+3)*(AO_LMAX+4)/2; + constexpr int nfj_max = (AO_LMAX+1)*(AO_LMAX+2)/2; + +#ifdef USE_SYCL + const int task_id = item.get_group(0); + double (&buf1)[nfi2_max*nfj_max] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + __shared__ double buf1[nfi2_max*nfj_max]; + extern __shared__ double smem[]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -37,17 +51,13 @@ void type2_cart_ipipv(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += ioff*nao + joff + 9*ecp_id*nao*nao; - constexpr int nfi2_max = (AO_LMAX+3)*(AO_LMAX+4)/2; - constexpr int nfj_max = (AO_LMAX+1)*(AO_LMAX+2)/2; - __shared__ double buf1[nfi2_max*nfj_max]; - type2_cart_kernel<2,0>( + type2_cart_kernel<2,0>(smem, buf1, LI+2, LJ, LC, ish, jsh, ksh, ecpbas, ecploc, atm, bas, env); constexpr int nfi1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; - extern __shared__ double smem[]; double *buf = smem; set_shared_memory(buf, 3*nfi1_max*nfj_max); _li_down(buf, buf1, LI+1, LJ); @@ -55,7 +65,7 @@ void type2_cart_ipipv(double *gctr, _li_down_and_write(gctr, buf, LI, LJ, nao); __syncthreads(); - type2_cart_kernel<1,0>( + type2_cart_kernel<1,0>(smem, buf1, LI, LJ, LC, ish, jsh, ksh, ecpbas, ecploc, @@ -73,7 +83,7 @@ void type2_cart_ipipv(double *gctr, _li_up_and_write(gctr, buf, LI, LJ, nao); __syncthreads(); if (LI > 1){ - type2_cart_kernel<0,0>( + type2_cart_kernel<0,0>(smem, buf1, LI-2, LJ, LC, ish, jsh, ksh, ecpbas, ecploc, @@ -93,9 +103,23 @@ void type2_cart_ipvip(double *gctr, const int *ao_loc, const int nao, const int *tasks, const int ntasks, const int *ecpbas, const int *ecploc, - const int *atm, const int *bas, const double *env) + const int *atm, const int *bas, const double *env +#ifdef USE_SYCL + , sycl::nd_item<1> &item, double* smem +#endif + ) { + constexpr int nfi1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; + constexpr int nfj1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; + +#ifdef USE_SYCL + const int task_id = item.get_group(0); + double (&buf1)[nfi1_max*nfj1_max] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else // USE_SYCL const int task_id = blockIdx.x; + __shared__ double buf1[nfi1_max*nfj1_max]; + extern __shared__ double smem[]; +#endif // USE_SYCL if (task_id >= ntasks){ return; } @@ -109,17 +133,13 @@ void type2_cart_ipvip(double *gctr, const int ecp_id = ecpbas[ECP_ATOM_ID+ecploc[ksh]*BAS_SLOTS]; gctr += ioff*nao + joff + 9*ecp_id*nao*nao; - constexpr int nfi1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; - constexpr int nfj1_max = (AO_LMAX+2)*(AO_LMAX+3)/2; - __shared__ double buf1[nfi1_max*nfj1_max]; - type2_cart_kernel<1,1>( + type2_cart_kernel<1,1>(smem, buf1, LI+1, LJ+1, LC, ish, jsh, ksh, ecpbas, ecploc, atm, bas, env); constexpr int nfi_max = (AO_LMAX+1)*(AO_LMAX+2)/2; - extern __shared__ double smem[]; double *buf = smem; set_shared_memory(buf, 3*nfi_max*nfj1_max); _li_down(buf, buf1, LI, LJ+1); @@ -127,7 +147,7 @@ void type2_cart_ipvip(double *gctr, _lj_down_and_write(gctr, buf, LI, LJ, nao); __syncthreads(); if (LI > 0){ - type2_cart_kernel<0,1>( + type2_cart_kernel<0,1>(smem, buf1, LI-1, LJ+1, LC, ish, jsh, ksh, ecpbas, ecploc, @@ -140,7 +160,7 @@ void type2_cart_ipvip(double *gctr, } if (LJ > 0){ - type2_cart_kernel<1,0>( + type2_cart_kernel<1,0>(smem, buf1, LI+1, LJ-1, LC, ish, jsh, ksh, ecpbas, ecploc, @@ -151,7 +171,7 @@ void type2_cart_ipvip(double *gctr, _lj_up_and_write(gctr, buf, LI, LJ, nao); __syncthreads(); if (LI > 0){ - type2_cart_kernel<0,0>( + type2_cart_kernel<0,0>(smem, buf1, LI-1, LJ-1, LC, ish, jsh, ksh, ecpbas, ecploc, diff --git a/gpu4pyscf/lib/ecp/nr_ecp_driver.cu b/gpu4pyscf/lib/ecp/nr_ecp_driver.cu index c68b982b4..e8a713e0a 100644 --- a/gpu4pyscf/lib/ecp/nr_ecp_driver.cu +++ b/gpu4pyscf/lib/ecp/nr_ecp_driver.cu @@ -29,6 +29,70 @@ #include "ecp_type1_ipip.cu" #include "ecp_type2_ipip.cu" +#define ECP_ARGS gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env + +#ifdef USE_SYCL +#define ECP_LAUNCH1(TAG, KPREFIX, LI, LJ) \ + stream.parallel_for( \ + sycl::nd_range<1>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + KPREFIX(ECP_ARGS); \ + }) +#else +#define ECP_LAUNCH1(TAG, KPREFIX, LI, LJ) \ + KPREFIX<<>>(ECP_ARGS) +#endif + +#ifdef USE_SYCL +#define ECP_LAUNCH2(TAG, KPREFIX, LI, LJ, LC) \ + stream.parallel_for( \ + sycl::nd_range<1>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + KPREFIX(ECP_ARGS); \ + }) +#else +#define ECP_LAUNCH2(TAG, KPREFIX, LI, LJ, LC) \ + KPREFIX<<>>(ECP_ARGS) +#endif + +#ifdef USE_SYCL +#define ECP_LAUNCH_SMEM(TAG, SMEM, KFUNC, ...) \ + stream.submit([&](sycl::handler &cgh) { \ + sycl::local_accessor local_acc(sycl::range<1>(SMEM), cgh); \ + cgh.parallel_for( \ + sycl::nd_range<1>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + KFUNC(__VA_ARGS__, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); \ + }); \ + }) +#else +#define ECP_LAUNCH_SMEM(TAG, SMEM, KFUNC, ...) \ + KFUNC<<>>(__VA_ARGS__) +#endif + +#ifdef USE_SYCL +#define ECP_LAUNCH_GENERAL(TAG, SMEM, KFUNC, ...) \ + stream.submit([&](sycl::handler &cgh) { \ + sycl::local_accessor local_acc(sycl::range<1>(SMEM), cgh); \ + cgh.parallel_for( \ + sycl::nd_range<1>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + KFUNC(__VA_ARGS__, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); \ + }); \ + }) +#else +#define ECP_LAUNCH_GENERAL(TAG, SMEM, KFUNC, ...) do { \ + cudaError_t _e = cudaFuncSetAttribute( \ + KFUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, (SMEM)*sizeof(double)); \ + if (_e != cudaSuccess) { \ + fprintf(stderr, "CUDA Error in cudaFuncSetAttribute %s: %s\n", \ + __func__, cudaGetErrorString(_e)); \ + return 1; \ + } \ + KFUNC<<>>(__VA_ARGS__); \ +} while(0) +#endif + extern "C" { int ECP_cart(double *gctr, const int *ao_loc, const int nao, @@ -37,26 +101,31 @@ int ECP_cart(double *gctr, const int *atm, const int *bas, const double *env, const int li, const int lj, const int lc){ // one task per thread block +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(ntasks); + sycl::queue& stream = *sycl_get_queue(); +#else dim3 threads(THREADS); dim3 blocks(ntasks); +#endif if (lc >= 0){ int task_type = li * 100 + lj * 10 + lc; switch (task_type) { - case 0: type2_cart<0,0,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 1: type2_cart<0,0,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 2: type2_cart<0,0,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 3: type2_cart<0,0,3><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 10: type2_cart<0,1,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 11: type2_cart<0,1,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 12: type2_cart<0,1,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 110: type2_cart<1,1,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 111: type2_cart<1,1,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 112: type2_cart<1,1,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 20: type2_cart<0,2,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 21: type2_cart<0,2,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 30: type2_cart<0,3,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 120: type2_cart<1,2,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - + case 0: ECP_LAUNCH2(type2_cart_000, type2_cart, 0,0,0); break; + case 1: ECP_LAUNCH2(type2_cart_001, type2_cart, 0,0,1); break; + case 2: ECP_LAUNCH2(type2_cart_002, type2_cart, 0,0,2); break; + case 3: ECP_LAUNCH2(type2_cart_003, type2_cart, 0,0,3); break; + case 10: ECP_LAUNCH2(type2_cart_010, type2_cart, 0,1,0); break; + case 11: ECP_LAUNCH2(type2_cart_011, type2_cart, 0,1,1); break; + case 12: ECP_LAUNCH2(type2_cart_012, type2_cart, 0,1,2); break; + case 110: ECP_LAUNCH2(type2_cart_110, type2_cart, 1,1,0); break; + case 111: ECP_LAUNCH2(type2_cart_111, type2_cart, 1,1,1); break; + case 112: ECP_LAUNCH2(type2_cart_112, type2_cart, 1,1,2); break; + case 20: ECP_LAUNCH2(type2_cart_020, type2_cart, 0,2,0); break; + case 21: ECP_LAUNCH2(type2_cart_021, type2_cart, 0,2,1); break; + case 30: ECP_LAUNCH2(type2_cart_030, type2_cart, 0,3,0); break; + case 120: ECP_LAUNCH2(type2_cart_120, type2_cart, 1,2,0); break; // General kernel default: { const int li1 = li+1; @@ -76,40 +145,31 @@ int ECP_cart(double *gctr, int smem_size4 = lj1*nfj*ljc1; // angj int smem_size = smem_size0 + smem_size1 + smem_size2 + smem_size3 + smem_size4; - type2_cart<<>>( - gctr, - li, lj, lc, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + ECP_LAUNCH_SMEM(type2_cart_sycl, smem_size, type2_cart, + gctr, li, lj, lc, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); }} } else { int task_type = li * 10 + lj; switch (task_type) { - case 0: type1_cart<0,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 1: type1_cart<0,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 11: type1_cart<1,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 2: type1_cart<0,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 3: type1_cart<0,3><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 12: type1_cart<1,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 4: type1_cart<0,4><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 13: type1_cart<1,3><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 22: type1_cart<2,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; + case 0: ECP_LAUNCH1(type1_cart_00, type1_cart, 0,0); break; + case 1: ECP_LAUNCH1(type1_cart_01, type1_cart, 0,1); break; + case 11: ECP_LAUNCH1(type1_cart_11, type1_cart, 1,1); break; + case 2: ECP_LAUNCH1(type1_cart_02, type1_cart, 0,2); break; + case 3: ECP_LAUNCH1(type1_cart_03, type1_cart, 0,3); break; + case 12: ECP_LAUNCH1(type1_cart_12, type1_cart, 1,2); break; + case 4: ECP_LAUNCH1(type1_cart_04, type1_cart, 0,4); break; + case 13: ECP_LAUNCH1(type1_cart_13, type1_cart, 1,3); break; + case 22: ECP_LAUNCH1(type1_cart_22, type1_cart, 2,2); break; default: { const int lij1 = li+lj+1; const int lij3 = lij1*lij1*lij1; + int smem_size = lij3 + lij1*lij1; - int smem_size = 0; - smem_size += lij3; // rad_ang - smem_size += lij1*lij1; // rad_all - type1_cart<<>>( - gctr, li, lj, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + ECP_LAUNCH_SMEM(type1_cart_kernel, smem_size, type1_cart, + gctr, li, lj, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); } } } @@ -128,52 +188,53 @@ int ECP_ip_cart(double *gctr, const int *atm, const int *bas, const double *env, const int li, const int lj, const int lc){ // one task per thread block +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(ntasks); + sycl::queue& stream = *sycl_get_queue(); +#else dim3 threads(THREADS); dim3 blocks(ntasks); +#endif if (lc < 0){ int task_type = li * 10 + lj; switch (task_type) { - case 0: type1_cart_ip1<0,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 1: type1_cart_ip1<0,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 11: type1_cart_ip1<1,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 2: type1_cart_ip1<0,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 3: type1_cart_ip1<0,3><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 12: type1_cart_ip1<1,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 4: type1_cart_ip1<0,4><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 13: type1_cart_ip1<1,3><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 22: type1_cart_ip1<2,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; + case 0: ECP_LAUNCH1(type1_cart_ip1_00, type1_cart_ip1, 0,0); break; + case 1: ECP_LAUNCH1(type1_cart_ip1_01, type1_cart_ip1, 0,1); break; + case 11: ECP_LAUNCH1(type1_cart_ip1_11, type1_cart_ip1, 1,1); break; + case 2: ECP_LAUNCH1(type1_cart_ip1_02, type1_cart_ip1, 0,2); break; + case 3: ECP_LAUNCH1(type1_cart_ip1_03, type1_cart_ip1, 0,3); break; + case 12: ECP_LAUNCH1(type1_cart_ip1_12, type1_cart_ip1, 1,2); break; + case 4: ECP_LAUNCH1(type1_cart_ip1_04, type1_cart_ip1, 0,4); break; + case 13: ECP_LAUNCH1(type1_cart_ip1_13, type1_cart_ip1, 1,3); break; + case 22: ECP_LAUNCH1(type1_cart_ip1_22, type1_cart_ip1, 2,2); break; default: { const int lij1 = li+lj+2; const int lij3 = lij1*lij1*lij1; + int smem_size = lij3 + lij1*lij1; - int smem_size = 0; - smem_size += lij3; // rad_ang - smem_size += lij1*lij1; // rad_all - type1_cart_ip1_general<<>>( - gctr, li, lj, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + ECP_LAUNCH_GENERAL(type1_cart_ip1_general_kernel, smem_size, + type1_cart_ip1_general, + gctr, li, lj, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); }} } else { int task_type = li * 100 + lj * 10 + lc; switch (task_type) { - case 0: type2_cart_ip1<0,0,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 1: type2_cart_ip1<0,0,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 2: type2_cart_ip1<0,0,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 3: type2_cart_ip1<0,0,3><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 10: type2_cart_ip1<0,1,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 11: type2_cart_ip1<0,1,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 12: type2_cart_ip1<0,1,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 110: type2_cart_ip1<1,1,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 111: type2_cart_ip1<1,1,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 112: type2_cart_ip1<1,1,2><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 20: type2_cart_ip1<0,2,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 21: type2_cart_ip1<0,2,1><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 30: type2_cart_ip1<0,3,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - case 120: type2_cart_ip1<1,2,0><<>>(gctr, ao_loc, nao, tasks, ntasks, ecpbas, ecploc, atm, bas, env); break; - + case 0: ECP_LAUNCH2(type2_cart_ip1_000, type2_cart_ip1, 0,0,0); break; + case 1: ECP_LAUNCH2(type2_cart_ip1_001, type2_cart_ip1, 0,0,1); break; + case 2: ECP_LAUNCH2(type2_cart_ip1_002, type2_cart_ip1, 0,0,2); break; + case 3: ECP_LAUNCH2(type2_cart_ip1_003, type2_cart_ip1, 0,0,3); break; + case 10: ECP_LAUNCH2(type2_cart_ip1_010, type2_cart_ip1, 0,1,0); break; + case 11: ECP_LAUNCH2(type2_cart_ip1_011, type2_cart_ip1, 0,1,1); break; + case 12: ECP_LAUNCH2(type2_cart_ip1_012, type2_cart_ip1, 0,1,2); break; + case 110: ECP_LAUNCH2(type2_cart_ip1_110, type2_cart_ip1, 1,1,0); break; + case 111: ECP_LAUNCH2(type2_cart_ip1_111, type2_cart_ip1, 1,1,1); break; + case 112: ECP_LAUNCH2(type2_cart_ip1_112, type2_cart_ip1, 1,1,2); break; + case 20: ECP_LAUNCH2(type2_cart_ip1_020, type2_cart_ip1, 0,2,0); break; + case 21: ECP_LAUNCH2(type2_cart_ip1_021, type2_cart_ip1, 0,2,1); break; + case 30: ECP_LAUNCH2(type2_cart_ip1_030, type2_cart_ip1, 0,3,0); break; + case 120: ECP_LAUNCH2(type2_cart_ip1_120, type2_cart_ip1, 1,2,0); break; // General kernel default: { const int li1 = li+2; @@ -192,24 +253,12 @@ int ECP_ip_cart(double *gctr, int smem_size2 = lj1*(lj1+1)*(lj1+2)/6 * blkj; // omegaj int smem_size3 = li1*lic1*nfi; // angi int smem_size4 = lj1*ljc1*nfj; // angj - int dynamic_smem_size = smem_size0 + smem_size1 + smem_size2 + smem_size3 + smem_size4; - cudaError_t err = cudaFuncSetAttribute( - type2_cart_ip1_general, - cudaFuncAttributeMaxDynamicSharedMemorySize, - dynamic_smem_size*sizeof(double)); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in cudaFuncSetAttribute %s (li,lj,lc = %d,%d,%d): %s\n", __func__, li,lj,lc, cudaGetErrorString(err)); - return 1; - } - - type2_cart_ip1_general<<>>( - gctr, li, lj, lc, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + ECP_LAUNCH_GENERAL(type2_cart_ip1_general_kernel, dynamic_smem_size, + type2_cart_ip1_general, + gctr, li, lj, lc, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); }} } cudaError_t err = cudaGetLastError(); @@ -227,23 +276,24 @@ int ECP_ipipv_cart(double *gctr, const int *atm, const int *bas, const double *env, const int li, const int lj, const int lc){ // one task per thread block +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(ntasks); + sycl::queue& stream = *sycl_get_queue(); +#else dim3 threads(THREADS); dim3 blocks(ntasks); +#endif if (lc < 0){ - const int lij1 = li+lj+3; // + const int lij1 = li+lj+3; const int lij3 = lij1*lij1*lij1; - int smem_size = 0; - smem_size += lij3; // rad_ang - smem_size += lij1*lij1; // rad_all - type1_cart_ipipv<<>>( - gctr, li, lj, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + int smem_size = lij3 + lij1*lij1; + ECP_LAUNCH_GENERAL(type1_cart_ipipv_kernel, smem_size, type1_cart_ipipv, + gctr, li, lj, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); } else { const int li1 = li+3; const int lj1 = lj+1; @@ -262,29 +312,14 @@ int ECP_ipipv_cart(double *gctr, int smem_size3 = li1*lic1*nfi; // angi int smem_size4 = lj1*ljc1*nfj; // angj - //int NF2_MAX = (AO_LMAX+3)*(AO_LMAX+4)/2; int NF1_MAX = (AO_LMAX+2)*(AO_LMAX+3)/2; int NF0_MAX = (AO_LMAX+1)*(AO_LMAX+2)/2; - //int static_smem_size = NF2_MAX*NF0_MAX; int dynamic_smem_size = smem_size0 + smem_size1 + smem_size2 + smem_size3 + smem_size4; dynamic_smem_size = max(dynamic_smem_size, 3*NF1_MAX*NF0_MAX); - //int total_smem_size = static_smem_size + dynamic_smem_size; - - cudaError_t err = cudaFuncSetAttribute(type2_cart_ipipv, - cudaFuncAttributeMaxDynamicSharedMemorySize, - dynamic_smem_size*sizeof(double)); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in cudaFuncSetAttribute %s (li,lj,lc = %d,%d,%d): %s\n", __func__, li,lj,lc, cudaGetErrorString(err)); - return 1; - } - - type2_cart_ipipv<<>>( - gctr, li, lj, lc, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + ECP_LAUNCH_GENERAL(type2_cart_ipipv_kernel, dynamic_smem_size, type2_cart_ipipv, + gctr, li, lj, lc, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -301,22 +336,24 @@ int ECP_ipvip_cart(double *gctr, const int *atm, const int *bas, const double *env, const int li, const int lj, const int lc){ // one task per thread block +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(ntasks); + sycl::queue& stream = *sycl_get_queue(); +#else dim3 threads(THREADS); dim3 blocks(ntasks); +#endif if (lc < 0){ - const int lij1 = li+lj+3; // + const int lij1 = li+lj+3; const int lij3 = lij1*lij1*lij1; - int smem_size = 0; - smem_size += lij3; // rad_ang - smem_size += lij1*lij1; // rad_all - type1_cart_ipvip<<>>( - gctr, li, lj, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + int smem_size = lij3 + lij1*lij1; + + ECP_LAUNCH_GENERAL(type1_cart_ipvip_kernel, smem_size, type1_cart_ipvip, + gctr, li, lj, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); } else { const int li1 = li+2; const int lj1 = lj+2; @@ -337,26 +374,12 @@ int ECP_ipvip_cart(double *gctr, int NF1_MAX = (AO_LMAX+2)*(AO_LMAX+3)/2; int NF0_MAX = (AO_LMAX+1)*(AO_LMAX+2)/2; - //int static_smem_size = NF1_MAX*NF1_MAX; int dynamic_smem_size = smem_size0 + smem_size1 + smem_size2 + smem_size3 + smem_size4; dynamic_smem_size = max(dynamic_smem_size, 3*NF0_MAX*NF1_MAX); - //int total_smem_size = static_smem_size + dynamic_smem_size; - cudaError_t err = cudaFuncSetAttribute( - type2_cart_ipvip, - cudaFuncAttributeMaxDynamicSharedMemorySize, - dynamic_smem_size*sizeof(double)); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in cudaFuncSetAttribute %s (li,lj,lc = %d,%d,%d): %s\n", __func__, li,lj,lc, cudaGetErrorString(err)); - return 1; - } - - type2_cart_ipvip<<>>( - gctr, li, lj, lc, - ao_loc, nao, - tasks, ntasks, - ecpbas, ecploc, - atm, bas, env); + ECP_LAUNCH_GENERAL(type2_cart_ipvip_kernel, dynamic_smem_size, type2_cart_ipvip, + gctr, li, lj, lc, ao_loc, nao, + tasks, ntasks, ecpbas, ecploc, atm, bas, env); } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -366,3 +389,8 @@ int ECP_ipvip_cart(double *gctr, return 0; } } + +#undef ECP_ARGS +#undef ECP_LAUNCH1 +#undef ECP_LAUNCH2 +#undef ECP_LAUNCH_GENERAL diff --git a/gpu4pyscf/lib/gdft/CMakeLists.txt b/gpu4pyscf/lib/gdft/CMakeLists.txt index 7d5f74a19..c181b5c88 100644 --- a/gpu4pyscf/lib/gdft/CMakeLists.txt +++ b/gpu4pyscf/lib/gdft/CMakeLists.txt @@ -12,16 +12,85 @@ # See the License for the specific language governing permissions and # limitations under the License. -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128") +# ExchCXX is built as a static dependency and linked into the shared +# libgdft.so / libxc.so below, so every object in this directory must be +# position-independent. +set(CMAKE_POSITION_INDEPENDENT_CODE ON) -add_library(gdft SHARED +set(GPU_SRCS nr_eval_gto.cu contract_rho.cu gen_grids.cu nr_numint_sparse.cu vv10.cu - libxc.cu ) +if(USE_SYCL) + # Step 1: compile exchcxx.cpp ONCE as an OBJECT library. + # Isolated from GPU_SRCS so changes to grid/numint kernels never + # retrigger this expensive AoT SPIR-V compilation (~35 min). + # Only rebuilt when exchcxx.cpp itself changes. + add_library(exchcxx_shim OBJECT + ${CMAKE_CURRENT_SOURCE_DIR}/exchcxx.cpp + ) + set_source_files_properties( + ${CMAKE_CURRENT_SOURCE_DIR}/exchcxx.cpp PROPERTIES LANGUAGE CXX + ) + set_target_properties(exchcxx_shim PROPERTIES + POSITION_INDEPENDENT_CODE ON + ) + target_compile_options(exchcxx_shim PRIVATE + -x c++ -nocudainc -nocudalib + ) + target_include_directories(exchcxx_shim PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} # exchcxx.h, gint/sycl_device.hpp + ${PROJECT_SOURCE_DIR} # project-wide includes + ) + # Provides exchcxx/exchcxx.hpp include path + EXCHCXX_HAS_CONFIG_H=1 + target_link_libraries(exchcxx_shim PRIVATE ExchCXX::ExchCXX) + add_dependencies(exchcxx_shim ExchCXX) + + # Step 2: libgdft.so. + # GPU grid/numint kernels + exchcxx.o via TARGET_OBJECTS (no recompile). + # Python libgdft handle: GDFT_xc_lda/gga/mgga entry points live here. + # Changing any GPU_SRCS file only recompiles that file + relinks gdft. + file(GLOB CUH_HEADERS "${CMAKE_CURRENT_SOURCE_DIR}/*.cuh") + list(APPEND CUH_HEADERS exchcxx.h) # .h suffix not caught by *.cuh glob + add_library(gdft SHARED + ${GPU_SRCS} + ${CUH_HEADERS} + $ # reuse compiled object, never recompiles + ) + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + target_compile_options(gdft PRIVATE -x c++ -nocudainc -nocudalib) + # Required: exchcxx.o inside libgdft.so calls ExchCXX device kernel symbols + # that must be resolved against libexchcxx.so at runtime + target_link_libraries(gdft PRIVATE sycl_compat gsycl ExchCXX::ExchCXX) + + # Step 3: libxc.so -- the libxc ABI shim Python loads as _libxc. + # Python _libxc handle: xc_func_alloc/init/end/free + + # xc_version_string + xc_number_of_functionals etc. + # Same exchcxx.o reused via TARGET_OBJECTS -- zero recompilation cost. + # Replaces the old create_symlink approach with a real DSO target. + add_library(xc SHARED + $ # reuse same compiled object + ) + # Required: exchcxx.o inside libxc.so calls ExchCXX device kernel symbols + target_link_libraries(xc PRIVATE ExchCXX::ExchCXX) + set_target_properties(xc PROPERTIES + OUTPUT_NAME xc + # Python searches gpu4pyscf/lib/deps/lib/libxc.so + LIBRARY_OUTPUT_DIRECTORY "${PROJECT_SOURCE_DIR}/deps/lib" + # libexchcxx.so is in the same deps/lib/ directory + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + ) + +else() + list(APPEND GPU_SRCS libxc.cu) + add_library(gdft SHARED ${GPU_SRCS}) + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v") +endif() + set_target_properties(gdft PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) diff --git a/gpu4pyscf/lib/gdft/contract_rho.cu b/gpu4pyscf/lib/gdft/contract_rho.cu index 1987fd492..59d4cdfc5 100644 --- a/gpu4pyscf/lib/gdft/contract_rho.cu +++ b/gpu4pyscf/lib/gdft/contract_rho.cu @@ -21,32 +21,50 @@ #include #include #include "contract_rho.cuh" +// Tree reduction along iy dimension in shared memory buf. +#define REDUCE_Y(buf, ixy, iy) \ + for (int _s_ = BLKSIZEY >> 1; _s_ > 0; _s_ >>= 1) { \ + if ((iy) < _s_) { \ + (buf)[(ixy)] += (buf)[(ixy) + BLKSIZEX * _s_]; \ + } \ + __syncthreads(); \ + } +static_assert((BLKSIZEY & (BLKSIZEY - 1)) == 0, "BLKSIZEY must be a power of 2"); + // TODO: improve this? __global__ -void GDFTcontract_rho_kernel(double *rho, double *bra, double *ket, int ngrids, int nao) +void GDFTcontract_rho_kernel(double *rho, const double *bra, const double *ket, int ngrids, int nao) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int grid_id = item.get_global_id(1); + sycl::group thread_block = item.get_group(); + using tile_t = double[BLKSIZEX*(BLKSIZEY+1)]; + tile_t& buf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + const int threadIdx_y = item.get_local_id(0); + int ix = item.get_local_id(1); + int iy = item.get_local_id(0); +#else int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; + int threadIdx_y = threadIdx.y; + int ix = threadIdx.x; + int iy = threadIdx.y; +#endif + const bool active = grid_id < ngrids; size_t Ngrids = ngrids; double v = 0; if (active){ - for (int ao_id = threadIdx.y; ao_id < nao; ao_id += BLKSIZEY) { + for (int ao_id = threadIdx_y; ao_id < nao; ao_id += BLKSIZEY) { int ket_idx = grid_id + ao_id * Ngrids; - v += bra[ket_idx] * ket[ket_idx]; + v += (bra[ket_idx] * ket[ket_idx]); } } - __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; - int ix = threadIdx.x; - int iy = threadIdx.y; int ixy = ix + BLKSIZEX * iy; buf[ixy] = v; __syncthreads(); - - if (blockDim.y >= 32 && iy < 16) buf[ixy] += buf[ixy + BLKSIZEX * 16]; __syncthreads(); - if (blockDim.y >= 16 && iy < 8) buf[ixy] += buf[ixy + BLKSIZEX * 8]; __syncthreads(); - if (blockDim.y >= 8 && iy < 4) buf[ixy] += buf[ixy + BLKSIZEX * 4]; __syncthreads(); - if (blockDim.y >= 4 && iy < 2) buf[ixy] += buf[ixy + BLKSIZEX * 2]; __syncthreads(); - if (blockDim.y >= 2 && iy < 1) buf[ixy] += buf[ixy + BLKSIZEX * 1]; __syncthreads(); + REDUCE_Y(buf, ixy, iy); if (iy == 0 && active) { rho[grid_id] = buf[ix]; @@ -57,17 +75,30 @@ void GDFTcontract_rho_kernel(double *rho, double *bra, double *ket, int ngrids, __global__ void GDFTcontract_rho4_kernel(double *rho, double *bra, double *ket, int ngrids, int nao, int count) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int grid_id = item.get_global_id(1); + sycl::group thread_block = item.get_group(); + using tile_t = double[BLKSIZEX*(BLKSIZEY+1)]; + tile_t& buf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + const int threadIdx_y = item.get_local_id(0); + int ix = item.get_local_id(1); + int iy = item.get_local_id(0); +#else int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; + int threadIdx_y = threadIdx.y; + int ix = threadIdx.x; + int iy = threadIdx.y; +#endif const bool active = grid_id < ngrids; size_t ket_stride = nao * ngrids; size_t rho_stride = count * ngrids; - __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; - for (int ia = 0; ia < count; ia++){ double v[4] = {0.0, 0.0, 0.0, 0.0}; if (active){ - for (int ao_id = threadIdx.y; ao_id < nao; ao_id += BLKSIZEY) { + for (int ao_id = threadIdx_y; ao_id < nao; ao_id += BLKSIZEY) { int ket_idx = grid_id + ao_id * ngrids; double bra_tmp = bra[ket_idx + ia * ket_stride]; v[0] += bra_tmp * ket[0*ket_stride + ket_idx]; @@ -77,16 +108,10 @@ void GDFTcontract_rho4_kernel(double *rho, double *bra, double *ket, int ngrids, } } - int ix = threadIdx.x; - int iy = threadIdx.y; int ixy = ix + BLKSIZEX * iy; for (int i = 0; i < 4; i++){ buf[ixy] = v[i]; __syncthreads(); - if (blockDim.y >= 32 && iy < 16) buf[ixy] += buf[ixy + BLKSIZEX * 16]; __syncthreads(); - if (blockDim.y >= 16 && iy < 8) buf[ixy] += buf[ixy + BLKSIZEX * 8]; __syncthreads(); - if (blockDim.y >= 8 && iy < 4) buf[ixy] += buf[ixy + BLKSIZEX * 4]; __syncthreads(); - if (blockDim.y >= 4 && iy < 2) buf[ixy] += buf[ixy + BLKSIZEX * 2]; __syncthreads(); - if (blockDim.y >= 2 && iy < 1) buf[ixy] += buf[ixy + BLKSIZEX * 1]; __syncthreads(); + REDUCE_Y(buf, ixy, iy); if (iy == 0 && active) { rho[grid_id + ia * ngrids + rho_stride * i] = buf[ix]; @@ -98,7 +123,21 @@ void GDFTcontract_rho4_kernel(double *rho, double *bra, double *ket, int ngrids, __global__ void GDFTcontract_rho_gga_kernel(double *rho, double *bra, double *ket, int ngrids, int nao) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int grid_id = item.get_global_id(1); + using tile_t = double[BLKSIZEX*(BLKSIZEY+1)]; + tile_t& buf = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + const int ix = item.get_local_id(1); + const int iy = item.get_local_id(0); + const int threadIdx_y = item.get_local_id(0); +#else int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; + int ix = threadIdx.x; + int iy = threadIdx.y; + const int threadIdx_y = threadIdx.y; +#endif const bool active = grid_id < ngrids; size_t Ngrids = ngrids; @@ -106,7 +145,7 @@ void GDFTcontract_rho_gga_kernel(double *rho, double *bra, double *ket, int ngri double v[4] = {0.0, 0.0, 0.0, 0.0}; if (active){ - for (int ao_id = threadIdx.y; ao_id < nao; ao_id += BLKSIZEY) { + for (int ao_id = threadIdx_y; ao_id < nao; ao_id += BLKSIZEY) { size_t ket_idx = grid_id + ao_id * Ngrids; double bra_tmp = bra[ket_idx]; double ket_tmp = ket[ket_idx]; @@ -127,18 +166,11 @@ void GDFTcontract_rho_gga_kernel(double *rho, double *bra, double *ket, int ngri } } - __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; - int ix = threadIdx.x; - int iy = threadIdx.y; int ixy = ix + BLKSIZEX * iy; for (int i = 0; i < 4; i++){ buf[ixy] = v[i]; __syncthreads(); - if (blockDim.y >= 32 && iy < 16) buf[ixy] += buf[ixy + BLKSIZEX * 16]; __syncthreads(); - if (blockDim.y >= 16 && iy < 8) buf[ixy] += buf[ixy + BLKSIZEX * 8]; __syncthreads(); - if (blockDim.y >= 8 && iy < 4) buf[ixy] += buf[ixy + BLKSIZEX * 4]; __syncthreads(); - if (blockDim.y >= 4 && iy < 2) buf[ixy] += buf[ixy + BLKSIZEX * 2]; __syncthreads(); - if (blockDim.y >= 2 && iy < 1) buf[ixy] += buf[ixy + BLKSIZEX * 1]; __syncthreads(); + REDUCE_Y(buf, ixy, iy); if (iy == 0 && active) { rho[grid_id + ngrids * i] = buf[ix]; @@ -150,7 +182,21 @@ void GDFTcontract_rho_gga_kernel(double *rho, double *bra, double *ket, int ngri __global__ void GDFTcontract_rho_mgga_kernel(double *rho, double *bra, double *ket, int ngrids, int nao) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int threadIdx_y = item.get_local_id(0); + const int grid_id = item.get_global_id(1); + using tile_t = double[BLKSIZEX*(BLKSIZEY+1)]; + tile_t& buf = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + const int ix = item.get_local_id(1); + const int iy = item.get_local_id(0); +#else + int threadIdx_y = threadIdx.y; int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; + int ix = threadIdx.x; + int iy = threadIdx.y; +#endif const bool active = grid_id < ngrids; size_t Ngrids = ngrids; @@ -158,7 +204,7 @@ void GDFTcontract_rho_mgga_kernel(double *rho, double *bra, double *ket, int ngr double v[5] = {0.0, 0.0, 0.0, 0.0, 0.0}; if (active){ - for (int ao_id = threadIdx.y; ao_id < nao; ao_id += BLKSIZEY) { + for (int ao_id = threadIdx_y; ao_id < nao; ao_id += BLKSIZEY) { size_t ket_idx = grid_id + ao_id * Ngrids; double bra_tmp0 = bra[ket_idx]; double ket_tmp0 = ket[ket_idx]; @@ -191,18 +237,11 @@ void GDFTcontract_rho_mgga_kernel(double *rho, double *bra, double *ket, int ngr v[4] *= 0.5; - __shared__ double buf[BLKSIZEX*(BLKSIZEY+1)]; - int ix = threadIdx.x; - int iy = threadIdx.y; int ixy = ix + BLKSIZEX * iy; for (int i = 0; i < 5; i++){ buf[ixy] = v[i]; __syncthreads(); - if (blockDim.y >= 32 && iy < 16) buf[ixy] += buf[ixy + BLKSIZEX * 16]; __syncthreads(); - if (blockDim.y >= 16 && iy < 8) buf[ixy] += buf[ixy + BLKSIZEX * 8]; __syncthreads(); - if (blockDim.y >= 8 && iy < 4) buf[ixy] += buf[ixy + BLKSIZEX * 4]; __syncthreads(); - if (blockDim.y >= 4 && iy < 2) buf[ixy] += buf[ixy + BLKSIZEX * 2]; __syncthreads(); - if (blockDim.y >= 2 && iy < 1) buf[ixy] += buf[ixy + BLKSIZEX * 1]; __syncthreads(); + REDUCE_Y(buf, ixy, iy); if (iy == 0 && active) { rho[grid_id + ngrids * i] = buf[ix]; @@ -214,8 +253,14 @@ static __global__ void dscale_ao_kernel(double *out, double *ket, double *wv, int ngrids, int nao, int nvar) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int grid_id = item.get_global_id(1); + int ao_id = item.get_global_id(0); +#else int grid_id = blockIdx.x * blockDim.x + threadIdx.x; int ao_id = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (grid_id >= ngrids || ao_id >= nao) { return; } @@ -235,8 +280,14 @@ static __global__ void zscale_ao_kernel(double *out, double *ket, double *wv, int ngrids, int nao, int nvar) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int grid_id = item.get_global_id(1); + int ao_id = item.get_global_id(0); +#else int grid_id = blockIdx.x * blockDim.x + threadIdx.x; int ao_id = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (grid_id >= ngrids || ao_id >= nao) { return; } @@ -263,10 +314,16 @@ void zscale_ao_kernel(double *out, double *ket, double *wv, __global__ void GDFT_make_dR_dao_w_kernel(double *out, double *ket, double *wv, - int ngrids, int nao) + int ngrids, int nao) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int grid_id = item.get_global_id(1); + int ao_id = item.get_global_id(0); + #else int grid_id = blockIdx.x * blockDim.x + threadIdx.x; int ao_id = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (grid_id >= ngrids || ao_id >= nao) { return; } @@ -306,8 +363,14 @@ void GDFT_make_dR_dao_w_kernel(double *out, double *ket, double *wv, extern "C"{ __host__ -int GDFTcontract_rho(cudaStream_t stream, double *rho, double *bra, double *ket, int ngrids, int nao) +int GDFTcontract_rho(cudaStream_t stream, double *rho, const double *bra, const double *ket, int ngrids, int nao) { +#ifdef USE_SYCL + sycl::range<2> threads(BLKSIZEY, BLKSIZEX); + sycl::range<2> blocks(1, (ngrids+BLKSIZEX-1)/BLKSIZEX); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTcontract_rho_kernel(rho, bra, ket, ngrids, nao); }); +#else dim3 threads(BLKSIZEX, BLKSIZEY); dim3 blocks((ngrids+BLKSIZEX-1)/BLKSIZEX); GDFTcontract_rho_kernel<<>>(rho, bra, ket, ngrids, nao); @@ -316,11 +379,19 @@ int GDFTcontract_rho(cudaStream_t stream, double *rho, double *bra, double *ket, fprintf(stderr, "CUDA Error of GDFTcontract_rho: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int GDFTcontract_rho4(cudaStream_t stream, double *rho, double *bra, double *ket, int ngrids, int nao, int count) { +#ifdef USE_SYCL + sycl::range<2> threads(BLKSIZEY, BLKSIZEX); + sycl::range<2> blocks(1, (ngrids+BLKSIZEX-1)/BLKSIZEX); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTcontract_rho4_kernel(rho, bra, ket, ngrids, nao, count); + }); +#else dim3 threads(BLKSIZEX, BLKSIZEY); dim3 blocks((ngrids+BLKSIZEX-1)/BLKSIZEX); GDFTcontract_rho4_kernel<<>>(rho, bra, ket, ngrids, nao, count); @@ -329,11 +400,19 @@ int GDFTcontract_rho4(cudaStream_t stream, double *rho, double *bra, double *ket fprintf(stderr, "CUDA Error of GDFTcontract_rho: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int GDFTcontract_rho_gga(cudaStream_t stream, double *rho, double *bra, double *ket, int ngrids, int nao) { +#ifdef USE_SYCL + sycl::range<2> threads(BLKSIZEY, BLKSIZEX); + sycl::range<2> blocks(1, (ngrids+BLKSIZEX-1)/BLKSIZEX); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTcontract_rho_gga_kernel(rho, bra, ket, ngrids, nao); + }); +#else dim3 threads(BLKSIZEX, BLKSIZEY); dim3 blocks((ngrids+BLKSIZEX-1)/BLKSIZEX); GDFTcontract_rho_gga_kernel<<>>(rho, bra, ket, ngrids, nao); @@ -342,11 +421,19 @@ int GDFTcontract_rho_gga(cudaStream_t stream, double *rho, double *bra, double * fprintf(stderr, "CUDA Error of GDFTcontract_rho_gga: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int GDFTcontract_rho_mgga(cudaStream_t stream, double *rho, double *bra, double *ket, int ngrids, int nao) { +#ifdef USE_SYCL + sycl::range<2> threads(BLKSIZEY, BLKSIZEX); + sycl::range<2> blocks(1, (ngrids+BLKSIZEX-1)/BLKSIZEX); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTcontract_rho_mgga_kernel(rho, bra, ket, ngrids, nao); + }); +#else dim3 threads(BLKSIZEX, BLKSIZEY); dim3 blocks((ngrids+BLKSIZEX-1)/BLKSIZEX); GDFTcontract_rho_mgga_kernel<<>>(rho, bra, ket, ngrids, nao); @@ -355,12 +442,20 @@ int GDFTcontract_rho_mgga(cudaStream_t stream, double *rho, double *bra, double fprintf(stderr, "CUDA Error of GDFTcontract_rho_mgga: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int GDFT_make_dR_dao_w(cudaStream_t stream, double *out, double *ket, double *wv, int ngrids, int nao) { +#ifdef USE_SYCL + sycl::range<2> threads(BLKSIZEY, BLKSIZEX); + sycl::range<2> blocks((nao+BLKSIZEY-1)/BLKSIZEY, (ngrids+BLKSIZEX-1)/BLKSIZEX); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFT_make_dR_dao_w_kernel(out, ket, wv, ngrids, nao); + }); +#else dim3 threads(BLKSIZEX, BLKSIZEY); dim3 blocks((ngrids+BLKSIZEX-1)/BLKSIZEX, (nao+BLKSIZEY-1)/BLKSIZEY); GDFT_make_dR_dao_w_kernel<<>>(out, ket, wv, ngrids, nao); @@ -369,12 +464,27 @@ int GDFT_make_dR_dao_w(cudaStream_t stream, double *out, double *ket, double *wv fprintf(stderr, "CUDA Error of GDFT_make_dR_dao_w: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int GDFTscale_ao(double *out, double *ket, double *wv, int ngrids, int nao, int nvar, int is_real) { +#ifdef USE_SYCL + sycl::queue& stream = *sycl_get_queue(); + sycl::range<2> threads(BLKSIZEY, BLKSIZEX); + sycl::range<2> blocks((nao+BLKSIZEY-1)/BLKSIZEY, (ngrids+BLKSIZEX-1)/BLKSIZEX); + if (is_real) { + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + dscale_ao_kernel(out, ket, wv, ngrids, nao, nvar); + }); + } else { + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + zscale_ao_kernel(out, ket, wv, ngrids, nao, nvar); + }); + } +#else dim3 threads(BLKSIZEX, BLKSIZEY); dim3 blocks((ngrids+BLKSIZEX-1)/BLKSIZEX, (nao+BLKSIZEY-1)/BLKSIZEY); if (is_real) { @@ -387,6 +497,7 @@ int GDFTscale_ao(double *out, double *ket, double *wv, fprintf(stderr, "CUDA Error of GDFTscale_ao: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } diff --git a/gpu4pyscf/lib/gdft/contract_rho.cuh b/gpu4pyscf/lib/gdft/contract_rho.cuh index 112d4a7cd..8d65dca9a 100644 --- a/gpu4pyscf/lib/gdft/contract_rho.cuh +++ b/gpu4pyscf/lib/gdft/contract_rho.cuh @@ -18,7 +18,7 @@ #define BLKSIZEY 16 __global__ -void GDFTcontract_rho_kernel(double *rho, double *bra, double *ket, int ngrids, int nao); +void GDFTcontract_rho_kernel(double *rho, const double *bra, const double *ket, int ngrids, int nao); __global__ void GDFTscale_ao_kernel(double *out, double *ket, double *wv, int ngrids, int nao, int nvar); diff --git a/gpu4pyscf/lib/gdft/exchcxx.cpp b/gpu4pyscf/lib/gdft/exchcxx.cpp new file mode 100644 index 000000000..9e00afa28 --- /dev/null +++ b/gpu4pyscf/lib/gdft/exchcxx.cpp @@ -0,0 +1,1855 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include "exchcxx.h" // ABI structs + +// Verbose tracing of functional construction and of every GGA evaluation. +// Off by default: the GGA dump copies np doubles device->host per call, so +// leaving it on would dominate the runtime of any real SCF. +// Build with -DGDFT_EXCHCXX_TRACE=1 to enable. +#ifndef GDFT_EXCHCXX_TRACE +#define GDFT_EXCHCXX_TRACE 0 +#endif +#if GDFT_EXCHCXX_TRACE +#define GDFT_TRACE(...) std::fprintf(stderr, __VA_ARGS__) +#else +#define GDFT_TRACE(...) ((void)0) +#endif + +namespace detail { + + static std::string to_upper(std::string s){ + for(auto &c : s) c = char(std::toupper(unsigned(c))); + return s; + } + + // Fast path: look up by canonical ExchCXX functional name (already in functional_map) + inline std::optional + functional_from_string(std::string_view s) { + const auto key = to_upper(std::string{s}); + try { + return ExchCXX::functional_map.value(key); // name -> enum + } catch (const std::out_of_range&) { + return std::nullopt; // not present + } + } + + // Minimal alias map: LibXC “family labels” → ExchCXX canonical name + static const std::unordered_map kLibXCAliases = { + // Hybrids / composites + {"HYB_GGA_XC_B3LYP", "B3LYP"}, // this is an issue since B3LYP uses VWN5 varient, but looks like pyscf 2.3.0> uses B3LYP VWN_RPA version + {"HYB_GGA_XC_PBEH", "PBE0"}, // ok + {"HYB_GGA_XC_HSE03", "HSE03"}, // ok + {"HYB_GGA_XC_HSE06", "HSE06"}, // ok + {"HYB_GGA_XC_CAM_B3LYP","CAMB3LYP"}, // ok + {"HYB_MGGA_X_SCAN0", "SCAN0"}, + {"HYB_GGA_XC_B3PW91", "B3PW91"}, // ok + {"HYB_GGA_XC_BHANDH", "BHANDH"}, + {"HYB_GGA_XC_O3LYP", "O3LYP"}, // ok + + // Pure composites + {"GGA_XC_PBE", "PBE"}, + {"GGA_XC_REVPBE", "REVPBE"}, + {"GGA_XC_BLYP", "BLYP"}, // needs checking + {"GGA_XC_BP86", "BP86"}, + {"GGA_XC_PW91", "PW91"}, // may be incorrect + {"GGA_XC_RPBE", "RPBE"}, + {"GGA_XC_X3LYP", "X3LYP"}, + {"GGA_XC_XLYP", "XLYP"}, + {"GGA_XC_OPBE", "OPBE"}, // incorrect + {"GGA_XC_OLYP", "OLYP"}, // incorrect + + // mGGA composites + {"MGGA_XC_SCAN", "SCAN"}, + {"MGGA_XC_R2SCAN", "R2SCAN"}, + {"MGGA_XC_TPSS", "TPSS"}, + {"MGGA_XC_REVTPSS", "REVTPSS"}, + {"MGGA_XC_M06_L", "M06L"}, + + // LDA packs + {"LDA_XC_VWN", "SPW92"}, // LibXC’s “VWN” combo equals Slater + VWN + {"LDA_XC_SVWN", "SVWN5"}, // Many LibXC builds alias SVWN→SVWN5 + {"LDA_XC_SVWN3", "SVWN3"}, + {"LDA_XC_SVWN5", "SVWN5"}, + // {"LDA_XC_LDA", "LDA"}, + }; + + // Robust LibXC → ExchCXX functional resolution. + // 1) Direct hit in alias table + // 2) Heuristics for common pairs (e.g. GGA_X_PBE + GGA_C_PBE ⇒ PBE) can be handled + // by your caller if it sees both half-labels; for single labels use the alias table. + inline std::optional + libxc_name_to_functional(std::string_view libxc_name) { + GDFT_TRACE("[gdft] libxc_name_to_functional(): %s\n", + std::string(libxc_name).c_str()); + auto key = to_upper(std::string{libxc_name}); + + // 1) Exact alias → canonical functional name + if (auto it = kLibXCAliases.find(key); it != kLibXCAliases.end()) { + if (auto f = functional_from_string(it->second)) return f; + } + + // 2) Loose pattern matches for families (cheap and safe) + // CAM-B3LYP spellings vary a bit across LibXC versions + if (key.find("CAM") != std::string::npos && key.find("B3LYP") != std::string::npos) { + if (auto f = functional_from_string("CAMB3LYP")) return f; + } + // LRC-ωPBE family + if (key.find("LC_WPBE") != std::string::npos || key.find("LRC_WPBE") != std::string::npos) { + if (key.find('H') != std::string::npos) { + if (auto f = functional_from_string("LRCWPBEH")) return f; + } else { + if (auto f = functional_from_string("LRCWPBE")) return f; + } + } + + // 3) Already a canonical ExchCXX name? (users sometimes pass that) + if (auto f = functional_from_string(key)) return f; + + return std::nullopt; + } + + // The following map is from https://github.com/wavefunction91/ExchCXX/blob/master/src/libxc.cxx#L73 + std::unordered_map< ExchCXX::Kernel, std::string > libxc_kernel_map { + // LDA Functionals + { ExchCXX::Kernel::SlaterExchange, "LDA_X" }, + { ExchCXX::Kernel::VWN3, "LDA_C_VWN_3" }, + { ExchCXX::Kernel::VWN5, "LDA_C_VWN_RPA" }, + { ExchCXX::Kernel::VWN, "LDA_C_VWN" }, + { ExchCXX::Kernel::PZ81, "LDA_C_PZ" }, + { ExchCXX::Kernel::PZ81_MOD, "LDA_C_PZ_MOD" }, + { ExchCXX::Kernel::PW91_LDA, "LDA_C_PW" }, + { ExchCXX::Kernel::PW91_LDA_MOD, "LDA_C_PW_MOD" }, + { ExchCXX::Kernel::PW91_LDA_RPA, "LDA_C_PW_RPA" }, + + // GGA Functionals + { ExchCXX::Kernel::PBE_X, "GGA_X_PBE" }, + { ExchCXX::Kernel::PBE_C, "GGA_C_PBE" }, + { ExchCXX::Kernel::revPBE_X, "GGA_X_PBE_R" }, + { ExchCXX::Kernel::B88, "GGA_X_B88" }, + { ExchCXX::Kernel::LYP, "GGA_C_LYP" }, + { ExchCXX::Kernel::B97_D, "GGA_XC_B97_D" }, + { ExchCXX::Kernel::ITYH_X, "GGA_X_ITYH" }, + { ExchCXX::Kernel::P86_C, "GGA_C_P86" }, + { ExchCXX::Kernel::P86VWN_FT_C, "GGA_C_P86VWN_FT" }, + { ExchCXX::Kernel::PW91_C, "GGA_C_PW91" }, + { ExchCXX::Kernel::PBE_SOL_C, "GGA_C_PBE_SOL" }, + { ExchCXX::Kernel::BMK_C, "GGA_C_BMK" }, + { ExchCXX::Kernel::N12_C, "GGA_C_N12" }, + { ExchCXX::Kernel::N12_SX_C, "GGA_C_N12_SX" }, + { ExchCXX::Kernel::SOGGA11_X_C, "GGA_C_SOGGA11_X" }, + { ExchCXX::Kernel::PW91_X, "GGA_X_PW91" }, + { ExchCXX::Kernel::MPW91_X, "GGA_X_MPW91" }, + { ExchCXX::Kernel::OPTX_X, "GGA_X_OPTX" }, + { ExchCXX::Kernel::RPBE_X, "GGA_X_RPBE" }, + { ExchCXX::Kernel::SOGGA11_X_X, "HYB_GGA_X_SOGGA11_X" }, + { ExchCXX::Kernel::PW86_X, "GGA_X_PW86" }, + { ExchCXX::Kernel::wB97_XC, "HYB_GGA_XC_WB97" }, + { ExchCXX::Kernel::wB97X_XC, "HYB_GGA_XC_WB97X" }, + { ExchCXX::Kernel::wB97X_V_XC, "HYB_GGA_XC_WB97X_V"}, + { ExchCXX::Kernel::wB97X_D_XC, "HYB_GGA_XC_WB97X_D"}, + { ExchCXX::Kernel::wB97X_D3_XC, "HYB_GGA_XC_WB97X_D3"}, + { ExchCXX::Kernel::HJS_PBE_X, "GGA_X_HJS_PBE" }, + { ExchCXX::Kernel::wPBEh_X_default0, "GGA_X_WPBEH" }, + + // MGGA Functionals + { ExchCXX::Kernel::SCAN_C, "MGGA_C_SCAN" }, + { ExchCXX::Kernel::SCAN_X, "MGGA_X_SCAN" }, + { ExchCXX::Kernel::SCANL_C, "MGGA_C_SCANL" }, + { ExchCXX::Kernel::SCANL_X, "MGGA_X_SCANL" }, + { ExchCXX::Kernel::R2SCAN_C, "MGGA_C_R2SCAN" }, + { ExchCXX::Kernel::R2SCAN_X, "MGGA_X_R2SCAN" }, + { ExchCXX::Kernel::R2SCANL_C, "MGGA_C_R2SCANL" }, + { ExchCXX::Kernel::R2SCANL_X, "MGGA_X_R2SCANL" }, + { ExchCXX::Kernel::FT98_X, "MGGA_X_FT98" }, + { ExchCXX::Kernel::M062X_X, "HYB_MGGA_X_M06_2X" }, + { ExchCXX::Kernel::M062X_C, "MGGA_C_M06_2X" }, + { ExchCXX::Kernel::PKZB_X, "MGGA_X_PKZB" }, + { ExchCXX::Kernel::PKZB_C, "MGGA_C_PKZB" }, + { ExchCXX::Kernel::TPSS_X, "MGGA_X_TPSS" }, + { ExchCXX::Kernel::revTPSS_X, "MGGA_X_REVTPSS" }, + { ExchCXX::Kernel::M06_L_X, "MGGA_X_M06_L" }, + { ExchCXX::Kernel::M06_X, "HYB_MGGA_X_M06" }, + { ExchCXX::Kernel::revM06_L_X, "MGGA_X_REVM06_L" }, + { ExchCXX::Kernel::M06_HF_X, "HYB_MGGA_X_M06_HF" }, + { ExchCXX::Kernel::M06_SX_X, "HYB_MGGA_X_M06_SX" }, + { ExchCXX::Kernel::M06_L_C, "MGGA_C_M06_L" }, + { ExchCXX::Kernel::M06_C, "MGGA_C_M06" }, + { ExchCXX::Kernel::revM06_L_C, "MGGA_C_REVM06_L" }, + { ExchCXX::Kernel::M06_HF_C, "MGGA_C_M06_HF" }, + { ExchCXX::Kernel::M06_SX_C, "MGGA_C_M06_SX" }, + { ExchCXX::Kernel::M05_2X_C, "MGGA_C_M05_2X" }, + { ExchCXX::Kernel::M05_C, "MGGA_C_M05" }, + { ExchCXX::Kernel::M08_HX_C, "MGGA_C_M08_HX" }, + { ExchCXX::Kernel::M08_SO_C, "MGGA_C_M08_SO" }, + { ExchCXX::Kernel::CF22D_C, "MGGA_C_CF22D" }, + { ExchCXX::Kernel::M11_C, "MGGA_C_M11" }, + { ExchCXX::Kernel::MN12_L_C, "MGGA_C_MN12_L" }, + { ExchCXX::Kernel::MN12_SX_C, "MGGA_C_MN12_SX" }, + { ExchCXX::Kernel::MN15_C, "MGGA_C_MN15" }, + { ExchCXX::Kernel::MN15_L_C, "MGGA_C_MN15_L" }, + { ExchCXX::Kernel::TPSS_C, "MGGA_C_TPSS" }, + { ExchCXX::Kernel::revTPSS_C, "MGGA_C_REVTPSS" }, + { ExchCXX::Kernel::RSCAN_C, "MGGA_C_RSCAN" }, + { ExchCXX::Kernel::BC95_C, "MGGA_C_BC95" }, + { ExchCXX::Kernel::mBEEF_X, "MGGA_X_MBEEF" }, + { ExchCXX::Kernel::RSCAN_X, "MGGA_X_RSCAN" }, + { ExchCXX::Kernel::BMK_X, "HYB_MGGA_X_BMK" }, + { ExchCXX::Kernel::M08_HX_X, "HYB_MGGA_X_M08_HX" }, + { ExchCXX::Kernel::M08_SO_X, "HYB_MGGA_X_M08_SO" }, + { ExchCXX::Kernel::MN12_L_X, "MGGA_X_MN12_L" }, + { ExchCXX::Kernel::MN15_L_X, "MGGA_X_MN15_L" }, + { ExchCXX::Kernel::MN15_X, "HYB_MGGA_X_MN15" }, + { ExchCXX::Kernel::CF22D_X, "HYB_MGGA_X_CF22D" }, + { ExchCXX::Kernel::MN12_SX_X, "HYB_MGGA_X_MN12_SX"}, + { ExchCXX::Kernel::M11_X, "HYB_MGGA_X_M11" }, + { ExchCXX::Kernel::M05_X, "HYB_MGGA_X_M05" }, + { ExchCXX::Kernel::M05_2X_X, "HYB_MGGA_X_M05_2X" }, + + // KEDFs + { ExchCXX::Kernel::PC07_K, "MGGA_K_PC07" }, + { ExchCXX::Kernel::PC07OPT_K, "MGGA_K_PC07_OPT" }, + }; + + std::once_flag g_exchcxx_init_once; + + inline void ensure_exchcxx_initialized(ExchCXX::Spin spin) { + std::call_once(g_exchcxx_init_once, [spin]{ ExchCXX::initialize(spin); }); + //g_exchcxx_users.fetch_add(1, std::memory_order_relaxed); + } + + inline void maybe_finalize_exchcxx() { + // Usually safer to never finalize until process exit. + // If you do want refcounted finalize, uncomment: + // if(g_exchcxx_users.fetch_sub(1, std::memory_order_relaxed) == 1) + // ExchCXX::finalize(); + } + + static bool is_composite_or_hybrid_name(const std::string& s_upper){ + // Heuristics: hybrids or explicit XC combos are not single kernels + if(s_upper.find("HYB_") != std::string::npos) return true; + if(s_upper.find("_XC_") != std::string::npos) return true; // exchange+correlation in one label + // Also common composites: e.g., "B88+LYP", "PBE0", "SCAN-RVV10", etc. + if(s_upper.find('+') != std::string::npos) return true; + if(s_upper.find("RVV10") != std::string::npos) return true; + if(s_upper.find("VV10") != std::string::npos) return true; + if(s_upper.find("D3") != std::string::npos) return true; + if(s_upper.find("D4") != std::string::npos) return true; + if(s_upper.find("DISP") != std::string::npos) return true; + if(s_upper.find("WB97") != std::string::npos) return true; + if(s_upper.find("CAM") != std::string::npos) return true; + return false; + } + + // case-insensitive equality + static bool iequals(const std::string& a, const std::string& b){ + if(a.size() != b.size()) return false; + for(size_t i=0;i Kernel (case-insensitive match) + static std::optional + kernel_from_libxc_name(const std::string& libxc_name){ + for(const auto& kv : libxc_kernel_map){ + if(iequals(kv.second, libxc_name)) return kv.first; + } + return std::nullopt; + } + + // True for Laplacian-requiring single-kernel variants (SCANL, R2SCANL, etc.) + static bool libxc_name_needs_lapl(const std::string& libxc_name_upper){ + return (libxc_name_upper.find("_R2SCANL") != std::string::npos) || + (libxc_name_upper.find("_SCANL") != std::string::npos); + } + +} // namespace detail + + +/* ---------------- Version / reference ---------------- */ +static const char* kRef = "ExchCXX GPU shim (libxc ABI)"; +static const char* kDOI = ""; +static const char* kKey = "ExchCXX"; +static const char* kVers = "ExchCXX-SYCL 1.0"; + +extern "C" { +const char *xc_reference(void) { return kRef; } +const char *xc_reference_doi(void) { return kDOI; } +const char *xc_reference_key(void) { return kKey; } +void xc_version(int *maj,int *min,int *mic){ if(maj) *maj=1; if(min)*min=0; if(mic)*mic=0; } +const char *xc_version_string(void) { return kVers; } +} + + +static const std::map libxc_id_to_name = { + // --- LDA Exchange --- + {1, "LDA_X"}, + {600, "LDA_X_1D_EXPONENTIAL"}, + {21, "LDA_X_1D_SOFT"}, + {19, "LDA_X_2D"}, + {546, "LDA_X_ERF"}, + {549, "LDA_X_RAE"}, + {532, "LDA_X_REL"}, + {692, "LDA_X_SLOC"}, + {641, "LDA_X_YUKAWA"}, + + // --- LDA Correlation --- + {18, "LDA_C_1D_CSC"}, + {26, "LDA_C_1D_LOOS"}, + {15, "LDA_C_2D_AMGB"}, + {16, "LDA_C_2D_PRM"}, + {552, "LDA_C_BR78"}, + {287, "LDA_C_CHACHIYO"}, + {307, "LDA_C_CHACHIYO_MOD"}, + {328, "LDA_C_EPC17"}, + {329, "LDA_C_EPC17_2"}, + {330, "LDA_C_EPC18_1"}, + {331, "LDA_C_EPC18_2"}, + {578, "LDA_C_GK72"}, + {5, "LDA_C_GL"}, + {24, "LDA_C_GOMBAS"}, + {4, "LDA_C_HL"}, + {579, "LDA_C_KARASIEV"}, + {308, "LDA_C_KARASIEV_MOD"}, + {551, "LDA_C_MCWEENY"}, + {22, "LDA_C_ML1"}, + {23, "LDA_C_ML2"}, + {14, "LDA_C_OB_PW"}, + {11, "LDA_C_OB_PZ"}, + {574, "LDA_C_OW"}, + {573, "LDA_C_OW_LYP"}, + {554, "LDA_C_PK09"}, + {590, "LDA_C_PMGB06"}, + {12, "LDA_C_PW"}, + {654, "LDA_C_PW_ERF"}, + {13, "LDA_C_PW_MOD"}, + {25, "LDA_C_PW_RPA"}, + {9, "LDA_C_PZ"}, + {10, "LDA_C_PZ_MOD"}, + {27, "LDA_C_RC04"}, + {3, "LDA_C_RPA"}, + {684, "LDA_C_RPW92"}, + {683, "LDA_C_UPW92"}, + {17, "LDA_C_VBH"}, + {7, "LDA_C_VWN"}, + {28, "LDA_C_VWN_1"}, + {29, "LDA_C_VWN_2"}, + {30, "LDA_C_VWN_3"}, + {31, "LDA_C_VWN_4"}, + {8, "LDA_C_VWN_RPA"}, + {317, "LDA_C_W20"}, + {2, "LDA_C_WIGNER"}, + {6, "LDA_C_XALPHA"}, + + // --- LDA Exchange–Correlation --- + {536, "LDA_XC_1D_EHWLRG_1"}, + {537, "LDA_XC_1D_EHWLRG_2"}, + {538, "LDA_XC_1D_EHWLRG_3"}, + {318, "LDA_XC_CORRKSDT"}, + {577, "LDA_XC_GDSMFB"}, + {259, "LDA_XC_KSDT"}, + {547, "LDA_XC_LP_A"}, + {548, "LDA_XC_LP_B"}, + {20, "LDA_XC_TETER93"}, + {599, "LDA_XC_TIH"}, + {43, "LDA_XC_ZLP"}, + + // --- LDA kinetic --- + {51, "LDA_K_LP"}, + {580, "LDA_K_LP96"}, + {50, "LDA_K_TF"}, + {550, "LDA_K_ZLP"}, + + // --- Hybrid LDA exchange --- + {653, "HYB_LDA_X_ERF"}, + + // --- Hybrid LDA exchange-correlation --- + {588, "HYB_LDA_XC_BN05"}, + {178, "HYB_LDA_XC_CAM_LDA0"}, + {177, "HYB_LDA_XC_LDA0"}, + + // --- GGA exchange --- + {128, "GGA_X_2D_B86"}, + {124, "GGA_X_2D_B86_MGC"}, + {127, "GGA_X_2D_B88"}, + {129, "GGA_X_2D_PBE"}, + {192, "GGA_X_AIRY"}, + {56, "GGA_X_AK13"}, + {120, "GGA_X_AM05"}, + {184, "GGA_X_APBE"}, + {103, "GGA_X_B86"}, + {105, "GGA_X_B86_MGC"}, + {41, "GGA_X_B86_R"}, + {106, "GGA_X_B88"}, + {179, "GGA_X_B88_6311G"}, + {570, "GGA_X_B88M"}, + {125, "GGA_X_BAYESIAN"}, + {38, "GGA_X_BCGP"}, + {285, "GGA_X_BEEFVDW"}, + {338, "GGA_X_BKL1"}, + {339, "GGA_X_BKL2"}, + {98, "GGA_X_BPCCAC"}, + {158, "GGA_X_C09X"}, + {270, "GGA_X_CAP"}, + {298, "GGA_X_CHACHIYO"}, + {111, "GGA_X_DK87_R1"}, + {112, "GGA_X_DK87_R2"}, + {271, "GGA_X_EB88"}, + {215, "GGA_X_ECMV92"}, + {35, "GGA_X_EV93"}, + {604, "GGA_X_FD_LB94"}, + {605, "GGA_X_FD_REVLB94"}, + {114, "GGA_X_FT97_A"}, + {115, "GGA_X_FT97_B"}, + {107, "GGA_X_G96"}, + {32, "GGA_X_GAM"}, + {535, "GGA_X_GG99"}, + {34, "GGA_X_HCTH_A"}, + {527, "GGA_X_HJS_B88"}, + {46, "GGA_X_HJS_B88_V2"}, + {528, "GGA_X_HJS_B97X"}, + {525, "GGA_X_HJS_PBE"}, + {526, "GGA_X_HJS_PBE_SOL"}, + {191, "GGA_X_HTBS"}, + {529, "GGA_X_ITYH"}, + {622, "GGA_X_ITYH_OPTX"}, + {623, "GGA_X_ITYH_PBE"}, + {544, "GGA_X_KGG99"}, + {145, "GGA_X_KT1"}, + {193, "GGA_X_LAG"}, + {44, "GGA_X_LAMBDA_CH_N"}, + {45, "GGA_X_LAMBDA_LO_N"}, + {40, "GGA_X_LAMBDA_OC2_N"}, + {160, "GGA_X_LB"}, + {182, "GGA_X_LBM"}, + {113, "GGA_X_LG93"}, + {168, "GGA_X_LSPBE"}, + {169, "GGA_X_LSRPBE"}, + {58, "GGA_X_LV_RPW86"}, + {149, "GGA_X_MB88"}, + {122, "GGA_X_MPBE"}, + {119, "GGA_X_MPW91"}, + {82, "GGA_X_N12"}, + {180, "GGA_X_NCAP"}, + {324, "GGA_X_NCAPR"}, + {183, "GGA_X_OL2"}, + {171, "GGA_X_OPTB86B_VDW"}, + {139, "GGA_X_OPTB88_VDW"}, + {141, "GGA_X_OPTPBE_VDW"}, + {110, "GGA_X_OPTX"}, + {101, "GGA_X_PBE"}, + {655, "GGA_X_PBE_ERF_GWS"}, + {321, "GGA_X_PBE_GAUSSIAN"}, + {126, "GGA_X_PBE_JSJR"}, + {320, "GGA_X_PBE_MOD"}, + {49, "GGA_X_PBE_MOL"}, + {102, "GGA_X_PBE_R"}, + {116, "GGA_X_PBE_SOL"}, + {59, "GGA_X_PBE_TCA"}, + {121, "GGA_X_PBEA"}, + {265, "GGA_X_PBEFE"}, + {60, "GGA_X_PBEINT"}, + {140, "GGA_X_PBEK1_VDW"}, + {539, "GGA_X_PBEPOW"}, + {291, "GGA_X_PBETRANS"}, + {108, "GGA_X_PW86"}, + {109, "GGA_X_PW91"}, + {316, "GGA_X_PW91_MOD"}, + {734, "GGA_X_Q1D"}, + {48, "GGA_X_Q2D"}, + {312, "GGA_X_REVSSB_D"}, + {142, "GGA_X_RGE2"}, + {117, "GGA_X_RPBE"}, + {144, "GGA_X_RPW86"}, + {495, "GGA_X_S12G"}, + {530, "GGA_X_SFAT"}, + {601, "GGA_X_SFAT_PBE"}, + {533, "GGA_X_SG4"}, + {150, "GGA_X_SOGGA"}, + {151, "GGA_X_SOGGA11"}, + {91, "GGA_X_SSB"}, + {92, "GGA_X_SSB_D"}, + {90, "GGA_X_SSB_SW"}, + {68, "GGA_X_VMT84_GE"}, + {69, "GGA_X_VMT84_PBE"}, + {70, "GGA_X_VMT_GE"}, + {71, "GGA_X_VMT_PBE"}, + {118, "GGA_X_WC"}, + {524, "GGA_X_WPBEH"}, + {123, "GGA_X_XPBE"}, + + // --- GGA correlation --- + {39, "GGA_C_ACGGA"}, + {176, "GGA_C_ACGGAP"}, + {135, "GGA_C_AM05"}, + {186, "GGA_C_APBE"}, + {280, "GGA_C_BMK"}, + {313, "GGA_C_CCDF"}, + {309, "GGA_C_CHACHIYO"}, + {565, "GGA_C_CS1"}, + {88, "GGA_C_FT97"}, + {33, "GGA_C_GAM"}, + {555, "GGA_C_GAPC"}, + {556, "GGA_C_GAPLOC"}, + {97, "GGA_C_HCTH_A"}, + {283, "GGA_C_HYB_TAU_HCTH"}, + {137, "GGA_C_LM"}, + {131, "GGA_C_LYP"}, + {624, "GGA_C_LYPR"}, + {712, "GGA_C_MGGAC"}, + {80, "GGA_C_N12"}, + {79, "GGA_C_N12_SX"}, + {87, "GGA_C_OP_B88"}, + {85, "GGA_C_OP_G96"}, + {86, "GGA_C_OP_PBE"}, + {262, "GGA_C_OP_PW91"}, + {84, "GGA_C_OP_XALPHA"}, + {200, "GGA_C_OPTC"}, + {132, "GGA_C_P86"}, + {217, "GGA_C_P86_FT"}, + {252, "GGA_C_P86VWN"}, + {253, "GGA_C_P86VWN_FT"}, + {130, "GGA_C_PBE"}, + {657, "GGA_C_PBE_ERF_GWS"}, + {322, "GGA_C_PBE_GAUSSIAN"}, + {138, "GGA_C_PBE_JRGX"}, + {272, "GGA_C_PBE_MOL"}, + {133, "GGA_C_PBE_SOL"}, + {216, "GGA_C_PBE_VWN"}, + {258, "GGA_C_PBEFE"}, + {62, "GGA_C_PBEINT"}, + {246, "GGA_C_PBELOC"}, + {134, "GGA_C_PW91"}, + {47, "GGA_C_Q2D"}, + {83, "GGA_C_REGTPSS"}, + {99, "GGA_C_REVTCA"}, + {143, "GGA_C_RGE2"}, + {553, "GGA_C_SCAN_E0"}, + {534, "GGA_C_SG4"}, + {152, "GGA_C_SOGGA11"}, + {159, "GGA_C_SOGGA11_X"}, + {89, "GGA_C_SPBE"}, + {281, "GGA_C_TAU_HCTH"}, + {100, "GGA_C_TCA"}, + {559, "GGA_C_TM_LYP"}, + {560, "GGA_C_TM_PBE"}, + {561, "GGA_C_W94"}, + {148, "GGA_C_WI"}, + {153, "GGA_C_WI0"}, + {147, "GGA_C_WL"}, + {136, "GGA_C_XPBE"}, + {61, "GGA_C_ZPBEINT"}, + {63, "GGA_C_ZPBESOL"}, + {557, "GGA_C_ZVPBEINT"}, + {606, "GGA_C_ZVPBELOC"}, + {558, "GGA_C_ZVPBESOL"}, + + // --- GGA exchange–correlation --- + {327, "GGA_XC_B97_3C"}, + {170, "GGA_XC_B97_D"}, + {96, "GGA_XC_B97_GGA1"}, + {286, "GGA_XC_BEEFVDW"}, + {165, "GGA_XC_EDF1"}, + {162, "GGA_XC_HCTH_120"}, + {163, "GGA_XC_HCTH_147"}, + {164, "GGA_XC_HCTH_407"}, + {93, "GGA_XC_HCTH_407P"}, + {161, "GGA_XC_HCTH_93"}, + {95, "GGA_XC_HCTH_P14"}, + {94, "GGA_XC_HCTH_P76"}, + {545, "GGA_XC_HLE16"}, + {167, "GGA_XC_KT1"}, + {146, "GGA_XC_KT2"}, + {587, "GGA_XC_KT3"}, + {194, "GGA_XC_MOHLYP"}, + {195, "GGA_XC_MOHLYP2"}, + {174, "GGA_XC_MPWLYP1W"}, + {181, "GGA_XC_NCAP"}, + {67, "GGA_XC_OBLYP_D"}, + {65, "GGA_XC_OPBE_D"}, + {66, "GGA_XC_OPWLYP_D"}, + {173, "GGA_XC_PBE1W"}, + {175, "GGA_XC_PBELYP1W"}, + {154, "GGA_XC_TH1"}, + {155, "GGA_XC_TH2"}, + {156, "GGA_XC_TH3"}, + {157, "GGA_XC_TH4"}, + {197, "GGA_XC_TH_FC"}, + {198, "GGA_XC_TH_FCFO"}, + {199, "GGA_XC_TH_FCO"}, + {196, "GGA_XC_TH_FL"}, + {255, "GGA_XC_VV10"}, + {166, "GGA_XC_XLYP"}, + + // --- GGA kinetic --- + {506, "GGA_K_ABSP1"}, + {507, "GGA_K_ABSP2"}, + {277, "GGA_K_ABSP3"}, + {278, "GGA_K_ABSP4"}, + {185, "GGA_K_APBE"}, + {54, "GGA_K_APBEINT"}, + {504, "GGA_K_BALTIN"}, + {516, "GGA_K_DK"}, + {520, "GGA_K_ERNZERHOF"}, + {597, "GGA_K_EXP4"}, + {514, "GGA_K_FR_B88"}, + {515, "GGA_K_FR_PW86"}, + {591, "GGA_K_GDS08"}, + {501, "GGA_K_GE2"}, + {592, "GGA_K_GHDS10"}, + {593, "GGA_K_GHDS10R"}, + {502, "GGA_K_GOLDEN"}, + {510, "GGA_K_GP85"}, + {508, "GGA_K_GR"}, + {521, "GGA_K_LC94"}, + {620, "GGA_K_LGAP"}, + {633, "GGA_K_LGAP_GE"}, + {505, "GGA_K_LIEB"}, + {613, "GGA_K_LKT"}, + {522, "GGA_K_LLP"}, + {509, "GGA_K_LUDENA"}, + {57, "GGA_K_MEYER"}, + {512, "GGA_K_OL1"}, + {513, "GGA_K_OL2"}, + {616, "GGA_K_PBE2"}, + {595, "GGA_K_PBE3"}, + {596, "GGA_K_PBE4"}, + {511, "GGA_K_PEARSON"}, + {517, "GGA_K_PERDEW"}, + {219, "GGA_K_PG1"}, + {218, "GGA_K_RATIONAL_P"}, + {55, "GGA_K_REVAPBE"}, + {53, "GGA_K_REVAPBEINT"}, + {52, "GGA_K_TFVW"}, + {635, "GGA_K_TFVW_OPT"}, + {523, "GGA_K_THAKKAR"}, + {594, "GGA_K_TKVLN"}, + {187, "GGA_K_TW1"}, + {188, "GGA_K_TW2"}, + {189, "GGA_K_TW3"}, + {190, "GGA_K_TW4"}, + {519, "GGA_K_VJKS"}, + {518, "GGA_K_VSK"}, + {619, "GGA_K_VT84F"}, + {500, "GGA_K_VW"}, + {503, "GGA_K_YT65"}, + + // --- HYB_GGA exchange --- + {646, "HYB_GGA_X_CAM_S12G"}, + {647, "HYB_GGA_X_CAM_S12H"}, + {81, "HYB_GGA_X_N12_SX"}, + {656, "HYB_GGA_X_PBE_ERF_GWS"}, + {496, "HYB_GGA_X_S12H"}, + {426, "HYB_GGA_X_SOGGA11_X"}, + + // --- HYB_GGA exchange–correlation --- + {607, "HYB_GGA_XC_APBE0"}, + {409, "HYB_GGA_XC_APF"}, + {416, "HYB_GGA_XC_B1LYP"}, + {417, "HYB_GGA_XC_B1PW91"}, + {412, "HYB_GGA_XC_B1WC"}, + {402, "HYB_GGA_XC_B3LYP"}, + {394, "HYB_GGA_XC_B3LYP3"}, + {475, "HYB_GGA_XC_B3LYP5"}, + {461, "HYB_GGA_XC_B3LYP_MCM1"}, + {462, "HYB_GGA_XC_B3LYP_MCM2"}, + {459, "HYB_GGA_XC_B3LYPS"}, + {403, "HYB_GGA_XC_B3P86"}, + {315, "HYB_GGA_XC_B3P86_NWCHEM"}, + {401, "HYB_GGA_XC_B3PW91"}, + {572, "HYB_GGA_XC_B5050LYP"}, + {407, "HYB_GGA_XC_B97"}, + {408, "HYB_GGA_XC_B97_1"}, + {266, "HYB_GGA_XC_B97_1P"}, + {410, "HYB_GGA_XC_B97_2"}, + {414, "HYB_GGA_XC_B97_3"}, + {413, "HYB_GGA_XC_B97_K"}, + {435, "HYB_GGA_XC_BHANDH"}, + {436, "HYB_GGA_XC_BHANDHLYP"}, + {499, "HYB_GGA_XC_BLYP35"}, + {433, "HYB_GGA_XC_CAM_B3LYP"}, + {395, "HYB_GGA_XC_CAM_O3LYP"}, + {681, "HYB_GGA_XC_CAM_PBEH"}, + {490, "HYB_GGA_XC_CAM_QTP_00"}, + {482, "HYB_GGA_XC_CAM_QTP_01"}, + {491, "HYB_GGA_XC_CAM_QTP_02"}, + {614, "HYB_GGA_XC_CAMH_B3LYP"}, + {470, "HYB_GGA_XC_CAMY_B3LYP"}, + {455, "HYB_GGA_XC_CAMY_BLYP"}, + {682, "HYB_GGA_XC_CAMY_PBEH"}, + {477, "HYB_GGA_XC_CAP0"}, + {390, "HYB_GGA_XC_CASE21"}, + {476, "HYB_GGA_XC_EDF2"}, + {608, "HYB_GGA_XC_HAPBE"}, + {314, "HYB_GGA_XC_HFLYP"}, + {431, "HYB_GGA_XC_HJS_B88"}, + {432, "HYB_GGA_XC_HJS_B97X"}, + {429, "HYB_GGA_XC_HJS_PBE"}, + {430, "HYB_GGA_XC_HJS_PBE_SOL"}, + {472, "HYB_GGA_XC_HPBEINT"}, + {427, "HYB_GGA_XC_HSE03"}, + {428, "HYB_GGA_XC_HSE06"}, + {479, "HYB_GGA_XC_HSE12"}, + {480, "HYB_GGA_XC_HSE12S"}, + {481, "HYB_GGA_XC_HSE_SOL"}, + {485, "HYB_GGA_XC_KMLYP"}, + {589, "HYB_GGA_XC_LB07"}, + {400, "HYB_GGA_XC_LC_BLYP"}, + {625, "HYB_GGA_XC_LC_BLYP_EA"}, + {639, "HYB_GGA_XC_LC_BLYPR"}, + {636, "HYB_GGA_XC_LC_BOP"}, + {637, "HYB_GGA_XC_LC_PBEOP"}, + {492, "HYB_GGA_XC_LC_QTP"}, + {469, "HYB_GGA_XC_LC_VV10"}, + {478, "HYB_GGA_XC_LC_WPBE"}, + {488, "HYB_GGA_XC_LC_WPBE08_WHS"}, + {486, "HYB_GGA_XC_LC_WPBE_WHS"}, + {487, "HYB_GGA_XC_LC_WPBEH_WHS"}, + {489, "HYB_GGA_XC_LC_WPBESOL_WHS"}, + {468, "HYB_GGA_XC_LCY_BLYP"}, + {467, "HYB_GGA_XC_LCY_PBE"}, + {473, "HYB_GGA_XC_LRC_WPBE"}, + {465, "HYB_GGA_XC_LRC_WPBEH"}, + {437, "HYB_GGA_XC_MB3LYP_RC04"}, + {640, "HYB_GGA_XC_MCAM_B3LYP"}, + {405, "HYB_GGA_XC_MPW1K"}, + {483, "HYB_GGA_XC_MPW1LYP"}, + {484, "HYB_GGA_XC_MPW1PBE"}, + {418, "HYB_GGA_XC_MPW1PW"}, + {419, "HYB_GGA_XC_MPW3LYP"}, + {415, "HYB_GGA_XC_MPW3PW"}, + {453, "HYB_GGA_XC_MPWLYP1M"}, + {404, "HYB_GGA_XC_O3LYP"}, + {386, "HYB_GGA_XC_OPB3LYP"}, + {456, "HYB_GGA_XC_PBE0_13"}, + {393, "HYB_GGA_XC_PBE38"}, + {290, "HYB_GGA_XC_PBE50"}, + {392, "HYB_GGA_XC_PBE_2X"}, + {273, "HYB_GGA_XC_PBE_MOL0"}, + {276, "HYB_GGA_XC_PBE_MOLB0"}, + {274, "HYB_GGA_XC_PBE_SOL0"}, + {275, "HYB_GGA_XC_PBEB0"}, + {406, "HYB_GGA_XC_PBEH"}, + {460, "HYB_GGA_XC_QTP17"}, + {610, "HYB_GGA_XC_RCAM_B3LYP"}, + {325, "HYB_GGA_XC_RELPBE0"}, + {454, "HYB_GGA_XC_REVB3LYP"}, + {420, "HYB_GGA_XC_SB98_1A"}, + {421, "HYB_GGA_XC_SB98_1B"}, + {422, "HYB_GGA_XC_SB98_1C"}, + {423, "HYB_GGA_XC_SB98_2A"}, + {424, "HYB_GGA_XC_SB98_2B"}, + {425, "HYB_GGA_XC_SB98_2C"}, + {434, "HYB_GGA_XC_TUNED_CAM_B3LYP"}, + {463, "HYB_GGA_XC_WB97"}, + {464, "HYB_GGA_XC_WB97X"}, + {471, "HYB_GGA_XC_WB97X_D"}, + {399, "HYB_GGA_XC_WB97X_D3"}, + {466, "HYB_GGA_XC_WB97X_V"}, + {611, "HYB_GGA_XC_WC04"}, + {615, "HYB_GGA_XC_WHPBE0"}, + {612, "HYB_GGA_XC_WP04"}, + {411, "HYB_GGA_XC_X3LYP"}, + + // --- MGGA exchange --- + {609, "MGGA_X_2D_JS17"}, + {210, "MGGA_X_2D_PRHG07"}, + {211, "MGGA_X_2D_PRHG07_PRP10"}, + {284, "MGGA_X_B00"}, + {207, "MGGA_X_BJ06"}, + {244, "MGGA_X_BLOC"}, + {206, "MGGA_X_BR89"}, + {214, "MGGA_X_BR89_1"}, + {586, "MGGA_X_BR89_EXPLICIT"}, + {602, "MGGA_X_BR89_EXPLICIT_1"}, + {686, "MGGA_X_EDMGGA"}, + {326, "MGGA_X_EEL"}, + {319, "MGGA_X_FT98"}, + {689, "MGGA_X_GDME_0"}, + {690, "MGGA_X_GDME_KOS"}, + {687, "MGGA_X_GDME_NV"}, + {691, "MGGA_X_GDME_VT"}, + {204, "MGGA_X_GVT4"}, + {575, "MGGA_X_GX"}, + {698, "MGGA_X_HLTA"}, + {256, "MGGA_X_JK"}, + {735, "MGGA_X_KTBM_0"}, + {736, "MGGA_X_KTBM_1"}, + {745, "MGGA_X_KTBM_10"}, + {746, "MGGA_X_KTBM_11"}, + {747, "MGGA_X_KTBM_12"}, + {748, "MGGA_X_KTBM_13"}, + {749, "MGGA_X_KTBM_14"}, + {750, "MGGA_X_KTBM_15"}, + {751, "MGGA_X_KTBM_16"}, + {752, "MGGA_X_KTBM_17"}, + {753, "MGGA_X_KTBM_18"}, + {754, "MGGA_X_KTBM_19"}, + {737, "MGGA_X_KTBM_2"}, + {755, "MGGA_X_KTBM_20"}, + {756, "MGGA_X_KTBM_21"}, + {757, "MGGA_X_KTBM_22"}, + {758, "MGGA_X_KTBM_23"}, + {759, "MGGA_X_KTBM_24"}, + {738, "MGGA_X_KTBM_3"}, + {739, "MGGA_X_KTBM_4"}, + {740, "MGGA_X_KTBM_5"}, + {741, "MGGA_X_KTBM_6"}, + {742, "MGGA_X_KTBM_7"}, + {743, "MGGA_X_KTBM_8"}, + {744, "MGGA_X_KTBM_9"}, + {760, "MGGA_X_KTBM_GAP"}, + {342, "MGGA_X_LAK"}, + {201, "MGGA_X_LTA"}, + {203, "MGGA_X_M06_L"}, + {226, "MGGA_X_M11_L"}, + {249, "MGGA_X_MBEEF"}, + {250, "MGGA_X_MBEEFVDW"}, + {716, "MGGA_X_MBR"}, + {696, "MGGA_X_MBRXC_BG"}, + {697, "MGGA_X_MBRXH_BG"}, + {644, "MGGA_X_MCML"}, + {711, "MGGA_X_MGGAC"}, + {230, "MGGA_X_MK00"}, + {243, "MGGA_X_MK00B"}, + {227, "MGGA_X_MN12_L"}, + {260, "MGGA_X_MN15_L"}, + {245, "MGGA_X_MODTPSS"}, + {221, "MGGA_X_MS0"}, + {222, "MGGA_X_MS1"}, + {223, "MGGA_X_MS2"}, + {228, "MGGA_X_MS2_REV"}, + {300, "MGGA_X_MS2B"}, + {301, "MGGA_X_MS2BS"}, + {765, "MGGA_X_MSB86BL"}, + {761, "MGGA_X_MSPBEL"}, + {763, "MGGA_X_MSRPBEL"}, + {724, "MGGA_X_MTASK"}, + {257, "MGGA_X_MVS"}, + {302, "MGGA_X_MVSB"}, + {303, "MGGA_X_MVSBS"}, + {576, "MGGA_X_PBE_GX"}, + {213, "MGGA_X_PKZB"}, + {497, "MGGA_X_R2SCAN"}, + {645, "MGGA_X_R2SCAN01"}, + {718, "MGGA_X_R2SCANL"}, + {650, "MGGA_X_R4SCAN"}, + {626, "MGGA_X_REGTM"}, + {603, "MGGA_X_REGTPSS"}, + {293, "MGGA_X_REVM06_L"}, + {581, "MGGA_X_REVSCAN"}, + {701, "MGGA_X_REVSCANL"}, + {693, "MGGA_X_REVTM"}, + {212, "MGGA_X_REVTPSS"}, + {688, "MGGA_X_RLDA"}, + {766, "MGGA_X_RMSB86BL"}, + {762, "MGGA_X_RMSPBEL"}, + {764, "MGGA_X_RMSRPBEL"}, + {209, "MGGA_X_RPP09"}, + {648, "MGGA_X_RPPSCAN"}, + {493, "MGGA_X_RSCAN"}, + {299, "MGGA_X_RTPSS"}, + {542, "MGGA_X_SA_TPSS"}, + {263, "MGGA_X_SCAN"}, + {700, "MGGA_X_SCANL"}, + {707, "MGGA_X_TASK"}, + {205, "MGGA_X_TAU_HCTH"}, + {208, "MGGA_X_TB09"}, + {225, "MGGA_X_TH"}, + {685, "MGGA_X_TLDA"}, + {540, "MGGA_X_TM"}, + {202, "MGGA_X_TPSS"}, + {651, "MGGA_X_VCML"}, + {541, "MGGA_X_VT84"}, + + // --- MGGA correlation --- + {571, "MGGA_C_B88"}, + {397, "MGGA_C_B94"}, + {240, "MGGA_C_BC95"}, + {387, "MGGA_C_CC"}, + {388, "MGGA_C_CCALDA"}, + {341, "MGGA_C_CF22D"}, + {72, "MGGA_C_CS"}, + {37, "MGGA_C_DLDF"}, + {699, "MGGA_C_HLTAPW"}, + {562, "MGGA_C_KCIS"}, + {638, "MGGA_C_KCISK"}, + {237, "MGGA_C_M05"}, + {238, "MGGA_C_M05_2X"}, + {235, "MGGA_C_M06"}, + {236, "MGGA_C_M06_2X"}, + {234, "MGGA_C_M06_HF"}, + {233, "MGGA_C_M06_L"}, + {311, "MGGA_C_M06_SX"}, + {78, "MGGA_C_M08_HX"}, + {77, "MGGA_C_M08_SO"}, + {76, "MGGA_C_M11"}, + {75, "MGGA_C_M11_L"}, + {74, "MGGA_C_MN12_L"}, + {73, "MGGA_C_MN12_SX"}, + {269, "MGGA_C_MN15"}, + {261, "MGGA_C_MN15_L"}, + {239, "MGGA_C_PKZB"}, + {498, "MGGA_C_R2SCAN"}, + {642, "MGGA_C_R2SCAN01"}, + {719, "MGGA_C_R2SCANL"}, + {306, "MGGA_C_REVM06"}, + {294, "MGGA_C_REVM06_L"}, + {172, "MGGA_C_REVM11"}, + {582, "MGGA_C_REVSCAN"}, + {585, "MGGA_C_REVSCAN_VV10"}, + {694, "MGGA_C_REVTM"}, + {241, "MGGA_C_REVTPSS"}, + {643, "MGGA_C_RMGGAC"}, + {649, "MGGA_C_RPPSCAN"}, + {391, "MGGA_C_RREGTM"}, + {494, "MGGA_C_RSCAN"}, + {267, "MGGA_C_SCAN"}, + {292, "MGGA_C_SCAN_RVV10"}, + {584, "MGGA_C_SCAN_VV10"}, + {702, "MGGA_C_SCANL"}, + {703, "MGGA_C_SCANL_RVV10"}, + {704, "MGGA_C_SCANL_VV10"}, + {251, "MGGA_C_TM"}, + {231, "MGGA_C_TPSS"}, + {323, "MGGA_C_TPSS_GAUSSIAN"}, + {247, "MGGA_C_TPSSLOC"}, + {232, "MGGA_C_VSXC"}, + + // --- MGGA exchange-correlation --- + {254, "MGGA_XC_B97M_V"}, + {229, "MGGA_XC_CC06"}, + {288, "MGGA_XC_HLE17"}, + {564, "MGGA_XC_LP90"}, + {64, "MGGA_XC_OTPSS_D"}, + {242, "MGGA_XC_TPSSLYP1W"}, + {652, "MGGA_XC_VCML_RVV10"}, + {42, "MGGA_XC_ZLP"}, + + // --- MGGA kinetic --- + {629, "MGGA_K_CSK1"}, + {630, "MGGA_K_CSK4"}, + {631, "MGGA_K_CSK_LOC1"}, + {632, "MGGA_K_CSK_LOC4"}, + {627, "MGGA_K_GEA2"}, + {628, "MGGA_K_GEA4"}, + {617, "MGGA_K_L04"}, + {618, "MGGA_K_L06"}, + {543, "MGGA_K_PC07"}, + {634, "MGGA_K_PC07_OPT"}, + {220, "MGGA_K_PGSL025"}, + {621, "MGGA_K_RDA"}, + + // --- HYB MGGA Exchange / XC --- + {279, "HYB_MGGA_X_BMK"}, + {340, "HYB_MGGA_X_CF22D"}, + {36, "HYB_MGGA_X_DLDF"}, + {705, "HYB_MGGA_X_JS18"}, + {438, "HYB_MGGA_X_M05"}, + {439, "HYB_MGGA_X_M05_2X"}, + {449, "HYB_MGGA_X_M06"}, + {450, "HYB_MGGA_X_M06_2X"}, + {444, "HYB_MGGA_X_M06_HF"}, + {310, "HYB_MGGA_X_M06_SX"}, + {295, "HYB_MGGA_X_M08_HX"}, + {296, "HYB_MGGA_X_M08_SO"}, + {297, "HYB_MGGA_X_M11"}, + {248, "HYB_MGGA_X_MN12_SX"}, + {268, "HYB_MGGA_X_MN15"}, + {224, "HYB_MGGA_X_MS2H"}, + {474, "HYB_MGGA_X_MVSH"}, + {706, "HYB_MGGA_X_PJS18"}, + {305, "HYB_MGGA_X_REVM06"}, + {304, "HYB_MGGA_X_REVM11"}, + {583, "HYB_MGGA_X_REVSCAN0"}, + {264, "HYB_MGGA_X_SCAN0"}, + {282, "HYB_MGGA_X_TAU_HCTH"}, + + // --- HYB MGGA XC --- + {563, "HYB_MGGA_XC_B0KCIS"}, + {441, "HYB_MGGA_XC_B86B95"}, + {440, "HYB_MGGA_XC_B88B95"}, + {398, "HYB_MGGA_XC_B94_HYB"}, + {598, "HYB_MGGA_XC_B98"}, + {443, "HYB_MGGA_XC_BB1K"}, + {389, "HYB_MGGA_XC_BR3P86"}, + {695, "HYB_MGGA_XC_EDMGGAH"}, + {658, "HYB_MGGA_XC_GAS22"}, + {720, "HYB_MGGA_XC_LC_TMLYP"}, + {445, "HYB_MGGA_XC_MPW1B95"}, + {566, "HYB_MGGA_XC_MPW1KCIS"}, + {446, "HYB_MGGA_XC_MPWB1K"}, + {567, "HYB_MGGA_XC_MPWKCIS1K"}, + {568, "HYB_MGGA_XC_PBE1KCIS"}, + {451, "HYB_MGGA_XC_PW6B95"}, + {442, "HYB_MGGA_XC_PW86B95"}, + {452, "HYB_MGGA_XC_PWB6K"}, + {660, "HYB_MGGA_XC_R2SCAN0"}, + {661, "HYB_MGGA_XC_R2SCAN50"}, + {659, "HYB_MGGA_XC_R2SCANH"}, + {458, "HYB_MGGA_XC_REVTPSSH"}, + {396, "HYB_MGGA_XC_TPSS0"}, + {569, "HYB_MGGA_XC_TPSS1KCIS"}, + {457, "HYB_MGGA_XC_TPSSH"}, + {531, "HYB_MGGA_XC_WB97M_V"}, + {447, "HYB_MGGA_XC_X1B95"}, + {448, "HYB_MGGA_XC_XB1K"} +}; + + +// ----- shared helpers ----- + +static ExchCXX::Kernel map_name_to_kernel(const std::string& in, + int* family_out, + bool* needs_lapl_out) +{ + using namespace detail; // for helpers & maps you showed + + auto set_family = [&](const std::string& u){ + int fam = -1; + // Exact family prefixes + if(u.rfind("LDA_", 0) == 0) fam = XC_FAMILY_LDA; + else if(u.rfind("GGA_", 0) == 0) fam = XC_FAMILY_GGA; + else if(u.rfind("MGGA_",0) == 0) fam = XC_FAMILY_MGGA; + // Hybrids: map to underlying family + else if(u.rfind("HYB_GGA_", 0) == 0) fam = XC_FAMILY_GGA; + else if(u.rfind("HYB_MGGA_",0) == 0) fam = XC_FAMILY_MGGA; + + if(family_out) *family_out = fam; + }; + + // 1) Normalize input + std::string s = in; + std::string su = to_upper(s); + + if(su.empty()){ + if(family_out) *family_out = -1; + if(needs_lapl_out) *needs_lapl_out = false; + throw std::invalid_argument("Empty functional name/id"); + } + + // 2) If user passed a numeric LibXC id, translate it to a LibXC name + bool is_numeric = std::all_of(su.begin(), su.end(), [](unsigned char c){ return std::isdigit(c); }); + if(is_numeric){ + int id = 0; + try { id = std::stoi(su); } catch(...) { /* fall through */ } + auto it = libxc_id_to_name.find(id); + if(it != libxc_id_to_name.end()){ + su = to_upper(it->second); + } else { + if(family_out) *family_out = -1; + if(needs_lapl_out) *needs_lapl_out = false; + std::ostringstream oss; + oss << "Unknown LibXC id: " << su; + throw std::invalid_argument(oss.str()); + } + } + + // 3) Reject obvious composite/hybrid-XC combos (this shim is for single kernels) + if(is_composite_or_hybrid_name(su)){ + if(family_out) *family_out = -1; + if(needs_lapl_out) *needs_lapl_out = false; + std::ostringstream oss; + oss << "Composite/Hybrid-XC label not supported as single kernel: " << su; + throw std::invalid_argument(oss.str()); + } + + // 4) Resolve LibXC name -> ExchCXX::Kernel using your libxc_kernel_map + auto optk = kernel_from_libxc_name(su); + if(!optk){ + if(family_out) *family_out = -1; + if(needs_lapl_out) *needs_lapl_out = false; + std::ostringstream oss; + oss << "No ExchCXX::Kernel mapping for LibXC name: " << su; + throw std::invalid_argument(oss.str()); + } + + // 5) Fill outs + set_family(su); + if(needs_lapl_out) *needs_lapl_out = libxc_name_needs_lapl(su); + + GDFT_TRACE("[gdft] map_name_to_kernel: %s -> %s | family=%d | needs_lapl=%d | kernel=%d\n", + in.c_str(), su.c_str(), + family_out ? *family_out : -999, + needs_lapl_out ? int(*needs_lapl_out) : 0, + static_cast(*optk)); + + return *optk; +} + +extern "C" int xc_functional_get_number(const char *name) { + if(!name) return 0; + int family=0; bool need_lapl=false; + auto k = map_name_to_kernel(name, &family, &need_lapl); + // You can’t encode need_lapl separately in this 16+16 scheme unless you dedicate a bit. + // If you need it later, recompute from the kernel at init time. + return (family << 16) | static_cast(k); +} +extern "C" const char *xc_functional_get_name(int number) { + // Look up the LibXC functional ID in our id->name mapping + auto it = libxc_id_to_name.find(number); + if (it != libxc_id_to_name.end()) { + return it->second.c_str(); + } + return nullptr; // LibXC ID not found +} +// LibXC IDs this shim can actually evaluate on the device. +// +// libxc_id_to_name is a full LibXC ID->name table, but ExchCXX's builtin +// backend implements only a subset of those functionals. Advertising the +// whole table makes gpu4pyscf/dft/libxc.py set XCfun.on_gpu = True for every +// functional; xc_func_init then returns 3 ("caller should fall back") and +// that caller raises RuntimeError instead, so the CPU path in +// numint.eval_xc_eff() is never reached. Reporting only the resolvable IDs +// keeps on_gpu honest, and everything else transparently uses PySCF's CPU +// libxc. +static const std::vector& supported_libxc_ids() { + static const std::vector ids = [] { + std::vector v; + for (const auto& kv : libxc_id_to_name) { + const auto name = detail::to_upper(kv.second); + if (detail::libxc_name_to_functional(name) || + detail::kernel_from_libxc_name(name)) { + v.push_back(kv.first); + } + } + return v; + }(); + return ids; +} + +// Highest derivative order the device path implements. +// +// GDFT_xc_lda/gga/mgga all evaluate up to fxc (order 2); kxc (order 3) has no +// ExchCXX device entry point. Without a way to ask, gpu4pyscf/dft/libxc.py +// requests order 3, gets a nonzero return, and raises RuntimeError. Exposing +// the limit lets the caller route those requests to PySCF's CPU libxc, the +// same way it already does for functionals this shim cannot evaluate. +extern "C" int xc_device_max_deriv_order(void) { return 2; } + +extern "C" int xc_number_of_functionals(void) { + return static_cast(supported_libxc_ids().size()); +} +extern "C" void xc_available_functional_numbers(int* list) { + if (!list) return; + int i = 0; + for (int id : supported_libxc_ids()) { + list[i++] = id; + } +} +/* ---------------- Shim state kept in xc_func_type::params ---------------- */ +struct ShimImpl { + ExchCXX::Spin spin{}; + int family = -1; + bool needs_lapl = false; + + // exactly one of these will be non-null + std::unique_ptr k; + std::unique_ptr f; + bool is_functional() const noexcept { return bool(f); } +}; + +static inline ShimImpl* get_impl(const xc_func_type* p){ + return reinterpret_cast(p ? p->params : nullptr); +} +template +static int with_xc(const xc_func_type* f, Fn&& fn) { + auto* impl = get_impl(f); + if(!impl) return 1; + if(impl->f) { fn(*impl->f); return 0; } + if(impl->k) { fn(*impl->k); return 0; } + return 1; // nothing to call +} +static inline int bad_args(...) { return 1; } + + +static void fill_dimensions(xc_dimensions* d, int family, int nspin, bool needs_lapl) { + std::memset(d, 0, sizeof(*d)); + + const bool unpol = (nspin == XC_UNPOLARIZED); + const int rho_dim = unpol ? 1 : 2; + const int sig_dim = unpol ? 1 : 3; // (aa, ab, bb) for pol + const int lap_dim = rho_dim; + const int tau_dim = rho_dim; + + // Inputs + d->rho = rho_dim; + d->sigma = sig_dim; + d->lapl = lap_dim; + d->tau = tau_dim; + + // 0th order + d->zk = 1; + + // 1st order + d->vrho = rho_dim; + if (family >= XC_FAMILY_GGA) d->vsigma = sig_dim; + if (family >= XC_FAMILY_MGGA) { + d->vtau = tau_dim; + d->vlapl = needs_lapl ? lap_dim : 0; + } + + // 2nd order - LDA + d->v2rho2 = unpol ? 1 : 3; // symmetric + + // 2nd order - GGA + if (family >= XC_FAMILY_GGA) { + d->v2rhosigma = rho_dim * sig_dim; // 2*3=6 for pol + d->v2sigma2 = unpol ? 1 : 6; // symmetric 3x3 + } + + // 2nd order - mGGA + if (family >= XC_FAMILY_MGGA) { + d->v2rholapl = needs_lapl ? (rho_dim * lap_dim) : 0; // 2*2=4 + d->v2rhotau = rho_dim * tau_dim; // 2*2=4 + d->v2sigmalapl = needs_lapl ? (sig_dim * lap_dim) : 0; // 3*2=6 + d->v2sigmatau = sig_dim * tau_dim; // 3*2=6 + d->v2lapl2 = needs_lapl ? (unpol ? 1 : 3) : 0; // symmetric + d->v2lapltau = needs_lapl ? (lap_dim * tau_dim) : 0; // 2*2=4 + d->v2tau2 = unpol ? 1 : 3; // symmetric + } +} + +extern "C" xc_func_type *xc_func_alloc(void) { + return (xc_func_type*) std::calloc(1, sizeof(xc_func_type)); +} + +extern "C" int xc_func_init(xc_func_type *p, int functional, int nspin) { + using detail::to_upper; + using detail::libxc_name_needs_lapl; + using detail::kernel_from_libxc_name; + using detail::libxc_name_to_functional; + + GDFT_TRACE("[gdft] xc_func_init: functional=%d nspin=%d\n", functional, nspin); + if (!p) return 1; + if (nspin != XC_UNPOLARIZED && nspin != XC_POLARIZED) return 2; + + auto impl = std::make_unique(); + impl->spin = (nspin == XC_UNPOLARIZED) + ? ExchCXX::Spin::Unpolarized + : ExchCXX::Spin::Polarized; + GDFT_TRACE("[DEBUG] GDFT_xc_gga: nspin=%d, spin=%d\n", + (impl->spin == ExchCXX::Spin::Polarized) ? 2 : 1, + int(impl->spin)); + + // Detect our packed form: (family<<16 | kernel_enum) + const int hi = (functional >> 16) & 0xFFFF; + const int lo = functional & 0xFFFF; + const bool looks_packed = + (hi == XC_FAMILY_LDA || hi == XC_FAMILY_GGA || hi == XC_FAMILY_MGGA); + + // If not packed, translate LibXC id -> name once + std::string libxc_name; + std::string name_upper; + + enum class Path { KernelPacked, FunctionalByName, KernelByName } path = Path::KernelByName; + + if (functional != 0 && looks_packed) { + path = Path::KernelPacked; + } else { + auto it_id = libxc_id_to_name.find(functional); + if (it_id == libxc_id_to_name.end()) { + std::fprintf(stderr, "ExchCXX: unknown LibXC ID %d\n", functional); + return 3; // let caller fallback + } + libxc_name = it_id->second; + name_upper = to_upper(libxc_name); + + // If this label is a composite/hybrid XC functional, use XCFunctional path + auto fopt = libxc_name_to_functional(name_upper); + if (fopt) { + path = Path::FunctionalByName; + } else { + path = Path::KernelByName; + } + } + + // Initialize ExchCXX device layer (spin-aware) before building kernels/functionals + detail::ensure_exchcxx_initialized(impl->spin); + + auto finalize = [&](int family, bool needs_lapl) { + impl->family = family; + impl->needs_lapl = needs_lapl; + p->nspin = nspin; + p->params = impl.release(); + p->params_size = sizeof(ShimImpl); + fill_dimensions(&p->dim, family, nspin, needs_lapl); + return 0; + }; + + try { + bool needs_lapl = false; + int family = XC_FAMILY_GGA; // corrected below + + // // ---- Custom hybrid override: B3LYP family ---- + // // Handle HYB_GGA_XC_B3LYP and its variants explicitly to control VWN flavor + // if (name_upper == "HYB_GGA_XC_B3LYP" || name_upper == "HYB_GGA_XC_B3LYP3" || name_upper == "HYB_GGA_XC_B3LYP5") { + // auto make_b3lyp_with = [&](ExchCXX::Kernel vwn_kernel) { + // ExchCXX::Backend backend = ExchCXX::Backend::builtin; + // ExchCXX::Spin spin = impl->spin; + + // const ExchCXX::HybCoeffs hyb_coefs = {0.20, 0.0, 0.0}; + // std::vector> terms = { + // {0.08, ExchCXX::XCKernel(backend, ExchCXX::Kernel::SlaterExchange, spin)}, + // {0.72, ExchCXX::XCKernel(backend, ExchCXX::Kernel::B88, spin)}, + // {0.19, ExchCXX::XCKernel(backend, vwn_kernel, spin)}, + // {0.81, ExchCXX::XCKernel(backend, ExchCXX::Kernel::LYP, spin)} + // }; + // return std::make_unique(terms, hyb_coefs); + // }; + + // // Map label → kernel flavor + // if (name_upper == "HYB_GGA_XC_B3LYP3") { + // std::cout << "exchcxx: kernel == VWN3 \n"; + // impl->f = make_b3lyp_with(ExchCXX::Kernel::VWN3); + // } else { + // // PySCF ≥ 2.3 default: VWN-RPA = VWN5 in this ExchCXX + // // LibXC's HYB_GGA_XC_B3LYP uses LDA_C_VWN_RPA (ID 8) = ExchCXX::Kernel::VWN5 + // std::cout << "exchcxx: kernel == VWN5 \n"; + // impl->f = make_b3lyp_with(ExchCXX::Kernel::VWN5); + // } + + // needs_lapl = false; + // family = XC_FAMILY_GGA; + // std::fprintf(stderr, "[gdft] Built XCFunctional (custom B3LYP) for '%s' (VWN5==RPA)\n", + // name_upper.c_str()); + + // return finalize(XC_FAMILY_GGA, needs_lapl); + // } + + + if (path == Path::KernelPacked) { + // Build from packed kernel enum + const auto kenum = static_cast(lo); + + // Heuristic laplacian flag from LibXC name if we can map back + auto it = detail::libxc_kernel_map.find(kenum); + if (it != detail::libxc_kernel_map.end()) + needs_lapl = libxc_name_needs_lapl(to_upper(it->second)); + + impl->k = std::make_unique(ExchCXX::Backend::builtin, kenum, impl->spin); + + family = impl->k->is_mgga() ? XC_FAMILY_MGGA + : impl->k->is_gga() ? XC_FAMILY_GGA + : XC_FAMILY_LDA; + + GDFT_TRACE("[gdft] 1. Built XCKernel: enum=%d family=%d spin=%d\n", + int(kenum), family, int(impl->spin)); + GDFT_TRACE("[gdft] 1. is_lda=%d is_gga=%d is_mgga=%d\n", + impl->k->is_lda(), impl->k->is_gga(), impl->k->is_mgga()); + + } else if (path == Path::FunctionalByName) { + // Build full XC functional (handles hybrids/composites like B3LYP, PBE, SCAN, …) + const auto fun_opt = libxc_name_to_functional(name_upper); + if (!fun_opt) { + std::fprintf(stderr, "ExchCXX: LibXC label '%s' not recognized as a composite functional\n", + name_upper.c_str()); + return 3; + } + const auto fun = *fun_opt; + needs_lapl = libxc_name_needs_lapl(name_upper); + + impl->f = std::make_unique(ExchCXX::Backend::builtin, fun, impl->spin); + + family = impl->f->is_mgga() ? XC_FAMILY_MGGA + : impl->f->is_gga() ? XC_FAMILY_GGA + : XC_FAMILY_LDA; + + GDFT_TRACE("[gdft] 2. Built XCFunctional: '%s' family=%d spin=%d\n", + name_upper.c_str(), family, int(impl->spin)); + GDFT_TRACE("[gdft] 2. f.is_lda=%d f.is_gga=%d f.is_mgga=%d\n", + impl->f->is_lda(), impl->f->is_gga(), impl->f->is_mgga()); + + } else { // KernelByName: single kernel by LibXC name + auto maybe_k = kernel_from_libxc_name(name_upper); + if (!maybe_k) { + std::fprintf(stderr, + "ExchCXX: LibXC name '%s' has no builtin single-kernel implementation\n", + name_upper.c_str()); + return 3; // let caller fallback + } + const auto kenum = *maybe_k; + needs_lapl = libxc_name_needs_lapl(name_upper); + + impl->k = std::make_unique( + ExchCXX::Backend::builtin, kenum, impl->spin); + + family = impl->k->is_mgga() ? XC_FAMILY_MGGA + : impl->k->is_gga() ? XC_FAMILY_GGA + : XC_FAMILY_LDA; + + GDFT_TRACE("[gdft] 3. Built XCKernel: name='%s' enum=%d family=%d spin=%d\n", + name_upper.c_str(), int(kenum), family, int(impl->spin)); + GDFT_TRACE("[gdft] 3. is_lda=%d is_gga=%d is_mgga=%d\n", + impl->k->is_lda(), impl->k->is_gga(), impl->k->is_mgga()); + } + + // Stash and finalize libxc-style handle + return finalize(family, needs_lapl); + + } catch (const std::exception& e) { + std::fprintf(stderr, "ExchCXX functional construction failed: %s\n", e.what()); + return 3; // signal caller to fallback + } +} + +extern "C" void xc_func_end(xc_func_type *p) { + if(!p) return; + auto *impl = get_impl(p); + if(impl) { + delete impl; + p->params = nullptr; + } +} + +extern "C" void xc_func_free(xc_func_type *p) { + // if(!p) return; + // xc_func_end(p); + // std::free(p); +} + +template +static inline int detect_order(const T* out) { + int order = -1; + if (out->zk != nullptr) order = 0; + if (out->vrho != nullptr) order = 1; + if (out->v2rho2 != nullptr) order = 2; + if (out->v3rho3 != nullptr) order = 3; + if (out->v4rho4 != nullptr) order = 4; + return order; +} + +static inline void zero_gga_out( + sycl::queue& q, + const xc_func_type* func, + const xc_gga_out_params* out, + std::size_t np, int order +) { + if(order >= 0) q.memset(out->zk, 0, sizeof(double)*np*func->dim.zk); + if(order >= 1) { + q.memset(out->vrho, 0, sizeof(double)*np*func->dim.vrho); + q.memset(out->vsigma, 0, sizeof(double)*np*func->dim.vsigma); // (sigma, lapl, tau) + } + if(order >= 2) { + q.memset(out->v2rho2, 0, sizeof(double)*np*func->dim.v2rho2); + q.memset(out->v2rhosigma, 0, sizeof(double)*np*func->dim.v2rhosigma); + q.memset(out->v2sigma2, 0, sizeof(double)*np*func->dim.v2sigma2); + } + if(order >= 3) { + q.memset(out->v3rho3, 0, sizeof(double)*np*func->dim.v3rho3); + q.memset(out->v3rho2sigma, 0, sizeof(double)*np*func->dim.v3rho2sigma); + q.memset(out->v3rhosigma2, 0, sizeof(double)*np*func->dim.v3rhosigma2); + q.memset(out->v3sigma3, 0, sizeof(double)*np*func->dim.v3sigma3); + } + if(order >= 4) { + q.memset(out->v4rho4, 0, sizeof(double)*np*func->dim.v4rho4); + q.memset(out->v4rho3sigma, 0, sizeof(double)*np*func->dim.v4rho3sigma); + q.memset(out->v4rho2sigma2, 0, sizeof(double)*np*func->dim.v4rho2sigma2); + q.memset(out->v4rhosigma3, 0, sizeof(double)*np*func->dim.v4rhosigma3); + q.memset(out->v4sigma4, 0, sizeof(double)*np*func->dim.v4sigma4); + } + q.wait(); +} + +extern "C" int GDFT_xc_lda( + void* stream_v, + const xc_func_type *func, int np, const double *rho, + xc_lda_out_params *out, xc_lda_out_params* /*buf*/ +){ + if(!func || !rho || !out || np <= 0) return bad_args(); + + const int order = detect_order(out); + if(order < 0) return 0; + if(order > 2){ + std::fprintf(stderr, "ExchCXX device: LDA order %d not implemented\n", order); + return 2; + } + + auto* stream = reinterpret_cast(stream_v); + double* eps = out->zk; + double* vrho = out->vrho; + double* v2rho2 = out->v2rho2; + + // 1) Derivatives first (some backends also write eps here; that’s fine, we’ll overwrite later) + if(order >= 1){ + int err = with_xc(func, [&](auto& xc){ + xc.eval_exc_vxc_device(np, rho, eps, vrho, stream); + }); + if(err) return err; + } + if(order >= 2){ + int err = with_xc(func, [&](auto& xc){ + xc.eval_vxc_fxc_device(np, rho, vrho, v2rho2, stream); + }); + if(err) return err; + } + + // 2) Energy last — this is the authoritative value Python expects to match CPU + if(eps){ + int err = with_xc(func, [&](auto& xc){ + xc.eval_exc_device(np, rho, eps, stream); + }); + if(err) return err; + } + + return 0; +} + + +// extern "C" int GDFT_xc_gga( +// void* stream_v, +// const xc_func_type *func, int np, const double *rho, const double *sigma, +// xc_gga_out_params *out, xc_gga_out_params* /*buf*/ +// ){ +// if(!func || !rho || !sigma || !out || np <= 0) return bad_args(); + +// const int order = detect_order(out); +// if(order < 0) return 0; +// if(order > 2){ +// std::fprintf(stderr, "ExchCXX device: GGA order %d not implemented\n", order); +// return 2; +// } + +// auto* stream = reinterpret_cast(stream_v); +// double* eps = out->zk; +// double* vrho = out->vrho; +// double* vsigma = out->vsigma; +// double* v2rho2 = out->v2rho2; +// double* v2rs = out->v2rhosigma; +// double* v2s2 = out->v2sigma2; + +// zero_gga_out(*stream, func, out, np, order); + +// // Step 2: Single evaluation — no redundant overwrites +// if(order == 0){ +// // Just energy +// int err = with_xc(func, [&](auto& xc){ +// xc.eval_exc_device(np, rho, sigma, eps, stream); +// }); +// if(err) return err; + +// } else if(order == 1){ +// // Energy + 1st derivatives — ONE call, no overwrite +// int err = with_xc(func, [&](auto& xc){ +// xc.eval_exc_vxc_device(np, rho, sigma, eps, vrho, vsigma, stream); +// }); +// if(err) return err; + +// } else if(order == 2){ +// // Energy + 1st derivatives +// int err = with_xc(func, [&](auto& xc){ +// xc.eval_exc_vxc_device(np, rho, sigma, eps, vrho, vsigma, stream); +// }); +// if(err) return err; + +// // 2nd derivatives ONLY — does NOT touch eps/vrho/vsigma +// err = with_xc(func, [&](auto& xc){ +// xc.eval_fxc_device(np, rho, sigma, v2rho2, v2rs, v2s2, stream); +// }); +// if(err) return err; +// } + +// return 0; +// } + +// extern "C" int GDFT_xc_gga( +// void* stream_v, +// const xc_func_type *func, int np, +// const double *rho, const double *sigma, +// xc_gga_out_params *out, xc_gga_out_params *buf /* workspace for mix */ +// ){ +// if(!func || !rho || !sigma || !out || np <= 0) return bad_args(); + +// const int order = detect_order(out); +// if(order < 0) return 0; +// if(order > 2){ +// std::fprintf(stderr, "ExchCXX device: GGA order %d not implemented\n", order); +// return 2; +// } + +// auto* qptr = reinterpret_cast(stream_v); +// auto& q = *qptr; +// const auto& dim = func->dim; + +// // ---------- Direct (non-mixed) path ---------- +// if(func->info && func->info->gga && !func->mix_coef){ +// // Zero outputs with correct sizes +// zero_gga_out(q, func, out, np, order); + +// // Derivatives first (these may write eps too; we’ll overwrite eps after) +// if(order >= 1){ +// int err = with_xc(func, [&](auto& xc){ +// return xc.eval_exc_vxc_device(np, rho, sigma, +// /*eps=*/out->zk, +// out->vrho, out->vsigma, qptr); +// }); +// if(err) return err; +// } +// if(order >= 2){ +// int err = with_xc(func, [&](auto& xc){ +// return xc.eval_vxc_fxc_device(np, rho, sigma, +// out->vrho, out->vsigma, +// out->v2rho2, out->v2rhosigma, out->v2sigma2, +// qptr); +// }); +// if(err) return err; +// } +// // Energy last — authoritative +// if(out->zk){ +// int err = with_xc(func, [&](auto& xc){ +// return xc.eval_exc_device(np, rho, sigma, out->zk, qptr); +// }); +// if(err) return err; +// } +// q.wait(); // ensure all kernels complete before returning +// return 0; +// } + +// // ---------- Mixed / hybrid path (e.g., B3LYP) ---------- +// if(!func->mix_coef){ +// // Defensive: libxc-like mixes should have mix_coef; if not, nothing to do +// // (CUDA code returns ierr=0 here) +// return 0; +// } + +// if(!buf){ +// std::fprintf(stderr, +// "ExchCXX device: GGA mix path requires 'buf' workspace (np=%d). " +// "Caller must provide a device-resident scratch buffer.\n", np); +// return 2; +// } + +// // 1) Zero the final accumulator +// zero_gga_out(q, func, out, np, order); + +// // 2) Loop over components, compute into buf, then out += coef * buf +// for(int i = 0; i < func->n_func_aux; ++i){ +// const double coef = func->mix_coef[i]; +// const xc_func_type* aux = func->func_aux[i]; + +// // Stage: clear buf for this component +// zero_gga_out(q, func, buf, np, order); + +// // Evaluate this component into buf +// if(order >= 1){ +// int err = with_xc(aux, [&](auto& xc){ +// return xc.eval_exc_vxc_device(np, rho, sigma, +// /*eps=*/buf->zk, +// buf->vrho, buf->vsigma, qptr); +// }); +// if(err) return err; +// } +// if(order >= 2){ +// int err = with_xc(aux, [&](auto& xc){ +// return xc.eval_vxc_fxc_device(np, rho, sigma, +// buf->vrho, buf->vsigma, +// buf->v2rho2, buf->v2rhosigma, buf->v2sigma2, +// qptr); +// }); +// if(err) return err; +// } +// // Energy last for this component +// if(buf->zk){ +// int err = with_xc(aux, [&](auto& xc){ +// return xc.eval_exc_device(np, rho, sigma, buf->zk, qptr); +// }); +// if(err) return err; +// } + +// // out += coef * buf (per-field AXPY with correct dimensions) +// axpy_gga_out(q, dim, out, buf, coef, order, static_cast(np)); +// } + +// //q.wait(); // ensure accumulations are done +// return 0; +// } + +#if GDFT_EXCHCXX_TRACE +static void debug_dump_gga(sycl::queue* stream, const char* tag, + int np, const double* rho, const double* sigma, + const double* eps, const double* vrho, const double* vsigma) { + const int N = 5; // print first N points + std::vector h_rho(N), h_sig(N), h_eps(N), h_vrho(N), h_vsig(N); + + stream->memcpy(h_rho.data(), rho, N*sizeof(double)); + stream->memcpy(h_sig.data(), sigma, N*sizeof(double)); + if(eps) stream->memcpy(h_eps.data(), eps, N*sizeof(double)); + if(vrho) stream->memcpy(h_vrho.data(), vrho, N*sizeof(double)); + if(vsigma) stream->memcpy(h_vsig.data(), vsigma, N*sizeof(double)); + stream->wait(); + + GDFT_TRACE("\n[%s] np=%d, first %d points:\n", tag, np, N); + GDFT_TRACE("%6s %20s %20s %20s %20s %20s\n", + "pt", "rho", "sigma", "eps", "vrho", "vsigma"); + for(int i = 0; i < N; i++) { + GDFT_TRACE("%6d %20.12e %20.12e %20.12e %20.12e %20.12e\n", + i, h_rho[i], h_sig[i], + eps ? h_eps[i] : 0.0, + vrho ? h_vrho[i] : 0.0, + vsigma ? h_vsig[i] : 0.0); + } + + // Also print sums (copy all np values) + std::vector all_eps(np), all_vrho(np), all_vsig(np); + if(eps) { stream->memcpy(all_eps.data(), eps, np*sizeof(double)); } + if(vrho) { stream->memcpy(all_vrho.data(), vrho, np*sizeof(double)); } + if(vsigma) { stream->memcpy(all_vsig.data(), vsigma, np*sizeof(double)); } + stream->wait(); + + double sum_eps=0, sum_vrho=0, sum_vsig=0; + for(int i=0; i 2){ + std::fprintf(stderr, "ExchCXX device: GGA order %d not implemented\n", order); + return 2; + } + + auto* stream = reinterpret_cast(stream_v); + double* eps = out->zk; + double* vrho = out->vrho; + double* vsigma = out->vsigma; + double* v2rho2 = out->v2rho2; + double* v2rs = out->v2rhosigma; + double* v2s2 = out->v2sigma2; + + zero_gga_out(*stream, func, out, np, order); + +#if GDFT_EXCHCXX_TRACE + debug_dump_gga(stream, "EXCHCXX-INPUT", np, rho, sigma, nullptr, nullptr, nullptr); +#endif + + + if(order >= 1){ + GDFT_TRACE("[gdft] GDFT_xc_gga: eval_exc_vxc_device (order >= 1)\n"); + int err = with_xc(func, [&](auto& xc){ + xc.eval_exc_vxc_device(np, rho, sigma, eps, vrho, vsigma, stream); + }); + if(err) return err; + } + if(order >= 2){ + GDFT_TRACE("[gdft] GDFT_xc_gga: eval_vxc_fxc_device (order >= 2)\n"); + int err = with_xc(func, [&](auto& xc){ + xc.eval_vxc_fxc_device(np, rho, sigma, vrho, vsigma, v2rho2, v2rs, v2s2, stream); + }); + if(err) return err; + } + + if(eps){ + GDFT_TRACE("[gdft] GDFT_xc_gga: eval_exc_device\n"); + int err = with_xc(func, [&](auto& xc){ + xc.eval_exc_device(np, rho, sigma, eps, stream); + }); + if(err) return err; + } + +#if GDFT_EXCHCXX_TRACE + debug_dump_gga(stream, "EXCHCXX-OUTPUT", np, rho, sigma, eps, vrho, vsigma); +#endif + + return 0; +} + +extern "C" int GDFT_xc_mgga( + void* stream_v, + const xc_func_type *func, int np, + const double *rho, const double *sigma, const double *lapl, const double *tau, + xc_mgga_out_params *out, xc_mgga_out_params* /*buf*/ +){ + if(!func || !rho || !sigma || !tau || !out || np <= 0) return bad_args(); + + const int order = detect_order(out); + if(order < 0) return 0; + if(order > 2){ + std::fprintf(stderr, "ExchCXX device: mGGA order %d not implemented\n", order); + return 2; + } + + auto* impl = get_impl(func); + if(!impl) return 1; + + // If this functional doesn't need the Laplacian, pass nullptr for lapl and skip vlapl/its Hessians + const bool need_lapl = impl->needs_lapl; + const double* lapl_in = need_lapl ? lapl : nullptr; + + auto* stream = reinterpret_cast(stream_v); + + // 1st-order outputs + double* eps = out->zk; + double* vrho = out->vrho; + double* vsigma = out->vsigma; + double* vlapl = need_lapl ? out->vlapl : nullptr; + double* vtau = out->vtau; + + // 2nd-order outputs + double* v2rho2 = out->v2rho2; + double* v2rhosigma = out->v2rhosigma; + double* v2rholapl = need_lapl ? out->v2rholapl : nullptr; + double* v2rhotau = out->v2rhotau; + double* v2sigma2 = out->v2sigma2; + double* v2sigmalapl = need_lapl ? out->v2sigmalapl : nullptr; + double* v2sigmatau = out->v2sigmatau; + double* v2lapl2 = need_lapl ? out->v2lapl2 : nullptr; + double* v2lapltau = (need_lapl ? out->v2lapltau : nullptr); + double* v2tau2 = out->v2tau2; + + if(order >= 1){ + int err = with_xc(func, [&](auto& xc){ + xc.eval_exc_vxc_device(np, rho, sigma, lapl_in, tau, eps, vrho, vsigma, vlapl, vtau, stream); + }); + if(err) return err; + } + if(order >= 2){ + int err = with_xc(func, [&](auto& xc){ + xc.eval_vxc_fxc_device(np, rho, sigma, lapl_in, tau, + vrho, vsigma, vlapl, vtau, + v2rho2, v2rhosigma, v2rholapl, v2rhotau, + v2sigma2, v2sigmalapl, v2sigmatau, + v2lapl2, v2lapltau, v2tau2, stream); + }); + if(err) return err; + } + + if(eps){ + int err = with_xc(func, [&](auto& xc){ + xc.eval_exc_device(np, rho, sigma, lapl_in, tau, eps, stream); + }); + if(err) return err; + } + + return 0; +} diff --git a/gpu4pyscf/lib/gdft/exchcxx.h b/gpu4pyscf/lib/gdft/exchcxx.h new file mode 100644 index 000000000..a5f415ae7 --- /dev/null +++ b/gpu4pyscf/lib/gdft/exchcxx.h @@ -0,0 +1,249 @@ +/* + Copyright (C) 2006-2007 M.A.L. Marques + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. +*/ + +#ifndef _XC_H +#define _XC_H + +#ifdef __cplusplus +extern "C" { +#endif + +/* ---- Version / reference (minimal ABI) ---- */ +const char *xc_reference(void); +const char *xc_reference_doi(void); +const char *xc_reference_key(void); +void xc_version(int *major, int *minor, int *micro); +const char *xc_version_string(void); + +/* ---- Common constants (kept identical to original) ---- */ +#include + +#define XC_UNPOLARIZED 1 +#define XC_POLARIZED 2 + +#define XC_NON_RELATIVISTIC 0 +#define XC_RELATIVISTIC 1 + +#define XC_EXCHANGE 0 +#define XC_CORRELATION 1 +#define XC_EXCHANGE_CORRELATION 2 +#define XC_KINETIC 3 + +#define XC_FAMILY_UNKNOWN -1 +#define XC_FAMILY_LDA 1 +#define XC_FAMILY_GGA 2 +#define XC_FAMILY_MGGA 4 +#define XC_FAMILY_LCA 8 +#define XC_FAMILY_OEP 16 +#define XC_FAMILY_HYB_GGA 32 +#define XC_FAMILY_HYB_MGGA 64 +#define XC_FAMILY_HYB_LDA 128 + +#define XC_FLAGS_HAVE_EXC (1 << 0) +#define XC_FLAGS_HAVE_VXC (1 << 1) +#define XC_FLAGS_HAVE_FXC (1 << 2) +#define XC_FLAGS_HAVE_KXC (1 << 3) +#define XC_FLAGS_HAVE_LXC (1 << 4) +#define XC_FLAGS_1D (1 << 5) +#define XC_FLAGS_2D (1 << 6) +#define XC_FLAGS_3D (1 << 7) +#define XC_FLAGS_HYB_CAM (1 << 8) +#define XC_FLAGS_HYB_CAMY (1 << 9) +#define XC_FLAGS_VV10 (1 << 10) +#define XC_FLAGS_HYB_LC (1 << 11) +#define XC_FLAGS_HYB_LCY (1 << 12) +#define XC_FLAGS_STABLE (1 << 13) +#define XC_FLAGS_DEVELOPMENT (1 << 14) +#define XC_FLAGS_NEEDS_LAPLACIAN (1 << 15) +#define XC_FLAGS_NEEDS_TAU (1 << 16) +#define XC_FLAGS_HAVE_ALL (XC_FLAGS_HAVE_EXC | XC_FLAGS_HAVE_VXC | XC_FLAGS_HAVE_FXC | XC_FLAGS_HAVE_KXC | XC_FLAGS_HAVE_LXC) + +#define XC_EXT_PARAMS_DEFAULT -999998888 +#define XC_MAX_REFERENCES 5 + +/* ---- Output param macros kept for ABI (unused in this shim) ---- */ +#define XC_NOARG +#define XC_COMMA , + +#define LDA_OUT_PARAMS_NO_EXC(P1_, P2_) \ + P1_ P2_ ## vrho \ + P1_ P2_ ## v2rho2 \ + P1_ P2_ ## v3rho3 \ + P1_ P2_ ## v4rho4 + +#define GGA_OUT_PARAMS_NO_EXC(P1_, P2_) \ + P1_ P2_ ## vrho P1_ P2_ ## vsigma \ + P1_ P2_ ## v2rho2 P1_ P2_ ## v2rhosigma \ + P1_ P2_ ## v2sigma2 \ + P1_ P2_ ## v3rho3 P1_ P2_ ## v3rho2sigma \ + P1_ P2_ ## v3rhosigma2 P1_ P2_ ## v3sigma3 \ + P1_ P2_ ## v4rho4 P1_ P2_ ## v4rho3sigma \ + P1_ P2_ ## v4rho2sigma2 P1_ P2_ ## v4rhosigma3 \ + P1_ P2_ ## v4sigma4 + +#define MGGA_OUT_PARAMS_NO_EXC(P1_, P2_) \ + P1_ P2_ ## vrho P1_ P2_ ## vsigma \ + P1_ P2_ ## vlapl P1_ P2_ ## vtau \ + P1_ P2_ ## v2rho2 P1_ P2_ ## v2rhosigma \ + P1_ P2_ ## v2rholapl P1_ P2_ ## v2rhotau \ + P1_ P2_ ## v2sigma2 P1_ P2_ ## v2sigmalapl \ + P1_ P2_ ## v2sigmatau P1_ P2_ ## v2lapl2 \ + P1_ P2_ ## v2lapltau P1_ P2_ ## v2tau2 \ + P1_ P2_ ## v3rho3 P1_ P2_ ## v3rho2sigma \ + P1_ P2_ ## v3rho2lapl P1_ P2_ ## v3rho2tau \ + P1_ P2_ ## v3rhosigma2 P1_ P2_ ## v3rhosigmalapl \ + P1_ P2_ ## v3rhosigmatau P1_ P2_ ## v3rholapl2 \ + P1_ P2_ ## v3rholapltau P1_ P2_ ## v3rhotau2 \ + P1_ P2_ ## v3sigma3 P1_ P2_ ## v3sigma2lapl \ + P1_ P2_ ## v3sigma2tau P1_ P2_ ## v3sigmalapl2 \ + P1_ P2_ ## v3sigmalapltau P1_ P2_ ## v3sigmatau2 \ + P1_ P2_ ## v3lapl3 P1_ P2_ ## v3lapl2tau \ + P1_ P2_ ## v3lapltau2 P1_ P2_ ## v3tau3 \ + P1_ P2_ ## v4rho4 P1_ P2_ ## v4rho3sigma \ + P1_ P2_ ## v4rho3lapl P1_ P2_ ## v4rho3tau \ + P1_ P2_ ## v4rho2sigma2 P1_ P2_ ## v4rho2sigmalapl \ + P1_ P2_ ## v4rho2sigmatau P1_ P2_ ## v4rho2lapl2 \ + P1_ P2_ ## v4rho2lapltau P1_ P2_ ## v4rho2tau2 \ + P1_ P2_ ## v4rhosigma3 P1_ P2_ ## v4rhosigma2lapl \ + P1_ P2_ ## v4rhosigma2tau P1_ P2_ ## v4rhosigmalapl2 \ + P1_ P2_ ## v4rhosigmalapltau P1_ P2_ ## v4rhosigmatau2 \ + P1_ P2_ ## v4rholapl3 P1_ P2_ ## v4rholapl2tau \ + P1_ P2_ ## v4rholapltau2 P1_ P2_ ## v4rhotau3 \ + P1_ P2_ ## v4sigma4 P1_ P2_ ## v4sigma3lapl \ + P1_ P2_ ## v4sigma3tau P1_ P2_ ## v4sigma2lapl2 \ + P1_ P2_ ## v4sigma2lapltau P1_ P2_ ## v4sigma2tau2 \ + P1_ P2_ ## v4sigmalapl3 P1_ P2_ ## v4sigmalapl2tau \ + P1_ P2_ ## v4sigmalapltau2 P1_ P2_ ## v4sigmatau3 \ + P1_ P2_ ## v4lapl4 P1_ P2_ ## v4lapl3tau \ + P1_ P2_ ## v4lapl2tau2 P1_ P2_ ## v4lapltau3 \ + P1_ P2_ ## v4tau4 + +/* ---- C structs (match your Python ctypes exactly) ---- */ +typedef struct{ + const char *ref, *doi, *bibtex, *key; +} func_reference_type; + +typedef struct{ + int n; + const char **names; + const char **descriptions; + const double *values; + void (*set)(struct xc_func_type *p, const double *ext_params); +} func_params_type; + +typedef struct { + int rho, sigma, lapl, tau; + int zk MGGA_OUT_PARAMS_NO_EXC(XC_COMMA, ); +} xc_dimensions; + +typedef struct xc_func_info_type { + int number; + int kind; + const char *name; + int family; + func_reference_type *refs[XC_MAX_REFERENCES]; + int flags; + double dens_threshold; + func_params_type ext_params; + void (*init)(struct xc_func_type *p); + void (*end) (struct xc_func_type *p); + const void *lda; /* unused by shim */ + const void *gga; /* unused by shim */ + const void *mgga; /* unused by shim */ +} xc_func_info_type; + +typedef struct xc_func_type{ + const xc_func_info_type *info; + int nspin; + int n_func_aux; + struct xc_func_type **func_aux; + double *mix_coef; + + double cam_omega, cam_alpha, cam_beta; + double nlc_b, nlc_C; + + xc_dimensions dim; + + double *ext_params; + void *params; /* shim stores internal ExchCXX state here */ + int params_size; + + double dens_threshold; + double zeta_threshold; + double sigma_threshold; + double tau_threshold; +} xc_func_type; + +/* ---- Output parameter containers ---- */ +typedef struct { + double *zk; + double *vrho; + double *v2rho2; + double *v3rho3; + double *v4rho4; +} xc_lda_out_params; + +typedef struct { + double *zk; + double *vrho, *vsigma; + double *v2rho2, *v2rhosigma, *v2sigma2; + double *v3rho3, *v3rho2sigma, *v3rhosigma2, *v3sigma3; + double *v4rho4, *v4rho3sigma, *v4rho2sigma2, *v4rhosigma3, *v4sigma4; +} xc_gga_out_params; + +typedef struct { + double *zk; + double *vrho, *vsigma, *vlapl, *vtau; + double *v2rho2, *v2rhosigma, *v2rholapl, *v2rhotau, *v2sigma2; + double *v2sigmalapl, *v2sigmatau, *v2lapl2, *v2lapltau, *v2tau2; + double *v3rho3, *v3rho2sigma, *v3rho2lapl, *v3rho2tau, *v3rhosigma2; + double *v3rhosigmalapl, *v3rhosigmatau, *v3rholapl2, *v3rholapltau; + double *v3rhotau2, *v3sigma3, *v3sigma2lapl, *v3sigma2tau; + double *v3sigmalapl2, *v3sigmalapltau, *v3sigmatau2, *v3lapl3; + double *v3lapl2tau, *v3lapltau2, *v3tau3; + double *v4rho4, *v4rho3sigma, *v4rho3lapl, *v4rho3tau, *v4rho2sigma2; + double *v4rho2sigmalapl, *v4rho2sigmatau, *v4rho2lapl2, *v4rho2lapltau; + double *v4rho2tau2, *v4rhosigma3, *v4rhosigma2lapl, *v4rhosigma2tau; + double *v4rhosigmalapl2, *v4rhosigmalapltau, *v4rhosigmatau2; + double *v4rholapl3, *v4rholapl2tau, *v4rholapltau2, *v4rhotau3; + double *v4sigma4, *v4sigma3lapl, *v4sigma3tau, *v4sigma2lapl2; + double *v4sigma2lapltau, *v4sigma2tau2, *v4sigmalapl3, *v4sigmalapl2tau; + double *v4sigmalapltau2, *v4sigmatau3, *v4lapl4, *v4lapl3tau; + double *v4lapl2tau2, *v4lapltau3, *v4tau4; +} xc_mgga_out_params; + +/* ---- Minimal API we implement ---- */ +xc_func_type *xc_func_alloc(void); +int xc_func_init(xc_func_type *p, int functional, int nspin); +void xc_func_end(xc_func_type *p); +void xc_func_free(xc_func_type *p); + +/* String <-> id helper (we implement a small table; extend as needed) */ +int xc_functional_get_number(const char *name); +const char *xc_functional_get_name(int number); +int xc_number_of_functionals(void); +void xc_available_functional_numbers(int *list); + +/* ---- Device entry points used by Python (unchanged ABI) ---- */ +int GDFT_xc_lda (void* stream, + const xc_func_type *func, int np, const double *rho, + xc_lda_out_params *out, xc_lda_out_params *buf); + +int GDFT_xc_gga (void* stream, + const xc_func_type *func, int np, const double *rho, const double *sigma, + xc_gga_out_params *out, xc_gga_out_params *buf); + +int GDFT_xc_mgga(void* stream, + const xc_func_type *func, int np, + const double *rho, const double *sigma, const double *lapl, const double *tau, + xc_mgga_out_params *out, xc_mgga_out_params *buf); + +#ifdef __cplusplus +} +#endif +#endif /* _XC_H */ diff --git a/gpu4pyscf/lib/gdft/gen_grids.cu b/gpu4pyscf/lib/gdft/gen_grids.cu index c534bd32f..27fa156ed 100644 --- a/gpu4pyscf/lib/gdft/gen_grids.cu +++ b/gpu4pyscf/lib/gdft/gen_grids.cu @@ -42,10 +42,38 @@ __global__ void GDFTgrid_weight_kernel(double *weight, const double *coords, const double *atm_coords, const double *a_factor, const int *atm_idx, const int ngrids, const int natm) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int tx = item.get_local_id(1); + int ty = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + + auto thread_block = item.get_group(); + using tile_t = double[TILE]; + tile_t& atom_xi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + tile_t& atom_yi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + tile_t& atom_zi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + tile_t& atom_xj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + tile_t& atom_yj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + tile_t& atom_zj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&a_smem)[if_radii_adjust ? (TILE*TILE) : 1] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&dij_smem)[TILE*TILE] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else int tx = threadIdx.x; int ty = threadIdx.y; + int blockIdx_x = blockIdx.x; + + __shared__ double atom_xi[TILE]; + __shared__ double atom_yi[TILE]; + __shared__ double atom_zi[TILE]; + __shared__ double atom_xj[TILE]; + __shared__ double atom_yj[TILE]; + __shared__ double atom_zj[TILE]; + __shared__ double a_smem[if_radii_adjust ? (TILE*TILE) : 1]; // CUDA doesn't allow zero-sized array + __shared__ double dij_smem[TILE*TILE]; +#endif int thread_id = ty * TILE + tx; - int grid_id = blockIdx.x * TILE*TILE + thread_id; + int grid_id = blockIdx_x * TILE*TILE + thread_id; double xg = 0.0; double yg = 0.0; double zg = 0.0; @@ -59,14 +87,6 @@ void GDFTgrid_weight_kernel(double *weight, const double *coords, const double * const double *atm_x = atm_coords; const double *atm_y = atm_x + natm; const double *atm_z = atm_y + natm; - __shared__ double atom_xi[TILE]; - __shared__ double atom_yi[TILE]; - __shared__ double atom_zi[TILE]; - __shared__ double atom_xj[TILE]; - __shared__ double atom_yj[TILE]; - __shared__ double atom_zj[TILE]; - __shared__ double a_smem[if_radii_adjust ? (TILE*TILE) : 1]; // CUDA doesn't allow zero-sized array - __shared__ double dij_smem[TILE*TILE]; double becke_self = 0.; double becke_sum = 0.; @@ -320,8 +340,14 @@ void GDFTgrid_weight_derivative_kernel(double* __restrict__ dwdG, const double* const double* __restrict__ Ar_distance, const double* __restrict__ PB, const double* __restrict__ invsumPB, const int ngrids, const int natm) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i_grid = item.get_global_id(1); + const int i_derivative_atom = item.get_global_id(0); +#else const int i_grid = blockIdx.x * blockDim.x + threadIdx.x; const int i_derivative_atom = blockIdx.y; +#endif if (i_grid >= ngrids || i_derivative_atom >= natm) return; const int i_associated_atom = atm_idx[i_grid]; @@ -393,7 +419,7 @@ typedef struct { double3 y; double3 z; } double9; -__device__ constexpr double9 identity_3 = { 1,0,0, 0,1,0, 0,0,1 }; +__device__ constexpr double9 identity_3 = { {1,0,0}, {0,1,0}, {0,0,1} }; __device__ double9 operator+(const double9& v1, const double9& v2) { return { v1.x + v2.x, v1.y + v2.y, v1.z + v2.z }; } __device__ double9 operator-(const double9& v1, const double9& v2) { return { v1.x - v2.x, v1.y - v2.y, v1.z - v2.z }; } __device__ double9 operator-(const double9& v) { return { -v.x, -v.y, -v.z }; } @@ -478,9 +504,16 @@ void GDFTgrid_weight_second_derivative_offdiagonal_kernel(double* __restrict__ d const int* __restrict__ atm_idx, const double* __restrict__ Ar_distance, const double* __restrict__ PB, const double* __restrict__ invsumPB, const int ngrids, const int natm) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<3>(); + const int i_grid = item.get_global_id(2); + const int i_atom_G = item.get_global_id(1); + const int i_atom_H = item.get_global_id(0); +#else const int i_grid = blockIdx.x * blockDim.x + threadIdx.x; const int i_atom_G = blockIdx.y * blockDim.y + threadIdx.y; const int i_atom_H = blockIdx.z * blockDim.z + threadIdx.z; +#endif if (i_grid >= ngrids || i_atom_G >= natm || i_atom_H >= natm) return; const int i_atom_A = atm_idx[i_grid]; @@ -509,7 +542,7 @@ void GDFTgrid_weight_second_derivative_offdiagonal_kernel(double* __restrict__ d const double P_H = PB[i_atom_H * ngrids + i_grid]; double3 dPH_dH = { 0.0, 0.0, 0.0 }; - double9 sum_d2PB_dGdH = { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 }; + double9 sum_d2PB_dGdH = { {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0} }; for (int i_atom_B = 0; i_atom_B < natm; i_atom_B++) { const double3 atom_B = { atm_coords[i_atom_B + 0 * natm], atm_coords[i_atom_B + 1 * natm], atm_coords[i_atom_B + 2 * natm] }; @@ -623,7 +656,7 @@ void GDFTgrid_weight_second_derivative_offdiagonal_kernel(double* __restrict__ d const double9 d2PA_dGdH = P_A * outer(dsAG_dG, dsAH_dH); const double sum_P_B_1 = invsumPB[i_grid]; - double9 d2wi_dGdH = { 0,0,0, 0,0,0, 0,0,0 }; + double9 d2wi_dGdH = { {0,0,0}, {0,0,0}, {0,0,0} }; d2wi_dGdH += sum_P_B_1 * d2PA_dGdH; d2wi_dGdH -= (sum_P_B_1 * sum_P_B_1) * outer(sum_dPB_dG, dPA_dH); d2wi_dGdH -= (sum_P_B_1 * sum_P_B_1) * outer(dPA_dG, sum_dPB_dH); @@ -650,8 +683,14 @@ void GDFTgrid_weight_second_derivative_diagonal_kernel(double* __restrict__ d2w_ const int* __restrict__ atm_idx, const double* __restrict__ Ar_distance, const double* __restrict__ PB, const double* __restrict__ invsumPB, const int ngrids, const int natm) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i_grid = item.get_global_id(1); + const int i_atom_G = item.get_global_id(0); +#else const int i_grid = blockIdx.x * blockDim.x + threadIdx.x; const int i_atom_G = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i_grid >= ngrids || i_atom_G >= natm) return; const int i_atom_A = atm_idx[i_grid]; @@ -670,8 +709,8 @@ void GDFTgrid_weight_second_derivative_diagonal_kernel(double* __restrict__ d2w_ const double P_G = PB[i_atom_G * ngrids + i_grid]; double3 dPG_dG = { 0.0, 0.0, 0.0 }; - double9 sum_d2PB_dG2 = { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 }; - double9 d2PG_dG2 = { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 }; + double9 sum_d2PB_dG2 = { {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0} }; + double9 d2PG_dG2 = { {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0} }; for (int i_atom_B = 0; i_atom_B < natm; i_atom_B++) { const double3 atom_B = { atm_coords[i_atom_B + 0 * natm], atm_coords[i_atom_B + 1 * natm], atm_coords[i_atom_B + 2 * natm] }; @@ -738,7 +777,7 @@ void GDFTgrid_weight_second_derivative_diagonal_kernel(double* __restrict__ d2w_ const double9 d2PA_dG2 = P_A * (dsdmu_dmu2dG2 + d2sdmu2_dmuAGdG_2); const double sum_P_B_1 = invsumPB[i_grid]; - double9 d2wi_dG2 = { 0,0,0, 0,0,0, 0,0,0 }; + double9 d2wi_dG2 = { {0,0,0}, {0,0,0}, {0,0,0} }; d2wi_dG2 += sum_P_B_1 * d2PA_dG2; d2wi_dG2 -= (sum_P_B_1 * sum_P_B_1) * outer(sum_dPB_dG, dPA_dG); d2wi_dG2 -= (sum_P_B_1 * sum_P_B_1) * outer(dPA_dG, sum_dPB_dG); @@ -764,8 +803,14 @@ void GDFTgrid_becke_eval_PB_kernel(double* __restrict__ PB, const double* __restrict__ a_factor, const double* __restrict__ inv_atom_distance, const double* __restrict__ Ar_distance, const int ngrids, const int natm) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i_grid = item.get_global_id(1); + const int i_atom_B = item.get_global_id(0); +#else const int i_grid = blockIdx.x * blockDim.x + threadIdx.x; const int i_atom_B = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i_grid >= ngrids || i_atom_B >= natm) return; @@ -789,7 +834,24 @@ void GDFTgrid_becke_eval_PB_kernel(double* __restrict__ PB, __global__ void GDFTgroup_grids_kernel(int* group_ids, const double* atom_coords, const double* coords, int natm, int ngrids){ +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int grid_id = item.get_global_id(0); + const int tx = item.get_local_id(0); + const int blockDim_x = item.get_local_range(0); + using tile_t = double[NATOM_PER_BLOCK]; + auto thread_block = item.get_group(); + tile_t& x_atom = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + tile_t& y_atom = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + tile_t& z_atom = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + const int tx = threadIdx.x; + const int blockDim_x = blockDim.x; + double __shared__ x_atom[NATOM_PER_BLOCK]; + double __shared__ y_atom[NATOM_PER_BLOCK]; + double __shared__ z_atom[NATOM_PER_BLOCK]; +#endif double xg = coords[grid_id]; double yg = coords[grid_id + ngrids]; @@ -797,11 +859,7 @@ void GDFTgroup_grids_kernel(int* group_ids, const double* atom_coords, const dou double r2min = 1e30; int idx = 0; - const int tx = threadIdx.x; - double __shared__ x_atom[NATOM_PER_BLOCK]; - double __shared__ y_atom[NATOM_PER_BLOCK]; - double __shared__ z_atom[NATOM_PER_BLOCK]; - for (int j = 0; j < natm; j+=blockDim.x){ + for (int j = 0; j < natm; j+=blockDim_x){ int atom_idx = j + tx; if (atom_idx < natm){ // distance between atom i and atom j @@ -831,6 +889,35 @@ __host__ int GDFTbecke_partition_weights(double *weights, const double *coords, const double *atm_coords, const double *a_factor, const int *atm_idx, const int ngrids, const int natm, const int scheme_id) { +#ifdef USE_SYCL + sycl::range<2> threads(TILE, TILE); + sycl::range<2> blocks(1, (ngrids+TILE*TILE-1)/(TILE*TILE)); + + const bool if_radii_adjust = a_factor != NULL; + const enum GridPartitionScheme scheme = get_grid_partition_sheme(scheme_id); + + if (scheme == GridPartitionScheme::original_becke) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_kernel< true, GridPartitionScheme::original_becke> (weights, coords, atm_coords, a_factor, atm_idx, ngrids, natm); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_kernel (weights, coords, atm_coords, a_factor, atm_idx, ngrids, natm); + }); + } + } else if (scheme == GridPartitionScheme::stratmann) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_kernel< true, GridPartitionScheme::stratmann> (weights, coords, atm_coords, a_factor, atm_idx, ngrids, natm); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_kernel (weights, coords, atm_coords, a_factor, atm_idx, ngrids, natm); + }); + } + } +#else const dim3 threads(TILE, TILE); const int blocks = (ngrids+TILE*TILE-1)/(TILE*TILE); @@ -849,7 +936,9 @@ int GDFTbecke_partition_weights(double *weights, const double *coords, const dou } else { GDFTgrid_weight_kernel <<>>(weights, coords, atm_coords, a_factor, atm_idx, ngrids, natm); } - } else { + } +#endif + else { cudaMemset(weights, 0xFF, ngrids * sizeof(double)); // Fill with NAN fprintf(stderr, "Incorrect scheme_id = %d in GDFTgrid_weight\n", scheme_id); return 1; @@ -868,6 +957,40 @@ int GDFTbecke_partition_weight_derivative(double *dwdG, const double *grid_coord const double *atm_coords, const double *a_factor, const double *inv_atom_distance, const int *atm_idx, const double *Ar_distance, const double* PB, const double* invsumPB, const int ngrids, const int natm, const int scheme_id) { +#ifdef USE_SYCL + const int n_thread_per_grid = 128; + sycl::range<2> threads(1, n_thread_per_grid); + sycl::range<2> blocks(natm, (ngrids + n_thread_per_grid - 1) / n_thread_per_grid); + + const bool if_radii_adjust = a_factor != NULL; + const enum GridPartitionScheme scheme = get_grid_partition_sheme(scheme_id); + + if (scheme == GridPartitionScheme::original_becke) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_derivative_kernel< true, GridPartitionScheme::original_becke> ( + dwdG, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_derivative_kernel ( + dwdG, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm); + }); + } + } else if (scheme == GridPartitionScheme::stratmann) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_derivative_kernel< true, GridPartitionScheme::stratmann> ( + dwdG, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_derivative_kernel ( + dwdG, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm); + }); + } + } +#else const int n_thread_per_grid = 128; const dim3 threads(n_thread_per_grid, 1); const dim3 blocks((ngrids + n_thread_per_grid - 1) / n_thread_per_grid, natm); @@ -895,7 +1018,9 @@ int GDFTbecke_partition_weight_derivative(double *dwdG, const double *grid_coord dwdG, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm ); } - } else { + } +#endif + else { cudaMemset(dwdG, 0xFF, ngrids * natm * 3 * sizeof(double)); // Fill with NAN fprintf(stderr, "Incorrect scheme_id = %d in GDFTgrid_weight_derivative\n", scheme_id); return 1; @@ -920,6 +1045,41 @@ int GDFTbecke_partition_weight_second_derivative(double *d2w_dG1dG2, const doubl { // Offdiagonal constexpr int n_grid_per_block = 16; constexpr int n_atom_per_block = 4; +#ifdef USE_SYCL + sycl::range<3> threads(n_atom_per_block, n_atom_per_block, n_grid_per_block); + sycl::range<3> blocks((natm + n_atom_per_block - 1) / n_atom_per_block, + (natm + n_atom_per_block - 1) / n_atom_per_block, + (ngrids + n_grid_per_block - 1) / n_grid_per_block); + if (scheme == GridPartitionScheme::original_becke) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<3>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_offdiagonal_kernel< true, GridPartitionScheme::original_becke> ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<3>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_offdiagonal_kernel ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } + } else if (scheme == GridPartitionScheme::stratmann) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<3>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_offdiagonal_kernel< true, GridPartitionScheme::stratmann> ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<3>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_offdiagonal_kernel ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } + } +#else const dim3 threads(n_grid_per_block, n_atom_per_block, n_atom_per_block); const dim3 blocks((ngrids + n_grid_per_block - 1) / n_grid_per_block, (natm + n_atom_per_block - 1) / n_atom_per_block, @@ -944,7 +1104,9 @@ int GDFTbecke_partition_weight_second_derivative(double *d2w_dG1dG2, const doubl d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm ); } - } else { + } +#endif + else { cudaMemset(d2w_dG1dG2, 0xFF, ngrids * natm * natm * 9 * sizeof(double)); // Fill with NAN fprintf(stderr, "Incorrect scheme_id = %d in GDFTgrid_weight_second_derivative\n", scheme_id); return 1; @@ -953,6 +1115,41 @@ int GDFTbecke_partition_weight_second_derivative(double *d2w_dG1dG2, const doubl { // Diagonal constexpr int n_grid_per_block = 64; constexpr int n_atom_per_block = 4; +#ifdef USE_SYCL + const sycl::range<2> threads(n_atom_per_block, n_grid_per_block); + const sycl::range<2> blocks((natm + n_atom_per_block - 1) / n_atom_per_block, + (ngrids + n_grid_per_block - 1) / n_grid_per_block); + + if (scheme == GridPartitionScheme::original_becke) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_diagonal_kernel< true, GridPartitionScheme::original_becke> ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_diagonal_kernel ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } + } else if (scheme == GridPartitionScheme::stratmann) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_diagonal_kernel< true, GridPartitionScheme::stratmann> ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_weight_second_derivative_diagonal_kernel ( + d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm + ); + }); + } + } +#else const dim3 threads(n_grid_per_block, n_atom_per_block); const dim3 blocks((ngrids + n_grid_per_block - 1) / n_grid_per_block, (natm + n_atom_per_block - 1) / n_atom_per_block); @@ -976,7 +1173,9 @@ int GDFTbecke_partition_weight_second_derivative(double *d2w_dG1dG2, const doubl d2w_dG1dG2, grid_coords, grid_quadrature_weights, atm_coords, a_factor, inv_atom_distance, atm_idx, Ar_distance, PB, invsumPB, ngrids, natm ); } - } else { + } +#endif + else { cudaMemset(d2w_dG1dG2, 0xFF, ngrids * natm * natm * 9 * sizeof(double)); // Fill with NAN fprintf(stderr, "Incorrect scheme_id = %d in GDFTgrid_weight_second_derivative\n", scheme_id); return 1; @@ -998,6 +1197,36 @@ int GDFTbecke_eval_PB(double *PB, { constexpr int n_grid_per_block = 64; constexpr int n_atom_per_block = 4; +#ifdef USE_SYCL + sycl::range<2> threads(n_atom_per_block, n_grid_per_block); + sycl::range<2> blocks((natm + n_atom_per_block - 1) / n_atom_per_block, + (ngrids + n_grid_per_block - 1) / n_grid_per_block); + + const bool if_radii_adjust = a_factor != NULL; + const enum GridPartitionScheme scheme = get_grid_partition_sheme(scheme_id); + + if (scheme == GridPartitionScheme::original_becke) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_becke_eval_PB_kernel< true, GridPartitionScheme::original_becke> (PB, a_factor, inv_atom_distance, Ar_distance, ngrids, natm); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_becke_eval_PB_kernel (PB, a_factor, inv_atom_distance, Ar_distance, ngrids, natm); + }); + } + } else if (scheme == GridPartitionScheme::stratmann) { + if (if_radii_adjust) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_becke_eval_PB_kernel< true, GridPartitionScheme::stratmann> (PB, a_factor, inv_atom_distance, Ar_distance, ngrids, natm); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgrid_becke_eval_PB_kernel (PB, a_factor, inv_atom_distance, Ar_distance, ngrids, natm); + }); + } + } +#else const dim3 threads(n_grid_per_block, n_atom_per_block); const dim3 blocks((ngrids + n_grid_per_block - 1) / n_grid_per_block, (natm + n_atom_per_block - 1) / n_atom_per_block); @@ -1017,7 +1246,9 @@ int GDFTbecke_eval_PB(double *PB, } else { GDFTgrid_becke_eval_PB_kernel <<>>(PB, a_factor, inv_atom_distance, Ar_distance, ngrids, natm); } - } else { + } +#endif + else { cudaMemset(PB, 0xFF, ngrids * natm * sizeof(double)); // Fill with NAN fprintf(stderr, "Incorrect scheme_id = %d in GDFTbecke_eval_PB\n", scheme_id); return 1; @@ -1038,6 +1269,13 @@ int GDFTgroup_grids(cudaStream_t stream, int* group_ids, const double* atom_coor fprintf(stderr, "CUDA Error of gen grids: grids alignment must be %d.", NATOM_PER_BLOCK); return 1; } +#ifdef USE_SYCL + sycl::range<1> threads(NATOM_PER_BLOCK); + sycl::range<1> blocks((ngrids+NATOM_PER_BLOCK-1)/NATOM_PER_BLOCK); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GDFTgroup_grids_kernel(group_ids, atom_coords, coords, natm, ngrids); + }); +#else dim3 threads(NATOM_PER_BLOCK); dim3 blocks((ngrids+NATOM_PER_BLOCK-1)/NATOM_PER_BLOCK); GDFTgroup_grids_kernel<<>>(group_ids, atom_coords, coords, natm, ngrids); @@ -1046,6 +1284,7 @@ int GDFTgroup_grids(cudaStream_t stream, int* group_ids, const double* atom_coor fprintf(stderr, "CUDA Error of group grids: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } diff --git a/gpu4pyscf/lib/gdft/nr_eval_gto.cu b/gpu4pyscf/lib/gdft/nr_eval_gto.cu index 0f37c2d24..16d4b31bb 100644 --- a/gpu4pyscf/lib/gdft/nr_eval_gto.cu +++ b/gpu4pyscf/lib/gdft/nr_eval_gto.cu @@ -25,6 +25,32 @@ #include "nr_eval_gto.cuh" #include "contract_rho.cuh" +// Abstracts 2D kernel launch/setup syntax. blocks/threads must be in scope. +// The envs argument is supplied by the macro: SYCL makes an on-host value copy +// of *gto_envs for lambda capture; CUDA passes *gto_envs directly. The SYCL +// kernel name is generated inline per source line (unique in this TU). +#define GDFT_CAT_(a, b) a##b +#define GDFT_CAT(a, b) GDFT_CAT_(a, b) +#ifdef USE_SYCL +#define LAUNCH_KERNEL(KERNEL, ...) { \ + auto dev_gto_envs = *gto_envs; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { KERNEL(__VA_ARGS__, dev_gto_envs); }); } + +#define KERNEL_PROLOGUE_BAS_GRID() \ + auto item = syclex::this_work_item::get_nd_item<2>(); \ + const int grid_id = item.get_global_id(1); \ + const int bas_id = item.get_group(0); +#else +#define LAUNCH_KERNEL(KERNEL, ...) \ + KERNEL<<>>(__VA_ARGS__, *gto_envs); + +#define KERNEL_PROLOGUE_BAS_GRID() \ + const int grid_id = blockIdx.x * blockDim.x + threadIdx.x; \ + int bas_id = blockIdx.y; +#endif + #define NG_PER_BLOCK 256 #define LMAX 8 @@ -51,12 +77,27 @@ static void _screen_index(int8_t *non0shl_mask, double log_cutoff, double *coords, int ngrids, int block_size, int *atm, int natm, int *bas, int nbas, double *env) { - int grid_block_id = blockIdx.x; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int blockIdx_x = item.get_group(1); + const int blockIdx_y = item.get_group(0); + const int blockDim_x = item.get_local_range(1); + const int threadIdx_x = item.get_local_id(1); + double (&gridx_cache)[NG_PER_BLOCK*3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else + const int blockIdx_x = blockIdx.x; + const int blockIdx_y = blockIdx.y; + const int blockDim_x = blockDim.x; + const int threadIdx_x = threadIdx.x; + __shared__ double gridx_cache[NG_PER_BLOCK*3]; +#endif + + int grid_block_id = blockIdx_x; int grid_start = grid_block_id * block_size; int grid_stop = min(grid_start+block_size, ngrids); - int shl_block_id = blockIdx.y; - int thread_id = threadIdx.x; - int ish = shl_block_id * blockDim.x + thread_id; + int shl_block_id = blockIdx_y; + int thread_id = threadIdx_x; + int ish = shl_block_id * blockDim_x + thread_id; if (ish >= nbas) { ish = 0; } @@ -71,7 +112,6 @@ static void _screen_index(int8_t *non0shl_mask, double log_cutoff, double atom_y = ri[1]; double atom_z = ri[2]; - __shared__ double gridx_cache[NG_PER_BLOCK*3]; double *gridy_cache = gridx_cache + NG_PER_BLOCK; double *gridz_cache = gridy_cache + NG_PER_BLOCK; @@ -100,16 +140,29 @@ static void _screen_index(int8_t *non0shl_mask, double log_cutoff, is_large |= log(fabs(gto_sup)) + ang*log(rr)/2 > log_cutoff; } } - if (shl_block_id * blockDim.x + thread_id < nbas) { + if (shl_block_id * blockDim_x + thread_id < nbas) { non0shl_mask[grid_block_id*nbas + ish] = is_large; } } __global__ -static void _screen_index_legacy(int *non0shl_idx, double cutoff, int ang, int nprim, +static void _screen_index_legacy(int *non0shl_idx, double cutoff, int ang, int nprim, double *coords, int ngrids, int bas_offset, GTOValEnvVars gto_envs){ - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int grid_id = item.get_global_id(1); + int ish = item.get_group(0) + bas_offset; + int (&sdata)[NG_PER_BLOCK] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + const int blockDim_x = item.get_local_range(1); + const int threadIdx_x = item.get_local_id(1); +#else + const int grid_id = blockIdx.x * blockDim.x + threadIdx.x; int ish = blockIdx.y + bas_offset; + __shared__ int sdata[NG_PER_BLOCK]; + const int blockDim_x = blockDim.x; + const int threadIdx_x = threadIdx.x; +#endif + const bool active = grid_id < ngrids; int natm = gto_envs.natm; @@ -151,11 +204,10 @@ static void _screen_index_legacy(int *non0shl_idx, double cutoff, int ang, int n int is_large = fabs(gto_sup) > cutoff; // Reduce and write to global memory - unsigned int tx = threadIdx.x; - __shared__ int sdata[NG_PER_BLOCK]; + unsigned int tx = threadIdx_x; sdata[tx] = active ? is_large : 0; __syncthreads(); - for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) { + for (unsigned int s = blockDim_x / 2; s > 0; s >>= 1) { if (tx < s) { sdata[tx] = sdata[tx] || sdata[tx + s]; } @@ -364,12 +416,11 @@ template __global__ static void _cart_kernel_deriv0(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int natm = gto_envs.natm; int local_ish = offsets.bas_off + bas_id; int glob_ish = offsets.bas_indices[local_ish]; @@ -464,12 +515,10 @@ template __global__ static void _cart_kernel_deriv1(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -715,12 +764,10 @@ template __global__ static void _cart_kernel_deriv2(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -793,12 +840,10 @@ template __global__ static void _cart_kernel_deriv3(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -893,12 +938,10 @@ template __global__ static void _cart_kernel_deriv4(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -1023,11 +1066,10 @@ template __global__ static void _sph_kernel_deriv0(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int natm = gto_envs.natm; int local_ish = offsets.bas_off + bas_id; int glob_ish = offsets.bas_indices[local_ish]; @@ -1159,12 +1201,10 @@ template __global__ static void _sph_kernel_deriv1(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -1462,7 +1502,7 @@ static void _sph_kernel_deriv1(BasOffsets offsets, GTOValEnvVars gto_envs) fz0[lx] = fz0[lx-1] * rz; } double fx1[ANG+1], fy1[ANG+1], fz1[ANG+1]; - + _memset_sph(gto+grid_id, 4, ngrids, nao); for (int ip = 0; ip < offsets.nprim; ++ip) { @@ -1481,12 +1521,10 @@ template __global__ static void _sph_kernel_deriv2(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -1554,12 +1592,10 @@ template __global__ static void _sph_kernel_deriv3(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -1649,12 +1685,10 @@ template __global__ static void _sph_kernel_deriv4(BasOffsets offsets, GTOValEnvVars gto_envs) { int ngrids = offsets.ngrids; - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + KERNEL_PROLOGUE_BAS_GRID(); if (grid_id >= ngrids) { return; } - - int bas_id = blockIdx.y; int natm = gto_envs.natm; int nao = offsets.nao; int local_ish = offsets.bas_off + bas_id; @@ -1798,8 +1832,8 @@ int GDFTeval_gto(cudaStream_t stream, double *ao, int deriv, int cart, offsets.bas_indices = bas_indices; offsets.nbas = local_ctr_offsets[nctr]; offsets.nao = nao; - dim3 threads(NG_PER_BLOCK); - dim3 blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); + auto threads = MAKE_RANGE_2D(NG_PER_BLOCK, 1); + auto blocks = MAKE_RANGE_2D((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK, 1); for (int ictr = 0; ictr < nctr; ++ictr) { int local_ish = local_ctr_offsets[ictr]; @@ -1808,150 +1842,153 @@ int GDFTeval_gto(cudaStream_t stream, double *ao, int deriv, int cart, offsets.bas_off = local_ish; offsets.nprim = bas[NPRIM_OF+glob_ish*BAS_SLOTS]; offsets.fac = CINTcommon_fac_sp(l); - blocks.y = local_ctr_offsets[ictr+1] - local_ctr_offsets[ictr]; - if (blocks.y == 0){ + + BLOCKS_SET_Y(local_ctr_offsets[ictr+1] - local_ctr_offsets[ictr]); + if (BLOCKS_GET_Y() == 0){ continue; } + switch (deriv) { case 0: if (cart == 1) { switch (l) { - case 0: _cart_kernel_deriv0<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv0<1> <<>>(offsets, *gto_envs); break; - case 2: _cart_kernel_deriv0<2> <<>>(offsets, *gto_envs); break; - case 3: _cart_kernel_deriv0<3> <<>>(offsets, *gto_envs); break; - case 4: _cart_kernel_deriv0<4> <<>>(offsets, *gto_envs); break; - case 5: _cart_kernel_deriv0<5> <<>>(offsets, *gto_envs); break; - case 6: _cart_kernel_deriv0<6> <<>>(offsets, *gto_envs); break; - case 7: _cart_kernel_deriv0<7> <<>>(offsets, *gto_envs); break; - case 8: _cart_kernel_deriv0<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv0<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv0<1>, offsets) break; + case 2: LAUNCH_KERNEL(_cart_kernel_deriv0<2>, offsets) break; + case 3: LAUNCH_KERNEL(_cart_kernel_deriv0<3>, offsets) break; + case 4: LAUNCH_KERNEL(_cart_kernel_deriv0<4>, offsets) break; + case 5: LAUNCH_KERNEL(_cart_kernel_deriv0<5>, offsets) break; + case 6: LAUNCH_KERNEL(_cart_kernel_deriv0<6>, offsets) break; + case 7: LAUNCH_KERNEL(_cart_kernel_deriv0<7>, offsets) break; + case 8: LAUNCH_KERNEL(_cart_kernel_deriv0<8>, offsets) break; default:fprintf(stderr, "l = %d not supported\n", l); } } else { switch (l) { - case 0: _cart_kernel_deriv0<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv0<1> <<>>(offsets, *gto_envs); break; - case 2: _sph_kernel_deriv0 <2> <<>>(offsets, *gto_envs); break; - case 3: _sph_kernel_deriv0 <3> <<>>(offsets, *gto_envs); break; - case 4: _sph_kernel_deriv0 <4> <<>>(offsets, *gto_envs); break; - case 5: _sph_kernel_deriv0 <5> <<>>(offsets, *gto_envs); break; - case 6: _sph_kernel_deriv0 <6> <<>>(offsets, *gto_envs); break; - case 7: _sph_kernel_deriv0 <7> <<>>(offsets, *gto_envs); break; - case 8: _sph_kernel_deriv0 <8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv0<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv0<1>, offsets) break; + case 2: LAUNCH_KERNEL(_sph_kernel_deriv0 <2>, offsets) break; + case 3: LAUNCH_KERNEL(_sph_kernel_deriv0 <3>, offsets) break; + case 4: LAUNCH_KERNEL(_sph_kernel_deriv0 <4>, offsets) break; + case 5: LAUNCH_KERNEL(_sph_kernel_deriv0 <5>, offsets) break; + case 6: LAUNCH_KERNEL(_sph_kernel_deriv0 <6>, offsets) break; + case 7: LAUNCH_KERNEL(_sph_kernel_deriv0 <7>, offsets) break; + case 8: LAUNCH_KERNEL(_sph_kernel_deriv0 <8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); } } break; case 1: if (cart == 1) { switch (l) { - case 0: _cart_kernel_deriv1<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv1<1> <<>>(offsets, *gto_envs); break; - case 2: _cart_kernel_deriv1<2> <<>>(offsets, *gto_envs); break; - case 3: _cart_kernel_deriv1<3> <<>>(offsets, *gto_envs); break; - case 4: _cart_kernel_deriv1<4> <<>>(offsets, *gto_envs); break; - case 5: _cart_kernel_deriv1<5> <<>>(offsets, *gto_envs); break; - case 6: _cart_kernel_deriv1<6> <<>>(offsets, *gto_envs); break; - case 7: _cart_kernel_deriv1<7> <<>>(offsets, *gto_envs); break; - case 8: _cart_kernel_deriv1<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv1<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv1<1>, offsets) break; + case 2: LAUNCH_KERNEL(_cart_kernel_deriv1<2>, offsets) break; + case 3: LAUNCH_KERNEL(_cart_kernel_deriv1<3>, offsets) break; + case 4: LAUNCH_KERNEL(_cart_kernel_deriv1<4>, offsets) break; + case 5: LAUNCH_KERNEL(_cart_kernel_deriv1<5>, offsets) break; + case 6: LAUNCH_KERNEL(_cart_kernel_deriv1<6>, offsets) break; + case 7: LAUNCH_KERNEL(_cart_kernel_deriv1<7>, offsets) break; + case 8: LAUNCH_KERNEL(_cart_kernel_deriv1<8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); } } else { switch (l) { - case 0: _cart_kernel_deriv1<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv1<1> <<>>(offsets, *gto_envs); break; - case 2: _sph_kernel_deriv1 <2> <<>>(offsets, *gto_envs); break; - case 3: _sph_kernel_deriv1 <3> <<>>(offsets, *gto_envs); break; - case 4: _sph_kernel_deriv1 <4> <<>>(offsets, *gto_envs); break; - case 5: _sph_kernel_deriv1 <5> <<>>(offsets, *gto_envs); break; - case 6: _sph_kernel_deriv1 <6> <<>>(offsets, *gto_envs); break; - case 7: _sph_kernel_deriv1 <7> <<>>(offsets, *gto_envs); break; - case 8: _sph_kernel_deriv1 <8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv1<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv1<1>, offsets) break; + case 2: LAUNCH_KERNEL(_sph_kernel_deriv1 <2>, offsets) break; + case 3: LAUNCH_KERNEL(_sph_kernel_deriv1 <3>, offsets) break; + case 4: LAUNCH_KERNEL(_sph_kernel_deriv1 <4>, offsets) break; + case 5: LAUNCH_KERNEL(_sph_kernel_deriv1 <5>, offsets) break; + case 6: LAUNCH_KERNEL(_sph_kernel_deriv1 <6>, offsets) break; + case 7: LAUNCH_KERNEL(_sph_kernel_deriv1 <7>, offsets) break; + case 8: LAUNCH_KERNEL(_sph_kernel_deriv1 <8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); } } break; case 2: if (cart == 1){ switch (l) { - case 0: _cart_kernel_deriv2<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv2<1> <<>>(offsets, *gto_envs); break; - case 2: _cart_kernel_deriv2<2> <<>>(offsets, *gto_envs); break; - case 3: _cart_kernel_deriv2<3> <<>>(offsets, *gto_envs); break; - case 4: _cart_kernel_deriv2<4> <<>>(offsets, *gto_envs); break; - case 5: _cart_kernel_deriv2<5> <<>>(offsets, *gto_envs); break; - case 6: _cart_kernel_deriv2<6> <<>>(offsets, *gto_envs); break; - case 7: _cart_kernel_deriv2<7> <<>>(offsets, *gto_envs); break; - case 8: _cart_kernel_deriv2<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv2<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv2<1>, offsets) break; + case 2: LAUNCH_KERNEL(_cart_kernel_deriv2<2>, offsets) break; + case 3: LAUNCH_KERNEL(_cart_kernel_deriv2<3>, offsets) break; + case 4: LAUNCH_KERNEL(_cart_kernel_deriv2<4>, offsets) break; + case 5: LAUNCH_KERNEL(_cart_kernel_deriv2<5>, offsets) break; + case 6: LAUNCH_KERNEL(_cart_kernel_deriv2<6>, offsets) break; + case 7: LAUNCH_KERNEL(_cart_kernel_deriv2<7>, offsets) break; + case 8: LAUNCH_KERNEL(_cart_kernel_deriv2<8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); break;} } else { switch(l){ - case 0: _cart_kernel_deriv2<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv2<1> <<>>(offsets, *gto_envs); break; - case 2: _sph_kernel_deriv2<2> <<>>(offsets, *gto_envs); break; - case 3: _sph_kernel_deriv2<3> <<>>(offsets, *gto_envs); break; - case 4: _sph_kernel_deriv2<4> <<>>(offsets, *gto_envs); break; - case 5: _sph_kernel_deriv2<5> <<>>(offsets, *gto_envs); break; - case 6: _sph_kernel_deriv2<6> <<>>(offsets, *gto_envs); break; - case 7: _sph_kernel_deriv2<7> <<>>(offsets, *gto_envs); break; - case 8: _sph_kernel_deriv2<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv2<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv2<1>, offsets) break; + case 2: LAUNCH_KERNEL(_sph_kernel_deriv2<2>, offsets) break; + case 3: LAUNCH_KERNEL(_sph_kernel_deriv2<3>, offsets) break; + case 4: LAUNCH_KERNEL(_sph_kernel_deriv2<4>, offsets) break; + case 5: LAUNCH_KERNEL(_sph_kernel_deriv2<5>, offsets) break; + case 6: LAUNCH_KERNEL(_sph_kernel_deriv2<6>, offsets) break; + case 7: LAUNCH_KERNEL(_sph_kernel_deriv2<7>, offsets) break; + case 8: LAUNCH_KERNEL(_sph_kernel_deriv2<8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); break; } } break; case 3: if (cart == 1){ switch (l) { - case 0: _cart_kernel_deriv3<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv3<1> <<>>(offsets, *gto_envs); break; - case 2: _cart_kernel_deriv3<2> <<>>(offsets, *gto_envs); break; - case 3: _cart_kernel_deriv3<3> <<>>(offsets, *gto_envs); break; - case 4: _cart_kernel_deriv3<4> <<>>(offsets, *gto_envs); break; - case 5: _cart_kernel_deriv3<5> <<>>(offsets, *gto_envs); break; - case 6: _cart_kernel_deriv3<6> <<>>(offsets, *gto_envs); break; - case 7: _cart_kernel_deriv3<7> <<>>(offsets, *gto_envs); break; - case 8: _cart_kernel_deriv3<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv3<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv3<1>, offsets) break; + case 2: LAUNCH_KERNEL(_cart_kernel_deriv3<2>, offsets) break; + case 3: LAUNCH_KERNEL(_cart_kernel_deriv3<3>, offsets) break; + case 4: LAUNCH_KERNEL(_cart_kernel_deriv3<4>, offsets) break; + case 5: LAUNCH_KERNEL(_cart_kernel_deriv3<5>, offsets) break; + case 6: LAUNCH_KERNEL(_cart_kernel_deriv3<6>, offsets) break; + case 7: LAUNCH_KERNEL(_cart_kernel_deriv3<7>, offsets) break; + case 8: LAUNCH_KERNEL(_cart_kernel_deriv3<8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); break; } } else { switch(l){ - case 0: _cart_kernel_deriv3<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv3<1> <<>>(offsets, *gto_envs); break; - case 2: _sph_kernel_deriv3<2> <<>>(offsets, *gto_envs); break; - case 3: _sph_kernel_deriv3<3> <<>>(offsets, *gto_envs); break; - case 4: _sph_kernel_deriv3<4> <<>>(offsets, *gto_envs); break; - case 5: _sph_kernel_deriv3<5> <<>>(offsets, *gto_envs); break; - case 6: _sph_kernel_deriv3<6> <<>>(offsets, *gto_envs); break; - case 7: _sph_kernel_deriv3<7> <<>>(offsets, *gto_envs); break; - case 8: _sph_kernel_deriv3<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv3<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv3<1>, offsets) break; + case 2: LAUNCH_KERNEL(_sph_kernel_deriv3<2>, offsets) break; + case 3: LAUNCH_KERNEL(_sph_kernel_deriv3<3>, offsets) break; + case 4: LAUNCH_KERNEL(_sph_kernel_deriv3<4>, offsets) break; + case 5: LAUNCH_KERNEL(_sph_kernel_deriv3<5>, offsets) break; + case 6: LAUNCH_KERNEL(_sph_kernel_deriv3<6>, offsets) break; + case 7: LAUNCH_KERNEL(_sph_kernel_deriv3<7>, offsets) break; + case 8: LAUNCH_KERNEL(_sph_kernel_deriv3<8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); break; } } break; case 4: if (cart == 1){ switch (l) { - case 0: _cart_kernel_deriv4<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv4<1> <<>>(offsets, *gto_envs); break; - case 2: _cart_kernel_deriv4<2> <<>>(offsets, *gto_envs); break; - case 3: _cart_kernel_deriv4<3> <<>>(offsets, *gto_envs); break; - case 4: _cart_kernel_deriv4<4> <<>>(offsets, *gto_envs); break; - case 5: _cart_kernel_deriv4<5> <<>>(offsets, *gto_envs); break; - case 6: _cart_kernel_deriv4<6> <<>>(offsets, *gto_envs); break; - case 7: _cart_kernel_deriv4<7> <<>>(offsets, *gto_envs); break; - case 8: _cart_kernel_deriv4<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv4<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv4<1>, offsets) break; + case 2: LAUNCH_KERNEL(_cart_kernel_deriv4<2>, offsets) break; + case 3: LAUNCH_KERNEL(_cart_kernel_deriv4<3>, offsets) break; + case 4: LAUNCH_KERNEL(_cart_kernel_deriv4<4>, offsets) break; + case 5: LAUNCH_KERNEL(_cart_kernel_deriv4<5>, offsets) break; + case 6: LAUNCH_KERNEL(_cart_kernel_deriv4<6>, offsets) break; + case 7: LAUNCH_KERNEL(_cart_kernel_deriv4<7>, offsets) break; + case 8: LAUNCH_KERNEL(_cart_kernel_deriv4<8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); break; } } else { switch(l){ - case 0: _cart_kernel_deriv4<0> <<>>(offsets, *gto_envs); break; - case 1: _cart_kernel_deriv4<1> <<>>(offsets, *gto_envs); break; - case 2: _sph_kernel_deriv4<2> <<>>(offsets, *gto_envs); break; - case 3: _sph_kernel_deriv4<3> <<>>(offsets, *gto_envs); break; - case 4: _sph_kernel_deriv4<4> <<>>(offsets, *gto_envs); break; - case 5: _sph_kernel_deriv4<5> <<>>(offsets, *gto_envs); break; - case 6: _sph_kernel_deriv4<6> <<>>(offsets, *gto_envs); break; - case 7: _sph_kernel_deriv4<7> <<>>(offsets, *gto_envs); break; - case 8: _sph_kernel_deriv4<8> <<>>(offsets, *gto_envs); break; + case 0: LAUNCH_KERNEL(_cart_kernel_deriv4<0>, offsets) break; + case 1: LAUNCH_KERNEL(_cart_kernel_deriv4<1>, offsets) break; + case 2: LAUNCH_KERNEL(_sph_kernel_deriv4<2>, offsets) break; + case 3: LAUNCH_KERNEL(_sph_kernel_deriv4<3>, offsets) break; + case 4: LAUNCH_KERNEL(_sph_kernel_deriv4<4>, offsets) break; + case 5: LAUNCH_KERNEL(_sph_kernel_deriv4<5>, offsets) break; + case 6: LAUNCH_KERNEL(_sph_kernel_deriv4<6>, offsets) break; + case 7: LAUNCH_KERNEL(_sph_kernel_deriv4<7>, offsets) break; + case 8: LAUNCH_KERNEL(_sph_kernel_deriv4<8>, offsets) break; default: fprintf(stderr, "l = %d not supported\n", l); break; } } break; default: fprintf(stderr, "deriv %d not supported\n", deriv); return 1; - } + } // switch + cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GDFTeval_gto_kernel: %s\n", cudaGetErrorString(err)); @@ -1966,13 +2003,20 @@ int GDFTscreen_index(cudaStream_t stream, int8_t *non0shl_mask, double log_cutof double *grids, int ngrids, int block_size, int *atm, int natm, int *bas, int nbas, double *env) { - dim3 threads(NG_PER_BLOCK); - dim3 blocks((ngrids+block_size-1)/block_size, - (nbas+NG_PER_BLOCK-1)/NG_PER_BLOCK); - _screen_index<<>> ( - non0shl_mask, log_cutoff, grids, ngrids, block_size, + auto threads = MAKE_RANGE_2D(NG_PER_BLOCK, 1); + auto blocks = MAKE_RANGE_2D((ngrids+block_size-1)/block_size, + (nbas+NG_PER_BLOCK-1)/NG_PER_BLOCK); + // _screen_index takes no gto_envs; launch directly (LAUNCH_KERNEL appends envs). +#ifdef USE_SYCL + stream.parallel_for( + sycl::nd_range<2>(blocks * threads, threads), + [=](auto item) [[intel::kernel_args_restrict]] { + _screen_index(non0shl_mask, log_cutoff, grids, ngrids, block_size, + atm, natm, bas, nbas, env); }); +#else + _screen_index<<>>(non0shl_mask, log_cutoff, grids, ngrids, block_size, atm, natm, bas, nbas, env); - +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GDFTscreen_index: %s\n", cudaGetErrorString(err)); @@ -1985,26 +2029,26 @@ int GDFTscreen_index_legacy(cudaStream_t stream, int *non0shl_idx, double cutoff double *grids, int ngrids, int *ctr_offsets, int nctr, int *bas, GTOValEnvVars *gto_envs) { - dim3 threads(NG_PER_BLOCK); - dim3 blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); + auto threads = MAKE_RANGE_2D(NG_PER_BLOCK, 1); + auto blocks = MAKE_RANGE_2D((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK, 1); for (int ictr = 0; ictr < nctr; ictr++){ int ish = ctr_offsets[ictr]; const int l = bas[ANG_OF+ish*BAS_SLOTS]; int nprim = bas[NPRIM_OF+ish*BAS_SLOTS]; int bas_offset = ctr_offsets[ictr]; - blocks.y = ctr_offsets[ictr+1] - bas_offset; - if (blocks.y == 0){ + BLOCKS_SET_Y(ctr_offsets[ictr+1] - bas_offset); + if (BLOCKS_GET_Y() == 0){ continue; } if (l > 8){ fprintf(stderr, "l = %d not supported\n", l); return 1; } - _screen_index_legacy<<>> (non0shl_idx, cutoff, l, nprim, - grids, ngrids, bas_offset, *gto_envs); + LAUNCH_KERNEL(_screen_index_legacy, + non0shl_idx, cutoff, l, nprim, + grids, ngrids, bas_offset); } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GDFTscreen_index: %s\n", cudaGetErrorString(err)); @@ -2013,4 +2057,7 @@ int GDFTscreen_index_legacy(cudaStream_t stream, int *non0shl_idx, double cutoff return 0; } -} +} // end extern "C" + +#undef LAUNCH_KERNEL +#undef KERNEL_PROLOGUE_BAS_GRID diff --git a/gpu4pyscf/lib/gdft/nr_eval_gto.cuh b/gpu4pyscf/lib/gdft/nr_eval_gto.cuh index 357068032..f48977d59 100644 --- a/gpu4pyscf/lib/gdft/nr_eval_gto.cuh +++ b/gpu4pyscf/lib/gdft/nr_eval_gto.cuh @@ -16,6 +16,33 @@ #pragma once +// Kernel-side & launch-config macros to unify CUDA and SYCL in gdft. +// All variants in ONE ifdef block - single pair of #ifdef / #else. +#ifdef USE_SYCL + +#define SHARED_ARRAY(T, name, SIZE) \ + using name##_tile_t = T[SIZE]; \ + name##_tile_t& name = *sycl::ext::oneapi:: \ + group_local_memory_for_overwrite(item.get_group()); + +#define MAKE_RANGE_2D(X, Y) sycl::range<2>((Y), (X)) +#define MAKE_RANGE_3D(X, Y, Z) sycl::range<3>((Z), (Y), (X)) + +#define BLOCKS_SET_Y(val) (blocks[0] = (val)) +#define BLOCKS_GET_Y() (blocks[0]) + +#else + +#define SHARED_ARRAY(T, name, SIZE) __shared__ T name[SIZE]; + +#define MAKE_RANGE_2D(X, Y) dim3((X), (Y)) +#define MAKE_RANGE_3D(X, Y, Z) dim3((X), (Y), (Z)) + +#define BLOCKS_SET_Y(val) (blocks.y = (val)) +#define BLOCKS_GET_Y() (blocks.y) + +#endif // USE_SYCL + typedef struct { int natm; int nbas; diff --git a/gpu4pyscf/lib/gdft/nr_numint_sparse.cu b/gpu4pyscf/lib/gdft/nr_numint_sparse.cu index d9a50762e..68be5de3c 100644 --- a/gpu4pyscf/lib/gdft/nr_numint_sparse.cu +++ b/gpu4pyscf/lib/gdft/nr_numint_sparse.cu @@ -21,6 +21,7 @@ #include #include #include "gint/cuda_alloc.cuh" +#include "nr_eval_gto.cuh" #define THREADSX 32 #define THREADSY 4 @@ -28,15 +29,63 @@ #define THREADSYY (THREADSY * THREADSY) #define DIVXY (THREADSX / THREADSY) +#ifdef USE_SYCL +#define LAUNCH_KERNEL_3D(TAG, KERNEL, ...) \ + sycl_get_queue()->parallel_for( \ + sycl::nd_range<3>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { KERNEL(__VA_ARGS__); }); + +#define KERNEL_PROLOGUE_3D_DM() \ + auto item = syclex::this_work_item::get_nd_item<3>(); \ + sycl::group thread_block = item.get_group(); \ + int tx = item.get_local_id(2); \ + int ty = item.get_local_id(1); \ + int grid_blk = thread_block.get_group_id(2); \ + int shell_blk = thread_block.get_group_id(1); \ + int blockIdx_z = thread_block.get_group_id(0); + +#define KERNEL_PROLOGUE_3D_AOW() \ + auto item = syclex::this_work_item::get_nd_item<3>(); \ + sycl::group thread_block = item.get_group(); \ + const int tx = item.get_local_id(2); \ + const int ty = item.get_local_id(1); \ + const int tz = item.get_local_id(0); \ + const int task_ij = thread_block.get_group_id(2); \ + const int blockIdx_y = thread_block.get_group_id(1); \ + const int blockIdx_z = thread_block.get_group_id(0); \ + const int gridDim_y = item.get_group_range(1); \ + const int gridDim_z = item.get_group_range(0); +#else +#define LAUNCH_KERNEL_3D(TAG, KERNEL, ...) \ + KERNEL<<>>(__VA_ARGS__); + +#define KERNEL_PROLOGUE_3D_DM() \ + int tx = threadIdx.x; \ + int ty = threadIdx.y; \ + int grid_blk = blockIdx.x; \ + int shell_blk = blockIdx.y; \ + int blockIdx_z = blockIdx.z; + +#define KERNEL_PROLOGUE_3D_AOW() \ + int task_ij = blockIdx.x; \ + int tx = threadIdx.x; \ + int ty = threadIdx.y; \ + int tz = threadIdx.z; \ + int blockIdx_y = blockIdx.y; \ + int blockIdx_z = blockIdx.z; \ + int gridDim_y = gridDim.y; \ + int gridDim_z = gridDim.z; +#endif + __global__ static void _dot_ao_dm(double *out, double *ao, double *dm, int jsh0, int jsh1, int ngrids, int nbas, int nbins, int nsegs, int *bas_segs, uint8_t *screen_index, uint8_t *pair_mask, int *ao_loc) { - int tx = threadIdx.x; - int ty = threadIdx.y; - int grid_blk = blockIdx.x; - int shell_blk = blockIdx.y; + KERNEL_PROLOGUE_3D_DM(); + SHARED_ARRAY(double, s_ao, THREADSX*THREADSY); + SHARED_ARRAY(double, s_dm, THREADSX*THREADSY); + int jsh = jsh0 + shell_blk * THREADSY + ty; if (jsh >= jsh1) { return; @@ -57,16 +106,13 @@ static void _dot_ao_dm(double *out, double *ao, double *dm, int jsh0, int jsh1, } int grid_id = grid_blk * THREADSX + tx; - int jp = blockIdx.z; + int jp = blockIdx_z; int j = ao_loc[jsh] + jp; int ishp, ip, k, i, seg; size_t Nao = ao_loc[nbas]; size_t Ngrids = ngrids; double val = 0; - __shared__ double s_ao[THREADSX*THREADSY]; - __shared__ double s_dm[THREADSX*THREADSY]; - for (seg = 0; seg < nsegs; seg++) { int ish0 = bas_segs[seg]; int ish1 = bas_segs[seg+1]; @@ -107,10 +153,10 @@ static void _dot_ao_dmT(double *out, double *ao, double *dm, int jsh0, int jsh1, int ngrids, int nbas, int nbins, int nsegs, int *bas_segs, uint8_t *screen_index, uint8_t *pair_mask, int *ao_loc) { - int tx = threadIdx.x; - int ty = threadIdx.y; - int grid_blk = blockIdx.x; - int shell_blk = blockIdx.y; + KERNEL_PROLOGUE_3D_DM(); + SHARED_ARRAY(double, s_ao, THREADSX*THREADSY); + SHARED_ARRAY(double, s_dm, THREADSX*THREADSY); + int jsh = jsh0 + shell_blk * THREADSY + ty; if (jsh >= jsh1) { return; @@ -131,16 +177,13 @@ static void _dot_ao_dmT(double *out, double *ao, double *dm, int jsh0, int jsh1, } int grid_id = grid_blk * THREADSX + tx; - int jp = blockIdx.z; + int jp = blockIdx_z; int j = ao_loc[jsh] + jp; int ishp, ip, k, i, seg; size_t Nao = ao_loc[nbas]; size_t Ngrids = ngrids; double val = 0; - __shared__ double s_ao[THREADSX*THREADSY]; - __shared__ double s_dm[THREADSX*THREADSY]; - for (seg = 0; seg < nsegs; seg++) { int ish0 = bas_segs[seg]; int ish1 = bas_segs[seg+1]; @@ -181,10 +224,10 @@ static void _dot_aow_ao(double *out, double *bra, double *ket, double *wv, int ngrids, int nbas, int nbins, uint8_t *screen_index, int *bas_pair2bra, int *bas_pair2ket, int *ao_loc) { - int task_ij = blockIdx.x; - int tx = threadIdx.x; - int ty = threadIdx.y; - int tz = threadIdx.z; + KERNEL_PROLOGUE_3D_AOW(); + SHARED_ARRAY(double, s_bra, THREADSXY); + SHARED_ARRAY(double, s_ket, THREADSXY); + int txy = ty * DIVXY + tx; int tyz = tz * THREADSY + ty; int ish0 = bas_pair2bra[task_ij]; @@ -193,19 +236,16 @@ static void _dot_aow_ao(double *out, double *bra, double *ket, double *wv, int j0 = ao_loc[jsh0]; int ish4 = ish0 / THREADSY; int jsh4 = jsh0 / THREADSY; - int degen_i = gridDim.y; - int degen_j = gridDim.z; - int ip = blockIdx.y; - int jp = blockIdx.z; + int degen_i = gridDim_y; + int degen_j = gridDim_z; + int ip = blockIdx_y; + int jp = blockIdx_z; int bas_blocks = (nbas + THREADSY - 1) / THREADSY; size_t Nao = ao_loc[nbas]; size_t Ngrids = ngrids; double val = 0; - __shared__ double s_bra[THREADSXY]; - __shared__ double s_ket[THREADSXY]; - int grid_blk; for (grid_blk = 0; grid_blk < ngrids/THREADSX; grid_blk++) { int grid0 = grid_blk * THREADSX; @@ -223,7 +263,11 @@ static void _dot_aow_ao(double *out, double *bra, double *ket, double *wv, double s2 = ket[j*Ngrids+grid_id]; double s = abs(s1 * s2); if (s > 1e-3 && si+sj < nbins){ + // #ifdef USE_SYCL + // sycl::ext::oneapi::experimental::printf("%f %f %f %d %d %d %d %d %d %d %d\n", s, s1, s2, si, sj, si+sj, grid_id, ish0, jsh0, i, j); + // #else printf("%f %f %f %d %d %d %d %d %d %d %d\n", s, s1, s2, si, sj, si+sj, grid_id, ish0, jsh0, i, j); + // #endif } } __syncthreads(); @@ -274,10 +318,10 @@ static void _dot_ao_ao(double *out, double *bra, double *ket, int ngrids, int nbas, int nbins, uint8_t *screen_index, int *bas_pair2bra, int *bas_pair2ket, int *ao_loc) { - int task_ij = blockIdx.x; - int tx = threadIdx.x; - int ty = threadIdx.y; - int tz = threadIdx.z; + KERNEL_PROLOGUE_3D_AOW(); + SHARED_ARRAY(double, s_bra, THREADSXY); + SHARED_ARRAY(double, s_ket, THREADSXY); + int txy = ty * DIVXY + tx; int tyz = tz * THREADSY + ty; int ish0 = bas_pair2bra[task_ij]; @@ -286,19 +330,16 @@ static void _dot_ao_ao(double *out, double *bra, double *ket, int j0 = ao_loc[jsh0]; int ish4 = ish0 / THREADSY; int jsh4 = jsh0 / THREADSY; - int degen_i = gridDim.y; - int degen_j = gridDim.z; - int ip = blockIdx.y; - int jp = blockIdx.z; + int degen_i = gridDim_y; + int degen_j = gridDim_z; + int ip = blockIdx_y; + int jp = blockIdx_z; int bas_blocks = (nbas + THREADSY - 1) / THREADSY; size_t Nao = ao_loc[nbas]; size_t Ngrids = ngrids; double val = 0; - __shared__ double s_bra[THREADSXY]; - __shared__ double s_ket[THREADSXY]; - int grid_blk; for (grid_blk = 0; grid_blk < ngrids/THREADSX; grid_blk++) { int grid0 = grid_blk * THREADSX; @@ -433,16 +474,18 @@ int GDFTdot_ao_dm_sparse(double *out, double *ao, double *dm, int trans_dm, assert(ish1 % THREADSY == 0); int degen = ao_loc[ish0+1] - ao_loc[ish0]; int nsh = ish1 - ish0; - dim3 threads(THREADSX, THREADSY); - dim3 blocks((ngrids+THREADSX-1)/THREADSX, (nsh+THREADSY-1)/THREADSY, degen); + auto threads = MAKE_RANGE_3D(THREADSX, THREADSY, 1); + auto blocks = MAKE_RANGE_3D(grid_blocks, (nsh+THREADSY-1)/THREADSY, degen); if (trans_dm) { - _dot_ao_dmT<<>>(out, ao, dm, ish0, ish1, ngrids, nbas, - nbins, nsegs, d_seg_loc, d_sindex, - d_pair_mask, d_ao_loc); + LAUNCH_KERNEL_3D(_dot_ao_dmT_sycl, _dot_ao_dmT, + out, ao, dm, ish0, ish1, ngrids, nbas, + nbins, nsegs, d_seg_loc, d_sindex, + d_pair_mask, d_ao_loc); } else { - _dot_ao_dm<<>>(out, ao, dm, ish0, ish1, ngrids, nbas, - nbins, nsegs, d_seg_loc, d_sindex, - d_pair_mask, d_ao_loc); + LAUNCH_KERNEL_3D(_dot_ao_dm_sycl, _dot_ao_dm, + out, ao, dm, ish0, ish1, ngrids, nbas, + nbins, nsegs, d_seg_loc, d_sindex, + d_pair_mask, d_ao_loc); } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -488,10 +531,11 @@ int GDFTdot_aow_ao_sparse(double *out, double *bra, double *ket, double *wv, assert(jsh0 % THREADSY == 0); int degen_i = ao_loc[ish0+1] - ao_loc[ish0]; int degen_j = ao_loc[jsh0+1] - ao_loc[jsh0]; - dim3 threads(DIVXY, THREADSY, THREADSY); - dim3 blocks(ntasks, degen_i, degen_j); - _dot_aow_ao<<>>(out, bra, ket, wv, ngrids, nbas, nbins, d_sindex, - d_pair2bra+task0, d_pair2ket+task0, d_ao_loc); + auto threads = MAKE_RANGE_3D(DIVXY, THREADSY, THREADSY); + auto blocks = MAKE_RANGE_3D(ntasks, degen_i, degen_j); + LAUNCH_KERNEL_3D(_dot_aow_ao_sycl, _dot_aow_ao, + out, bra, ket, wv, ngrids, nbas, nbins, d_sindex, + d_pair2bra+task0, d_pair2ket+task0, d_ao_loc); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GDFTdot_aow_ao_sparse: %s\n", @@ -535,10 +579,11 @@ int GDFTdot_ao_ao_sparse(double *out, double *bra, double *ket, assert(jsh0 % THREADSY == 0); int degen_i = ao_loc[ish0+1] - ao_loc[ish0]; int degen_j = ao_loc[jsh0+1] - ao_loc[jsh0]; - dim3 threads(DIVXY, THREADSY, THREADSY); - dim3 blocks(ntasks, degen_i, degen_j); - _dot_ao_ao<<>>(out, bra, ket, ngrids, nbas, nbins, d_sindex, - d_pair2bra+task0, d_pair2ket+task0, d_ao_loc); + auto threads = MAKE_RANGE_3D(DIVXY, THREADSY, THREADSY); + auto blocks = MAKE_RANGE_3D(ntasks, degen_i, degen_j); + LAUNCH_KERNEL_3D(_dot_ao_sycl, _dot_ao_ao, + out, bra, ket, ngrids, nbas, nbins, d_sindex, + d_pair2bra+task0, d_pair2ket+task0, d_ao_loc); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GDFTdot_ao_ao_sparse: %s\n", @@ -554,3 +599,7 @@ cleanup: return err_code; } } + +#undef KERNEL_PROLOGUE_3D_DM +#undef KERNEL_PROLOGUE_3D_AOW +#undef LAUNCH_KERNEL_3D diff --git a/gpu4pyscf/lib/gdft/vv10.cu b/gpu4pyscf/lib/gdft/vv10.cu index 81fc046d7..aa516c469 100644 --- a/gpu4pyscf/lib/gdft/vv10.cu +++ b/gpu4pyscf/lib/gdft/vv10.cu @@ -33,7 +33,20 @@ static void vv10_fock_eval_UWE_kernel(double* __restrict__ U, double* __restrict const double* __restrict__ omega, const double* __restrict__ kappa, const int ngrids) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int i = item.get_global_id(0); + const int threadIdx_x = item.get_local_id(0); + sycl::group thread_block = item.get_group(); + double3 (&shared_omega_kappa_rhow_j)[NG_PER_BLOCK] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double3 (&shared_r_j)[NG_PER_BLOCK] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int threadIdx_x = threadIdx.x; + __shared__ double3 shared_omega_kappa_rhow_j[NG_PER_BLOCK]; + __shared__ double3 shared_r_j[NG_PER_BLOCK]; +#endif + const bool active = i < ngrids; double omega_i = NAN; @@ -51,18 +64,15 @@ static void vv10_fock_eval_UWE_kernel(double* __restrict__ U, double* __restrict double W_i = 0; double E_i = 0; - __shared__ double3 shared_omega_kappa_rhow_j[NG_PER_BLOCK]; - __shared__ double3 shared_r_j[NG_PER_BLOCK]; - for (int j_block_offset = 0; j_block_offset < ngrids; j_block_offset += NG_PER_BLOCK) { - const int j = j_block_offset + threadIdx.x; + const int j = j_block_offset + threadIdx_x; if (j < ngrids) { - shared_omega_kappa_rhow_j[threadIdx.x].x = omega[j]; - shared_omega_kappa_rhow_j[threadIdx.x].y = kappa[j]; - shared_omega_kappa_rhow_j[threadIdx.x].z = rho_weight[j]; - shared_r_j[threadIdx.x].x = grid_coord[j * 3 + 0]; - shared_r_j[threadIdx.x].y = grid_coord[j * 3 + 1]; - shared_r_j[threadIdx.x].z = grid_coord[j * 3 + 2]; + shared_omega_kappa_rhow_j[threadIdx_x].x = omega[j]; + shared_omega_kappa_rhow_j[threadIdx_x].y = kappa[j]; + shared_omega_kappa_rhow_j[threadIdx_x].z = rho_weight[j]; + shared_r_j[threadIdx_x].x = grid_coord[j * 3 + 0]; + shared_r_j[threadIdx_x].y = grid_coord[j * 3 + 1]; + shared_r_j[threadIdx_x].z = grid_coord[j * 3 + 2]; } __syncthreads(); @@ -103,7 +113,12 @@ static void vv10_fock_eval_omega_derivative_kernel(double* __restrict__ omega, d const double* __restrict__ rho, const double* __restrict__ gamma, const double C_factor, const int ngrids) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int i = item.get_global_id(0); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (i >= ngrids) return; @@ -132,8 +147,14 @@ static void vv10_grad_eval_E_grid_response_offdiagonal_kernel(double* __restrict const int* __restrict__ grid_associated_atom, const int* __restrict__ grid_offsets_of_atom, const int natoms, const int i_grid_begin, const int ngrids) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i_unoffset = item.get_global_id(1); + const int B_atom = item.get_group(0); +#else const int i_unoffset = blockIdx.x * blockDim.x + threadIdx.x; const int B_atom = blockIdx.y; +#endif if (i_unoffset >= ngrids || B_atom >= natoms) return; const int i = i_unoffset + i_grid_begin; @@ -194,7 +215,12 @@ static void vv10_hess_eval_UWABCE_kernel(double* __restrict__ U, double* __restr const double* __restrict__ omega, const double* __restrict__ kappa, const int ngrids) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int i = item.get_global_id(0); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (i >= ngrids) return; @@ -252,7 +278,12 @@ static void vv10_hess_eval_omega_derivative_kernel(double* __restrict__ omega, d const double* __restrict__ rho, const double* __restrict__ gamma, const double C_factor, const int ngrids) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int i = item.get_global_id(0); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (i >= ngrids) return; @@ -292,8 +323,14 @@ static void vv10_hess_eval_f_t_offdiagonal_kernel(double* __restrict__ f_rho_t, const double* __restrict__ rho_t, const double* __restrict__ gamma_t, const int ngrids, const int ntrial) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = item.get_global_id(1); + const int i_trial_start = (item.get_global_id(0)) * n_trial_per_thread; +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; const int i_trial_start = (blockIdx.y * blockDim.y + threadIdx.y) * n_trial_per_thread; +#endif if (i >= ngrids || i_trial_start >= ntrial) return; @@ -377,7 +414,12 @@ static void vv10_hess_eval_f_t_diagonal_kernel(double* __restrict__ f_rho_t, dou const double* __restrict__ rho_t, const double* __restrict__ gamma_t, const int ngrids, const int ntrial) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int i = item.get_global_id(0); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (i >= ngrids) return; @@ -420,8 +462,14 @@ static void vv10_hess_eval_EUW_grid_response_offdiagonal_kernel(double* __restri const int* __restrict__ grid_associated_atom, const int* __restrict__ grid_offsets_of_atom, const int ngrids, const int natoms) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = item.get_global_id(1); + const int B_atom = item.get_group(0); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; const int B_atom = blockIdx.y; +#endif if (i >= ngrids || B_atom >= natoms) return; const int i_associated_atom = grid_associated_atom[i]; @@ -506,8 +554,14 @@ static void vv10_hess_eval_EUW_with_weight1_kernel(double* __restrict__ Ew, doub const double* __restrict__ rho, const double* __restrict__ omega, const double* __restrict__ kappa, const int ngrids, const int nderivative) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = item.get_global_id(1); + const int i_derivative_start = (item.get_global_id(0)) * n_derivative_per_thread; +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; const int i_derivative_start = (blockIdx.y * blockDim.y + threadIdx.y) * n_derivative_per_thread; +#endif if (i >= ngrids || i_derivative_start >= nderivative) return; @@ -562,8 +616,14 @@ static void vv10_hess_eval_D_B_in_double_grid_response_offdiagonal_kernel(double const int* __restrict__ grid_associated_atom, const int* __restrict__ grid_offsets_of_atom, const int ngrids, const int natoms) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = item.get_global_id(1); + const int B_atom = item.get_group(0); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; const int B_atom = blockIdx.y; +#endif if (i >= ngrids || B_atom >= natoms) return; const int i_associated_atom = grid_associated_atom[i]; @@ -654,6 +714,13 @@ int VXC_vv10nlc_fock_eval_UWE(const cudaStream_t stream, const double* omega, const double* kappa, const int ngrids) { +#ifdef USE_SYCL + const sycl::range<1> threads(NG_PER_BLOCK); + const sycl::range<1> blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_fock_eval_UWE_kernel(U, W, E, + grid_coord, rho_weight, omega, kappa, ngrids); }); +#else const dim3 threads(NG_PER_BLOCK); const dim3 blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); vv10_fock_eval_UWE_kernel<<>>(U, W, E, @@ -663,6 +730,7 @@ int VXC_vv10nlc_fock_eval_UWE(const cudaStream_t stream, fprintf(stderr, "CUDA Error of vv10 fock eval_UWE: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -672,6 +740,13 @@ int VXC_vv10nlc_fock_eval_omega_derivative(const cudaStream_t stream, const double* rho, const double* gamma, const double C_factor, const int ngrids) { +#ifdef USE_SYCL + const sycl::range<1> threads(NG_PER_BLOCK); + const sycl::range<1> blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_fock_eval_omega_derivative_kernel(omega, domega_drho, domega_dgamma, + rho, gamma, C_factor, ngrids); }); +#else const dim3 threads(NG_PER_BLOCK); const dim3 blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); vv10_fock_eval_omega_derivative_kernel<<>>(omega, domega_drho, domega_dgamma, @@ -681,6 +756,7 @@ int VXC_vv10nlc_fock_eval_omega_derivative(const cudaStream_t stream, fprintf(stderr, "CUDA Error of vv10 fock eval_omega_derivative: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -693,6 +769,13 @@ int VXC_vv10nlc_grad_eval_E_grid_response_offdiagonal(const cudaStream_t stream, const int natm, const int i_grid_begin, const int ngrids) { constexpr int n_grids_per_block = 128; +#ifdef USE_SYCL + const sycl::range<2> threads(1, n_grids_per_block); + const sycl::range<2> blocks(natm, (ngrids + n_grids_per_block - 1) / n_grids_per_block); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_grad_eval_E_grid_response_offdiagonal_kernel(Egr, grid_coord, rho_weight, omega, kappa, grid_associated_atom, grid_offsets_of_atom, natm, i_grid_begin, ngrids); + }); +#else const dim3 threads(n_grids_per_block, 1); const dim3 blocks((ngrids + n_grids_per_block - 1) / n_grids_per_block, natm); vv10_grad_eval_E_grid_response_offdiagonal_kernel<<>>( @@ -703,6 +786,7 @@ int VXC_vv10nlc_grad_eval_E_grid_response_offdiagonal(const cudaStream_t stream, fprintf(stderr, "CUDA Error of vv10 grad eval_E_grid_response: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -713,6 +797,14 @@ int VXC_vv10nlc_hess_eval_UWABCE(const cudaStream_t stream, const double* omega, const double* kappa, const int ngrids) { +#ifdef USE_SYCL + const sycl::range<1> threads(NG_PER_BLOCK); + const sycl::range<1> blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_hess_eval_UWABCE_kernel(U, W, A, B, C, E, + grid_coord, rho_weight, omega, kappa, ngrids); + }); +#else const dim3 threads(NG_PER_BLOCK); const dim3 blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); vv10_hess_eval_UWABCE_kernel<<>>(U, W, A, B, C, E, @@ -722,6 +814,7 @@ int VXC_vv10nlc_hess_eval_UWABCE(const cudaStream_t stream, fprintf(stderr, "CUDA Error of vv10 hess eval_UWABCE: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -732,6 +825,15 @@ int VXC_vv10nlc_hess_eval_omega_derivative(const cudaStream_t stream, const double* rho, const double* gamma, const double C_factor, const int ngrids) { +#ifdef USE_SYCL + const sycl::range<1> threads(NG_PER_BLOCK); + const sycl::range<1> blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_hess_eval_omega_derivative_kernel(omega, domega_drho, domega_dgamma, + d2omega_drho2, d2omega_dgamma2, d2omega_drho_dgamma, + rho, gamma, C_factor, ngrids); + }); +#else const dim3 threads(NG_PER_BLOCK); const dim3 blocks((ngrids+NG_PER_BLOCK-1)/NG_PER_BLOCK); vv10_hess_eval_omega_derivative_kernel<<>>(omega, domega_drho, domega_dgamma, @@ -742,6 +844,7 @@ int VXC_vv10nlc_hess_eval_omega_derivative(const cudaStream_t stream, fprintf(stderr, "CUDA Error of vv10 hess eval_omega_derivative: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -758,6 +861,18 @@ int VXC_vv10nlc_hess_eval_f_t(const cudaStream_t stream, { { constexpr int n_trial_per_thread = 6; // Notice: ntrial is likely a multiple of 3 + #ifdef USE_SYCL + const sycl::range<2> threads(1, NG_PER_BLOCK); + const sycl::range<2> blocks((ntrial + n_trial_per_thread - 1) / n_trial_per_thread, + (ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_hess_eval_f_t_offdiagonal_kernel (f_rho_t, f_gamma_t, + grid_coord, grid_weight, rho, omega, kappa, + domega_drho, domega_dgamma, dkappa_drho, + rho_t, gamma_t, ngrids, ntrial + ); + }); + #else const dim3 threads(NG_PER_BLOCK, 1); const dim3 blocks((ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK, (ntrial + n_trial_per_thread - 1) / n_trial_per_thread); @@ -767,8 +882,22 @@ int VXC_vv10nlc_hess_eval_f_t(const cudaStream_t stream, domega_drho, domega_dgamma, dkappa_drho, rho_t, gamma_t, ngrids, ntrial ); + #endif } { + #ifdef USE_SYCL + const sycl::range<1> threads(NG_PER_BLOCK); + const sycl::range<1> blocks((ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_hess_eval_f_t_diagonal_kernel(f_rho_t, f_gamma_t, + rho, + U, W, A, B, C, + domega_drho, domega_dgamma, dkappa_drho, + d2omega_drho2, d2omega_dgamma2, d2omega_drho_dgamma, d2kappa_drho2, + rho_t, gamma_t, ngrids, ntrial + ); + }); + #else const dim3 threads(NG_PER_BLOCK); const dim3 blocks((ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK); vv10_hess_eval_f_t_diagonal_kernel<<>> ( @@ -779,6 +908,7 @@ int VXC_vv10nlc_hess_eval_f_t(const cudaStream_t stream, d2omega_drho2, d2omega_dgamma2, d2omega_drho_dgamma, d2kappa_drho2, rho_t, gamma_t, ngrids, ntrial ); + #endif } const cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -797,6 +927,13 @@ int VXC_vv10nlc_hess_eval_EUW_grid_response_offdiagonal(const cudaStream_t strea const int ngrids, const int natm) { constexpr int n_grids_per_block = 128; + #ifdef USE_SYCL + const sycl::range<2> threads(1, n_grids_per_block); + const sycl::range<2> blocks(natm, (ngrids + n_grids_per_block - 1) / n_grids_per_block); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_hess_eval_EUW_grid_response_offdiagonal_kernel(Egr, Ugr, Wgr, grid_coord, rho_weight, omega, kappa, grid_associated_atom, grid_offsets_of_atom, ngrids, natm); + }); + #else const dim3 threads(n_grids_per_block, 1); const dim3 blocks((ngrids + n_grids_per_block - 1) / n_grids_per_block, natm); vv10_hess_eval_EUW_grid_response_offdiagonal_kernel<<>>( @@ -806,6 +943,7 @@ int VXC_vv10nlc_hess_eval_EUW_grid_response_offdiagonal(const cudaStream_t strea fprintf(stderr, "CUDA Error of vv10 hess eval_EUW_grid_response: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -817,6 +955,16 @@ int VXC_vv10nlc_hess_eval_EUW_with_weight1(const cudaStream_t stream, const int ngrids, const int nderivative) { constexpr int n_derivative_per_thread = 6; // Notice: ntrial is always a multiple of 3 +#ifdef USE_SYCL + const sycl::range<2> threads(1, NG_PER_BLOCK); + const sycl::range<2> blocks((nderivative + n_derivative_per_thread - 1) / n_derivative_per_thread, + (ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_hess_eval_EUW_with_weight1_kernel (Ew, Uw, Ww, + grid_coord, grid_weight1, rho, omega, kappa, + ngrids, nderivative); + }); +#else const dim3 threads(NG_PER_BLOCK, 1); const dim3 blocks((ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK, (nderivative + n_derivative_per_thread - 1) / n_derivative_per_thread); @@ -830,6 +978,7 @@ int VXC_vv10nlc_hess_eval_EUW_with_weight1(const cudaStream_t stream, fprintf(stderr, "CUDA Error of vv10 hess eval_EUW_with_weight1: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -842,6 +991,14 @@ int VXC_vv10nlc_hess_eval_D_B_in_double_grid_response_offdiagonal(const cudaStre const int ngrids, const int natm) { constexpr int n_grids_per_block = 128; + #ifdef USE_SYCL + const sycl::range<2> threads(1, n_grids_per_block); + const sycl::range<2> blocks(natm, (ngrids + n_grids_per_block - 1) / n_grids_per_block); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + vv10_hess_eval_D_B_in_double_grid_response_offdiagonal_kernel( + D_B, grid_coord, rho_weight, omega, kappa, grid_associated_atom, grid_offsets_of_atom, ngrids, natm); + }); + #else const dim3 threads(n_grids_per_block, 1); const dim3 blocks((ngrids + n_grids_per_block - 1) / n_grids_per_block, natm); vv10_hess_eval_D_B_in_double_grid_response_offdiagonal_kernel<<>>( @@ -852,6 +1009,7 @@ int VXC_vv10nlc_hess_eval_D_B_in_double_grid_response_offdiagonal(const cudaStre fprintf(stderr, "CUDA Error of vv10 hess eval_E_grgr_AB: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/gint/CMakeLists.txt b/gpu4pyscf/lib/gint/CMakeLists.txt index e81cc1deb..de911b6c6 100644 --- a/gpu4pyscf/lib/gint/CMakeLists.txt +++ b/gpu4pyscf/lib/gint/CMakeLists.txt @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v") - set(CMAKE_VERBOSE_MAKEFILE ON) -add_library(gint SHARED + +set(GPU_SRCS g2e.c pair_data.c bpcache.cu @@ -36,12 +35,22 @@ add_library(gint SHARED cart2sph.cu ) +add_library(gint SHARED ${GPU_SRCS}) + +set_target_properties(gint PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + +if (USE_SYCL) + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + target_compile_options(gint PRIVATE -x c++ -nocudainc -nocudalib) + target_link_libraries(gint PRIVATE sycl_compat gsycl) +else() + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v") + set_target_properties(gint PROPERTIES CUDA_SEPARABLE_COMPILATION ON) +endif() + #option(BUILD_SHARED_LIBS "build shared libraries" 1) #option(ENABLE_STATIC "Enforce static library build" 0) #if(ENABLE_STATIC) # set(BUILD_SHARED_LIBS 0) #endif() - -set_target_properties(gint PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_SEPARABLE_COMPILATION ON) diff --git a/gpu4pyscf/lib/gint/cart2sph.cu b/gpu4pyscf/lib/gint/cart2sph.cu index e4b2170b8..282fdab95 100644 --- a/gpu4pyscf/lib/gint/cart2sph.cu +++ b/gpu4pyscf/lib/gint/cart2sph.cu @@ -16,6 +16,36 @@ #include +// Abstracts 2D kernel launch syntax. blocks/threads must be in scope. +// TAG: unique SYCL class name for separate compilation (ignored on CUDA) +// KERNEL: kernel function, with template args if needed (e.g. func) +// ...: kernel arguments +#ifdef USE_SYCL +#include + +#define KERNEL_SETUP() \ + auto item = syclex::this_work_item::get_nd_item<2>(); \ + const int gid_x = item.get_global_id(1); \ + const int gid_y = item.get_global_id(0); + +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +// ARGS: parenthesized runtime args. Kernel-id (with any template args) is the +// trailing __VA_ARGS__ so its commas survive macro expansion. SYCL kernel name +// is generated inline per source line (unique within this translation unit). +#define LAUNCH_KERNEL(ARGS, ...) \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) { __VA_ARGS__ ARGS; }); +#else // USE_SYCL +#define KERNEL_SETUP() \ + const int gid_x = blockIdx.x * blockDim.x + threadIdx.x; \ + const int gid_y = blockIdx.y * blockDim.y + threadIdx.y; + +#define LAUNCH_KERNEL(ARGS, ...) \ + __VA_ARGS__ <<>> ARGS; +#endif // USE_SYCL + template __device__ static void cart2sph(const double *gcart, double *gsph, const int cart_stride) @@ -482,8 +512,9 @@ static void left_cart2sph_inplace(double* cartesian_matrix, const int n_ao_carte constexpr int n_cartesian_of_l = (L + 1) * (L + 2) / 2; constexpr int n_spherical_of_l = 2 * L + 1; - const int i_ao = blockIdx.x * blockDim.x + threadIdx.x; - const int i_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_ao = gid_x; // AO index (Cartesian or spherical) + const int i_bas = gid_y; // Shell (basis function) index if (i_ao >= n_ao_cartesian || i_bas >= n_bas) return; @@ -502,8 +533,9 @@ static void left_sph2cart_inplace(double* cartesian_matrix, const int n_ao_carte constexpr int n_cartesian_of_l = (L + 1) * (L + 2) / 2; constexpr int n_spherical_of_l = 2 * L + 1; - const int i_ao = blockIdx.x * blockDim.x + threadIdx.x; - const int i_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_ao = gid_x; // AO index (Cartesian or spherical) + const int i_bas = gid_y; // Shell (basis function) index if (i_ao >= n_ao_cartesian || i_bas >= n_bas) return; @@ -524,8 +556,9 @@ static void left_sph2cart(double* cartesian_matrix, const double* spherical_matr constexpr int n_cartesian_of_l = (L + 1) * (L + 2) / 2; constexpr int n_spherical_of_l = 2 * L + 1; - const int i_ao = blockIdx.x * blockDim.x + threadIdx.x; - const int i_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_ao = gid_x; // AO index (Cartesian or spherical) + const int i_bas = gid_y; // Shell (basis function) index if (i_ao >= n_right || i_bas >= n_bas) return; @@ -544,8 +577,9 @@ static void right_cart2sph_inplace(double* cartesian_matrix, const int n_ao_cart constexpr int n_cartesian_of_l = (L + 1) * (L + 2) / 2; constexpr int n_spherical_of_l = 2 * L + 1; - const int i_ao = blockIdx.x * blockDim.x + threadIdx.x; - const int i_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_ao = gid_x; // AO index (Cartesian or spherical) + const int i_bas = gid_y; // Shell (basis function) index if (i_ao >= n_ao_cartesian || i_bas >= n_bas) return; @@ -564,8 +598,9 @@ static void right_sph2cart_inplace(double* cartesian_matrix, const int n_ao_cart constexpr int n_cartesian_of_l = (L + 1) * (L + 2) / 2; constexpr int n_spherical_of_l = 2 * L + 1; - const int i_ao = blockIdx.x * blockDim.x + threadIdx.x; - const int i_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_ao = gid_x; // AO index (Cartesian or spherical) + const int i_bas = gid_y; // Shell (basis function) index if (i_ao >= n_ao_cartesian || i_bas >= n_bas) return; @@ -584,8 +619,9 @@ static void copy_spherical_cart2sph(const double* cartesian_matrix, double* sphe const int l_j, const int n_bas_j, const int cartesian_offset_j, const int spherical_offset_j, const int* d_ao_idx) { - const int i_bas = blockIdx.x * blockDim.x + threadIdx.x; - const int j_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_bas = gid_x; // i-shell (row basis function) index + const int j_bas = gid_y; // j-shell (column basis function) index if (i_bas >= n_bas_i || j_bas >= n_bas_j) return; @@ -612,8 +648,9 @@ static void copy_spherical_sph2cart(double* cartesian_matrix, const double* sphe const int l_j, const int n_bas_j, const int cartesian_offset_j, const int spherical_offset_j, const int* d_ao_idx) { - const int i_bas = blockIdx.x * blockDim.x + threadIdx.x; - const int j_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_bas = gid_x; // i-shell (row basis function) index + const int j_bas = gid_y; // j-shell (column basis function) index if (i_bas >= n_bas_i || j_bas >= n_bas_j) return; @@ -640,8 +677,9 @@ static void copy_cartesian_pad_to_unpad(const double* cartesian_matrix, double* const int l_j, const int n_bas_j, const int j_pad_offset, const int j_unpad_offset, const int* d_ao_idx) { - const int i_bas = blockIdx.x * blockDim.x + threadIdx.x; - const int j_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_bas = gid_x; // i-shell (row basis function) index + const int j_bas = gid_y; // j-shell (column basis function) index if (i_bas >= n_bas_i || j_bas >= n_bas_j) return; @@ -666,8 +704,9 @@ static void copy_cartesian_unpad_to_pad(double* cartesian_matrix, const double* const int l_j, const int n_bas_j, const int j_pad_offset, const int j_unpad_offset, const int* d_ao_idx) { - const int i_bas = blockIdx.x * blockDim.x + threadIdx.x; - const int j_bas = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); + const int i_bas = gid_x; // i-shell (row basis function) index + const int j_bas = gid_y; // j-shell (column basis function) index if (i_bas >= n_bas_i || j_bas >= n_bas_j) return; @@ -690,8 +729,14 @@ static void left_cart2cart(double* destination_matrix, const double* source_matr const int n_right, const int n_ao_copy, const int i_destination_offset, const int i_source_offset, const int* d_ao_idx) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i_right = item.get_global_id(1); + const int i_left = item.get_global_id(0); + #else const int i_right = blockIdx.x * blockDim.x + threadIdx.x; const int i_left = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (i_right >= n_right || i_left >= n_ao_copy) return; @@ -714,20 +759,25 @@ extern "C" { const int l_i = l_of_group[i_group]; const int n_bas = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(16, 16); + const sycl::range<2> blocks((n_bas + threads[0] - 1) / threads[0], (n_ao_cartesian + threads[1] - 1) / threads[1]); + #else const dim3 threads(16, 16); const dim3 blocks((n_ao_cartesian + threads.x - 1) / threads.x, (n_bas + threads.y - 1) / threads.y); + #endif switch (l_i) { - case 0: left_cart2sph_inplace< 0> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 1: left_cart2sph_inplace< 1> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 2: left_cart2sph_inplace< 2> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 3: left_cart2sph_inplace< 3> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 4: left_cart2sph_inplace< 4> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 5: left_cart2sph_inplace< 5> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 6: left_cart2sph_inplace< 6> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 7: left_cart2sph_inplace< 7> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 8: left_cart2sph_inplace< 8> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 9: left_cart2sph_inplace< 9> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 10: left_cart2sph_inplace<10> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; + case 0: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 0>) break; + case 1: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 1>) break; + case 2: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 2>) break; + case 3: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 3>) break; + case 4: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 4>) break; + case 5: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 5>) break; + case 6: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 6>) break; + case 7: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 7>) break; + case 8: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 8>) break; + case 9: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace< 9>) break; + case 10: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_cart2sph_inplace<10>) break; default: printf("l_i = %d not supported for cart2sph_C_mat_CT_with_padding(), max_L = 10\n", l_i); fprintf(stderr, "l_i = %d not supported for cart2sph_C_mat_CT_with_padding(), max_L = 10\n", l_i); @@ -742,20 +792,25 @@ extern "C" { const int l_i = l_of_group[i_group]; const int n_bas = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(16, 16); + const sycl::range<2> blocks((n_bas + threads[0] - 1) / threads[0], (n_ao_cartesian + threads[1] - 1) / threads[1]); + #else const dim3 threads(16, 16); const dim3 blocks((n_ao_cartesian + threads.x - 1) / threads.x, (n_bas + threads.y - 1) / threads.y); + #endif switch (l_i) { - case 0: right_cart2sph_inplace< 0> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 1: right_cart2sph_inplace< 1> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 2: right_cart2sph_inplace< 2> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 3: right_cart2sph_inplace< 3> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 4: right_cart2sph_inplace< 4> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 5: right_cart2sph_inplace< 5> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 6: right_cart2sph_inplace< 6> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 7: right_cart2sph_inplace< 7> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 8: right_cart2sph_inplace< 8> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 9: right_cart2sph_inplace< 9> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 10: right_cart2sph_inplace<10> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; + case 0: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 0>) break; + case 1: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 1>) break; + case 2: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 2>) break; + case 3: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 3>) break; + case 4: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 4>) break; + case 5: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 5>) break; + case 6: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 6>) break; + case 7: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 7>) break; + case 8: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 8>) break; + case 9: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace< 9>) break; + case 10: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_cart2sph_inplace<10>) break; default: printf("l_i = %d not supported for cart2sph_C_mat_CT_with_padding(), max_L = 10\n", l_i); fprintf(stderr, "l_i = %d not supported for cart2sph_C_mat_CT_with_padding(), max_L = 10\n", l_i); @@ -770,19 +825,21 @@ extern "C" { for (int i_group = 0; i_group < n_l_ctr_group; i_group++) { const int l_i = l_of_group[i_group]; const int n_bas_i = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; - + int j_cartesian_offset = 0; int j_spherical_offset = 0; for (int j_group = 0; j_group < n_l_ctr_group; j_group++) { const int l_j = l_of_group[j_group]; const int n_bas_j = n_total_bas_of_group[j_group] - n_pad_bas_of_group[j_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(32, 32); + const sycl::range<2> blocks((n_bas_j + threads[0] - 1) / threads[0], (n_bas_i + threads[1] - 1) / threads[1]); + #else const dim3 threads(32, 32); const dim3 blocks((n_bas_i + threads.x - 1) / threads.x, (n_bas_j + threads.y - 1) / threads.y); - copy_spherical_cart2sph<<>>(cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, - l_i, n_bas_i, i_cartesian_offset, i_spherical_offset, - l_j, n_bas_j, j_cartesian_offset, j_spherical_offset, - d_ao_idx); + #endif + LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, l_i, n_bas_i, i_cartesian_offset, i_spherical_offset, l_j, n_bas_j, j_cartesian_offset, j_spherical_offset, d_ao_idx), copy_spherical_cart2sph) j_cartesian_offset += n_total_bas_of_group[j_group] * ((l_j + 1) * (l_j + 2) / 2); j_spherical_offset += n_bas_j * (l_j * 2 + 1); @@ -796,19 +853,21 @@ extern "C" { for (int i_group = 0; i_group < n_l_ctr_group; i_group++) { const int l_i = l_of_group[i_group]; const int n_bas_i = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; - + int j_pad_offset = 0; int j_unpad_offset = 0; for (int j_group = 0; j_group < n_l_ctr_group; j_group++) { const int l_j = l_of_group[j_group]; const int n_bas_j = n_total_bas_of_group[j_group] - n_pad_bas_of_group[j_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(32, 32); + const sycl::range<2> blocks((n_bas_j + threads[1] - 1) / threads[1], (n_bas_i + threads[0] - 1) / threads[0]); + #else const dim3 threads(32, 32); const dim3 blocks((n_bas_i + threads.x - 1) / threads.x, (n_bas_j + threads.y - 1) / threads.y); - copy_cartesian_pad_to_unpad<<>>(cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, - l_i, n_bas_i, i_pad_offset, i_unpad_offset, - l_j, n_bas_j, j_pad_offset, j_unpad_offset, - d_ao_idx); + #endif + LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, l_i, n_bas_i, i_pad_offset, i_unpad_offset, l_j, n_bas_j, j_pad_offset, j_unpad_offset, d_ao_idx), copy_cartesian_pad_to_unpad) j_pad_offset += n_total_bas_of_group[j_group] * ((l_j + 1) * (l_j + 2) / 2); j_unpad_offset += n_bas_j * ((l_j + 1) * (l_j + 2) / 2); @@ -833,19 +892,21 @@ extern "C" { for (int i_group = 0; i_group < n_l_ctr_group; i_group++) { const int l_i = l_of_group[i_group]; const int n_bas_i = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; - + int j_cartesian_offset = 0; int j_spherical_offset = 0; for (int j_group = 0; j_group < n_l_ctr_group; j_group++) { const int l_j = l_of_group[j_group]; const int n_bas_j = n_total_bas_of_group[j_group] - n_pad_bas_of_group[j_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(32, 32); + const sycl::range<2> blocks((n_bas_j + threads[1] - 1) / threads[1], (n_bas_i + threads[0] - 1) / threads[0]); + #else const dim3 threads(32, 32); const dim3 blocks((n_bas_i + threads.x - 1) / threads.x, (n_bas_j + threads.y - 1) / threads.y); - copy_spherical_sph2cart<<>>(cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, - l_i, n_bas_i, i_cartesian_offset, i_spherical_offset, - l_j, n_bas_j, j_cartesian_offset, j_spherical_offset, - d_ao_idx); + #endif + LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, l_i, n_bas_i, i_cartesian_offset, i_spherical_offset, l_j, n_bas_j, j_cartesian_offset, j_spherical_offset, d_ao_idx), copy_spherical_sph2cart) j_cartesian_offset += n_total_bas_of_group[j_group] * ((l_j + 1) * (l_j + 2) / 2); j_spherical_offset += n_bas_j * (l_j * 2 + 1); @@ -859,20 +920,25 @@ extern "C" { const int l_i = l_of_group[i_group]; const int n_bas = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(16, 16); + const sycl::range<2> blocks((n_bas + threads[0] - 1) / threads[0], (n_ao_cartesian + threads[1] - 1) / threads[1]); + #else const dim3 threads(16, 16); const dim3 blocks((n_ao_cartesian + threads.x - 1) / threads.x, (n_bas + threads.y - 1) / threads.y); + #endif switch (l_i) { - case 0: left_sph2cart_inplace< 0> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 1: left_sph2cart_inplace< 1> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 2: left_sph2cart_inplace< 2> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 3: left_sph2cart_inplace< 3> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 4: left_sph2cart_inplace< 4> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 5: left_sph2cart_inplace< 5> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 6: left_sph2cart_inplace< 6> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 7: left_sph2cart_inplace< 7> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 8: left_sph2cart_inplace< 8> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 9: left_sph2cart_inplace< 9> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 10: left_sph2cart_inplace<10> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; + case 0: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 0>) break; + case 1: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 1>) break; + case 2: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 2>) break; + case 3: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 3>) break; + case 4: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 4>) break; + case 5: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 5>) break; + case 6: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 6>) break; + case 7: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 7>) break; + case 8: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 8>) break; + case 9: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace< 9>) break; + case 10: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), left_sph2cart_inplace<10>) break; default: printf("l_i = %d not supported for cart2sph_CT_mat_C_with_padding(), max_L = 10\n", l_i); fprintf(stderr, "l_i = %d not supported for cart2sph_CT_mat_C_with_padding(), max_L = 10\n", l_i); @@ -887,20 +953,25 @@ extern "C" { const int l_i = l_of_group[i_group]; const int n_bas = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; +#ifdef USE_SYCL + const sycl::range<2> threads(16, 16); + const sycl::range<2> blocks((n_bas + threads[0] - 1) / threads[0], (n_ao_cartesian + threads[1] - 1) / threads[1]); + #else const dim3 threads(16, 16); const dim3 blocks((n_ao_cartesian + threads.x - 1) / threads.x, (n_bas + threads.y - 1) / threads.y); + #endif switch (l_i) { - case 0: right_sph2cart_inplace< 0> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 1: right_sph2cart_inplace< 1> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 2: right_sph2cart_inplace< 2> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 3: right_sph2cart_inplace< 3> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 4: right_sph2cart_inplace< 4> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 5: right_sph2cart_inplace< 5> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 6: right_sph2cart_inplace< 6> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 7: right_sph2cart_inplace< 7> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 8: right_sph2cart_inplace< 8> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 9: right_sph2cart_inplace< 9> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; - case 10: right_sph2cart_inplace<10> <<>>(cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset); break; + case 0: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 0>) break; + case 1: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 1>) break; + case 2: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 2>) break; + case 3: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 3>) break; + case 4: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 4>) break; + case 5: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 5>) break; + case 6: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 6>) break; + case 7: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 7>) break; + case 8: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 8>) break; + case 9: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace< 9>) break; + case 10: LAUNCH_KERNEL((cartesian_matrix, n_ao_cartesian, n_bas, i_cartesian_offset), right_sph2cart_inplace<10>) break; default: printf("l_i = %d not supported for cart2sph_CT_mat_C_with_padding(), max_L = 10\n", l_i); fprintf(stderr, "l_i = %d not supported for cart2sph_CT_mat_C_with_padding(), max_L = 10\n", l_i); @@ -915,19 +986,21 @@ extern "C" { for (int i_group = 0; i_group < n_l_ctr_group; i_group++) { const int l_i = l_of_group[i_group]; const int n_bas_i = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; - + int j_pad_offset = 0; int j_unpad_offset = 0; for (int j_group = 0; j_group < n_l_ctr_group; j_group++) { const int l_j = l_of_group[j_group]; const int n_bas_j = n_total_bas_of_group[j_group] - n_pad_bas_of_group[j_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(32, 32); + const sycl::range<2> blocks((n_bas_j + threads[1] - 1) / threads[1], (n_bas_i + threads[0] - 1) / threads[0]); + #else const dim3 threads(32, 32); const dim3 blocks((n_bas_i + threads.x - 1) / threads.x, (n_bas_j + threads.y - 1) / threads.y); - copy_cartesian_unpad_to_pad<<>>(cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, - l_i, n_bas_i, i_pad_offset, i_unpad_offset, - l_j, n_bas_j, j_pad_offset, j_unpad_offset, - d_ao_idx); + #endif + LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_ao_cartesian, n_ao_spherical, l_i, n_bas_i, i_pad_offset, i_unpad_offset, l_j, n_bas_j, j_pad_offset, j_unpad_offset, d_ao_idx), copy_cartesian_unpad_to_pad) j_pad_offset += n_total_bas_of_group[j_group] * ((l_j + 1) * (l_j + 2) / 2); j_unpad_offset += n_bas_j * ((l_j + 1) * (l_j + 2) / 2); @@ -953,20 +1026,25 @@ extern "C" { const int l_i = l_of_group[i_group]; const int n_bas = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; + #ifdef USE_SYCL + const sycl::range<2> threads(16, 16); + const sycl::range<2> blocks((n_bas + threads[0] - 1) / threads[0], (n_right + threads[1] - 1) / threads[1]); + #else const dim3 threads(16, 16); const dim3 blocks((n_right + threads.x - 1) / threads.x, (n_bas + threads.y - 1) / threads.y); + #endif switch (l_i) { - case 0: left_sph2cart< 0> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 1: left_sph2cart< 1> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 2: left_sph2cart< 2> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 3: left_sph2cart< 3> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 4: left_sph2cart< 4> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 5: left_sph2cart< 5> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 6: left_sph2cart< 6> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 7: left_sph2cart< 7> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 8: left_sph2cart< 8> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 9: left_sph2cart< 9> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; - case 10: left_sph2cart<10> <<>>(cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx); break; + case 0: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 0>) break; + case 1: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 1>) break; + case 2: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 2>) break; + case 3: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 3>) break; + case 4: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 4>) break; + case 5: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 5>) break; + case 6: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 6>) break; + case 7: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 7>) break; + case 8: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 8>) break; + case 9: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart< 9>) break; + case 10: LAUNCH_KERNEL((cartesian_matrix, spherical_matrix, n_right, n_bas, i_cartesian_offset, i_spherical_offset, d_ao_idx), left_sph2cart<10>) break; default: printf("l_i = %d not supported for cart2sph_C_mat_with_padding(), max_L = 10\n", l_i); fprintf(stderr, "l_i = %d not supported for cart2sph_C_mat_with_padding(), max_L = 10\n", l_i); @@ -984,10 +1062,18 @@ extern "C" { const int n_bas = n_total_bas_of_group[i_group] - n_pad_bas_of_group[i_group]; const int n_cartesian_of_l = (l_i + 1) * (l_i + 2) / 2; + #ifdef USE_SYCL + const sycl::range<2> threads(16, 16); + const sycl::range<2> blocks((n_bas * n_cartesian_of_l + threads[0] - 1) / threads[0], (n_right + threads[1] - 1) / threads[1]); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + left_cart2cart(cartesian_matrix, spherical_matrix, + n_right, n_bas * n_cartesian_of_l, i_pad_offset, i_unpad_offset, d_ao_idx); }); + #else const dim3 threads(16, 16); const dim3 blocks((n_right + threads.x - 1) / threads.x, (n_bas * n_cartesian_of_l + threads.y - 1) / threads.y); left_cart2cart<<>>(cartesian_matrix, spherical_matrix, n_right, n_bas * n_cartesian_of_l, i_pad_offset, i_unpad_offset, d_ao_idx); + #endif i_pad_offset += n_total_bas_of_group[i_group] * n_cartesian_of_l; i_unpad_offset += n_bas * n_cartesian_of_l; @@ -997,3 +1083,6 @@ extern "C" { return 0; } } + +#undef KERNEL_SETUP +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/cint2e.cuh b/gpu4pyscf/lib/gint/cint2e.cuh index 17dd5525d..ce2667b06 100644 --- a/gpu4pyscf/lib/gint/cint2e.cuh +++ b/gpu4pyscf/lib/gint/cint2e.cuh @@ -18,9 +18,58 @@ #include "gint.h" +#ifdef USE_SYCL + +extern SYCL_EXTERNAL sycl_device_global s_bpcache; + +// Generated with GINTinit_index1d_xyz +// Look into constant.cu for details +inline constexpr int c_idx[TOT_NF*3] = { + 0, 1, 0, 0, 2, 1, 1, 0, 0, 0, 3, 2, 2, 1, 1, 1, 0, 0, 0, 0, 4, 3, 3, + 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 5, 4, 4, 3, 3, 3, 2, 2, 2, 2, 1, + 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 6, 5, 5, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, + 2, 2, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 2, + 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, + 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, + 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, 1, 0, + 6, 5, 4, 3, 2, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 2, 0, 0, 1, 0, 1, 2, + 0, 1, 2, 3, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 0, 1, 0, + 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 0, 1, 0, 1, 2, + 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6}; + +inline constexpr int c_l_locs[GPU_LMAX+2] = {0, 1, 4, 10, 20, 35, 56, 84}; + +#else // USE_SYCL //extern __constant__ GINTEnvVars c_envs; extern __constant__ BasisProdCache c_bpcache; //extern __constant__ int16_t c_idx4c[NFffff*3]; extern __constant__ int c_idx[TOT_NF*3]; extern __constant__ int c_l_locs[GPU_LMAX+2]; +#endif // USE_SYCL + +// Abstracts 2D kernel thread-index setup for task_ij/task_kl kernels. Used 79x across gint/. +#ifdef USE_SYCL +#define KERNEL_SETUP() \ + auto item = syclex::this_work_item::get_nd_item<2>(); \ + const int task_ij = item.get_global_id(1); \ + const int task_kl = item.get_global_id(0); \ + const auto& c_bpcache = s_bpcache.get(); +#else +#define KERNEL_SETUP() \ + const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; \ + const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; +#endif + +// Abstracts 2D kernel local thread-index setup for threadIdx_x/blockDim_x kernels. Used 9x across gint/. +#ifdef USE_SYCL +#define KERNEL_SETUP_LOCAL() \ + auto item = syclex::this_work_item::get_nd_item<2>(); \ + const int threadIdx_x = item.get_local_id(1); \ + const int blockDim_x = item.get_local_range(1); \ + const auto& c_bpcache = s_bpcache.get(); +#else +#define KERNEL_SETUP_LOCAL() \ + const int threadIdx_x = threadIdx.x; \ + const int blockDim_x = blockDim.x; +#endif diff --git a/gpu4pyscf/lib/gint/constant.cu b/gpu4pyscf/lib/gint/constant.cu index 4c71be92e..b488b6360 100644 --- a/gpu4pyscf/lib/gint/constant.cu +++ b/gpu4pyscf/lib/gint/constant.cu @@ -16,22 +16,31 @@ #include "gint.h" +#ifdef USE_SYCL + +//__constant__ GINTEnvVars c_envs; +SYCL_EXTERNAL sycl_device_global s_bpcache; +//__constant__ int16_t c_idx4c[NFffff*3]; + +#else // USE_SYCL + //__constant__ GINTEnvVars c_envs; __constant__ BasisProdCache c_bpcache; //__constant__ int16_t c_idx4c[NFffff*3]; // Generated with GINTinit_index1d_xyz __constant__ int c_idx[TOT_NF*3] = { - 0, 1, 0, 0, 2, 1, 1, 0, 0, 0, 3, 2, 2, 1, 1, 1, 0, 0, 0, 0, 4, 3, 3, - 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 5, 4, 4, 3, 3, 3, 2, 2, 2, 2, 1, - 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 6, 5, 5, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, - 2, 2, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 2, - 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, - 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, - 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, 1, 0, - 6, 5, 4, 3, 2, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 2, 0, 0, 1, 0, 1, 2, - 0, 1, 2, 3, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 0, 1, 0, - 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 0, 1, 0, 1, 2, + 0, 1, 0, 0, 2, 1, 1, 0, 0, 0, 3, 2, 2, 1, 1, 1, 0, 0, 0, 0, 4, 3, 3, + 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 5, 4, 4, 3, 3, 3, 2, 2, 2, 2, 1, + 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 6, 5, 5, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, + 2, 2, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 2, + 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, + 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, + 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, 1, 0, + 6, 5, 4, 3, 2, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 2, 0, 0, 1, 0, 1, 2, + 0, 1, 2, 3, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 0, 1, 0, + 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6}; __constant__ int c_l_locs[GPU_LMAX+2] = {0, 1, 4, 10, 20, 35, 56, 84}; +#endif // USE_SYCL diff --git a/gpu4pyscf/lib/gint/cuda_alloc.cuh b/gpu4pyscf/lib/gint/cuda_alloc.cuh index a33196660..c6b289457 100644 --- a/gpu4pyscf/lib/gint/cuda_alloc.cuh +++ b/gpu4pyscf/lib/gint/cuda_alloc.cuh @@ -13,12 +13,13 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - #pragma once #include #include +#ifndef USE_SYCL + // copy from samples/common/inc/helper_cuda.h template int check(T result, char const *const func, const char *const file, @@ -48,3 +49,65 @@ int check(T result, char const *const func, const char *const file, MALLOC(type, dst, size); \ checkCudaErrors(cudaMemcpy(dst, src, sizeof(type) * (size), cudaMemcpyHostToDevice)) +#else // !USE_SYCL + +// Function to check SYCL errors +template +void check(T result, char const *const func, const char *const file, int const line) { + if (result) { + std::cerr << "SYCL error at " << file << ":" << line << " code=" << result << " \"" << func << "\" \n"; + std::exit(EXIT_FAILURE); + } +} +#define checkCudaErrors(val) (val) + + +#define MALLOC(type, var, size) \ + type *var = sycl::malloc_device(size, *(sycl_get_queue())); \ + if (var == nullptr) { \ + std::cerr << "Memory allocation failed for " #var " at " __FILE__ ":" << __LINE__ << std::endl; \ + std::exit(EXIT_FAILURE); \ + } + +// Drain the queue before releasing device memory. +// +// The CUDA path uses cudaFree(), which IMPLICITLY SYNCHRONIZES the device, so +// no kernel can still be reading the buffer when it is unmapped. sycl::free() +// has no such guarantee: it unmaps immediately, even with work in flight. +// +// This matters because bpcache pointers (a12/e12/x12/... aliased into the +// single d_aexyz block, bas_coords, bas_atm, bas_pair2bra, ao_loc) are handed +// to the gint kernels INDIRECTLY -- the kernels are launched with +// zeKernelSetIndirectAccess(flags=0x7), so the runtime cannot see them as +// arguments and cannot keep them alive. When GINTdel_basis_prod() runs from +// Python teardown (intopt.clear() / __del__) while int3c1e/int3c2e kernels are +// still executing, the pages are unmapped underneath them: +// +// zeEventQueryStatus(...) -> ZE_RESULT_NOT_READY (kernel still running) +// zeMemFree(0xff000002e8a00000) (freed anyway) +// Segmentation fault from GPU at 0xff000002e8b05000 (base + 0x105000) +// +// The wait restores cudaFree's implicit-sync semantics. These frees are on +// teardown paths, not hot paths, so the cost is negligible. +#define FREE(var) \ + do { \ + sycl_get_queue()->wait(); \ + sycl::free(var, *(sycl_get_queue())); \ + } while (0) + +#define MEMSET(addr, val, size) \ + { \ + sycl_get_queue()->submit([&](sycl::handler& cgh) { \ + cgh.memset(addr, val, size); \ + }).wait(); \ + } + +#define DEVICE_INIT(type, dst, src, size) \ + MALLOC(type, dst, size); \ + { \ + sycl_get_queue()->submit([&](sycl::handler& cgh) { \ + cgh.memcpy(dst, src, sizeof(type) * (size)); \ + }).wait(); \ + } + +#endif // USE_SYCL diff --git a/gpu4pyscf/lib/gint/fill_ints.cu b/gpu4pyscf/lib/gint/fill_ints.cu index 461d45021..f1827db84 100644 --- a/gpu4pyscf/lib/gint/fill_ints.cu +++ b/gpu4pyscf/lib/gint/fill_ints.cu @@ -22,6 +22,9 @@ __device__ void GINTwrite_ints_s2(ERITensor eri, double* __restrict__ gout, int ish, int jsh, int ksh, int lsh) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif int *ao_loc = c_bpcache.ao_loc; size_t jstride = eri.stride_j; size_t kstride = eri.stride_k; @@ -56,6 +59,9 @@ __device__ void GINTwrite_ints_sph_s2(ERITensor eri, double* __restrict__ gout, int ish, int jsh, int ksh, int lsh) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif int *ao_loc = c_bpcache.ao_loc; size_t jstride = eri.stride_j; size_t kstride = eri.stride_k; diff --git a/gpu4pyscf/lib/gint/g1e.cu b/gpu4pyscf/lib/gint/g1e.cu index af377139e..d5ca2648c 100644 --- a/gpu4pyscf/lib/gint/g1e.cu +++ b/gpu4pyscf/lib/gint/g1e.cu @@ -24,6 +24,9 @@ static void GINT_g1e(double* __restrict__ g, const double* __restrict__ grid_poi const int ish, const int jsh, const int prim_ij, const int i_l, const int j_l, const double charge_exponent, const double omega) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif const double* __restrict__ a12 = c_bpcache.a12; const double* __restrict__ e12 = c_bpcache.e12; const double* __restrict__ x12 = c_bpcache.x12; @@ -141,6 +144,9 @@ static void GINT_g1e_save_u2(double* __restrict__ g, double* __restrict__ u2_sav const int ish, const int jsh, const int prim_ij, const int i_l, const int j_l, const double charge_exponent, const double omega) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif const double* __restrict__ a12 = c_bpcache.a12; const double* __restrict__ e12 = c_bpcache.e12; const double* __restrict__ x12 = c_bpcache.x12; @@ -257,6 +263,9 @@ __device__ static void GINT_g1e_without_hrr(double* __restrict__ g, const double grid_x, const double grid_y, const double grid_z, const int ish, const int prim_ij, const double charge_exponent, const double omega) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif constexpr int NROOTS = L_SUM / 2 + 1; const double* __restrict__ a12 = c_bpcache.a12; @@ -354,6 +363,9 @@ __device__ static void GINT_g1e_without_hrr_save_u2(double* __restrict__ g, double* __restrict__ u2_save, const double grid_x, const double grid_y, const double grid_z, const int ish, const int prim_ij, const double charge_exponent, const double omega) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif constexpr int NROOTS = L_SUM / 2 + 1; const double* __restrict__ a12 = c_bpcache.a12; diff --git a/gpu4pyscf/lib/gint/g1e_ip_root_1.cu b/gpu4pyscf/lib/gint/g1e_ip_root_1.cu index 935c06028..325fd017c 100644 --- a/gpu4pyscf/lib/gint/g1e_ip_root_1.cu +++ b/gpu4pyscf/lib/gint/g1e_ip_root_1.cu @@ -24,8 +24,15 @@ static void GINTfill_int3c1e_ip_kernel00(double* output, const BasisProdOffsets { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; @@ -124,7 +131,17 @@ static void GINTfill_int3c1e_ip1_charge_contracted_kernel00(double* output, cons { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -154,7 +171,7 @@ static void GINTfill_int3c1e_ip1_charge_contracted_kernel00(double* output, cons double deri_dAx_grid_sum = 0; double deri_dAy_grid_sum = 0; double deri_dAz_grid_sum = 0; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double Cx = grid_point[0]; const double Cy = grid_point[1]; @@ -218,11 +235,19 @@ __global__ static void GINTfill_int3c1e_ip1_density_contracted_kernel00(double* output, const BasisProdOffsets offsets, const int nprim_ij, const double* density, const int* aoslice, const int nao, const double omega, const double* grid_points, const double* charge_exponents) + { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; @@ -317,7 +342,18 @@ static void GINTfill_int3c1e_ip2_density_contracted_kernel00(double* output, con { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_grid = item.get_global_id(0); + const int thread_x_id = item.get_global_id(1); + const int total_threads_x = item.get_global_range(1); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + const int thread_x_id = blockIdx.x * blockDim.x + threadIdx.x; + const int total_threads_x = gridDim.x * blockDim.x; + #endif + if (task_grid >= ngrids) { return; } @@ -331,7 +367,7 @@ static void GINTfill_int3c1e_ip2_density_contracted_kernel00(double* output, con double deri_dCx_pair_sum = 0.0; double deri_dCy_pair_sum = 0.0; double deri_dCz_pair_sum = 0.0; - for (int task_ij = blockIdx.x * blockDim.x + threadIdx.x; task_ij < ntasks_ij; task_ij += gridDim.x * blockDim.x) { + for (int task_ij = thread_x_id; task_ij < ntasks_ij; task_ij += total_threads_x) { const int bas_ij = offsets.bas_ij + task_ij; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; @@ -392,9 +428,15 @@ static void GINTfill_int3c1e_ip2_charge_contracted_kernel00(double* output, cons { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; - + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; } diff --git a/gpu4pyscf/lib/gint/g1e_root_1.cu b/gpu4pyscf/lib/gint/g1e_root_1.cu index ba359017d..0ccd76bf3 100644 --- a/gpu4pyscf/lib/gint/g1e_root_1.cu +++ b/gpu4pyscf/lib/gint/g1e_root_1.cu @@ -24,8 +24,15 @@ static void GINTfill_int3c1e_kernel00(double* output, const BasisProdOffsets off { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; @@ -91,7 +98,17 @@ static void GINTfill_int3c1e_charge_contracted_kernel00(double* output, const Ba { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -110,7 +127,7 @@ static void GINTfill_int3c1e_charge_contracted_kernel00(double* output, const Ba const double* __restrict__ z12 = c_bpcache.z12; double eri_grid_sum = 0.0; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double Cx = grid_point[0]; const double Cy = grid_point[1]; @@ -164,7 +181,18 @@ static void GINTfill_int3c1e_density_contracted_kernel00(double* output, const d { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_grid = item.get_global_id(0); + const int thread_x_id = item.get_global_id(1); + const int total_threads_x = item.get_global_range(1); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + const int thread_x_id = blockIdx.x * blockDim.x + threadIdx.x; + const int total_threads_x = gridDim.x * blockDim.x; + #endif + if (task_grid >= ngrids) { return; } @@ -176,7 +204,7 @@ static void GINTfill_int3c1e_density_contracted_kernel00(double* output, const d const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; double eri_pair_sum = 0.0; - for (int task_ij = blockIdx.x * blockDim.x + threadIdx.x; task_ij < ntasks_ij; task_ij += gridDim.x * blockDim.x) { + for (int task_ij = thread_x_id; task_ij < ntasks_ij; task_ij += total_threads_x) { const int bas_ij = offsets.bas_ij + task_ij; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; // const int* bas_pair2bra = c_bpcache.bas_pair2bra; @@ -232,8 +260,15 @@ static void GINTfill_int3c1e_kernel10(double* output, const BasisProdOffsets off { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; @@ -327,7 +362,18 @@ static void GINTfill_int3c1e_charge_contracted_kernel10(double* output, const Ba { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif + if (task_ij >= ntasks_ij) { return; } @@ -356,7 +402,7 @@ static void GINTfill_int3c1e_charge_contracted_kernel10(double* output, const Ba double eri_grid_sum_x = 0; double eri_grid_sum_y = 0; double eri_grid_sum_z = 0; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double Cx = grid_point[0]; const double Cy = grid_point[1]; @@ -431,7 +477,18 @@ static void GINTfill_int3c1e_density_contracted_kernel10(double* output, const d { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_grid = item.get_global_id(0); + const int thread_x_id = item.get_global_id(1); + const int total_threads_x = item.get_global_range(1); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + const int thread_x_id = blockIdx.x * blockDim.x + threadIdx.x; + const int total_threads_x = gridDim.x * blockDim.x; + #endif + if (task_grid >= ngrids) { return; } @@ -443,7 +500,7 @@ static void GINTfill_int3c1e_density_contracted_kernel10(double* output, const d const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; double eri_pair_sum = 0.0; - for (int task_ij = blockIdx.x * blockDim.x + threadIdx.x; task_ij < ntasks_ij; task_ij += gridDim.x * blockDim.x) { + for (int task_ij = thread_x_id; task_ij < ntasks_ij; task_ij += total_threads_x) { const int bas_ij = offsets.bas_ij + task_ij; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; const int* bas_pair2bra = c_bpcache.bas_pair2bra; diff --git a/gpu4pyscf/lib/gint/g2e.cu b/gpu4pyscf/lib/gint/g2e.cu index 00b6418ad..60deacba3 100644 --- a/gpu4pyscf/lib/gint/g2e.cu +++ b/gpu4pyscf/lib/gint/g2e.cu @@ -26,6 +26,9 @@ template __device__ static void GINTg0_2e_2d4d(GINTEnvVars envs, double* __restrict__ g, double norm, int ish, int jsh, int ksh, int lsh, int prim_ij, int prim_kl) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif double* __restrict__ a12 = c_bpcache.a12; double* __restrict__ e12 = c_bpcache.e12; double* __restrict__ x12 = c_bpcache.x12; @@ -430,6 +433,15 @@ static void GINTg0_int3c2e_shared(GINTEnvVars envs, double* __restrict__ g0, const int ish, const int jsh, const int ksh, const int prim_ij, const int prim_kl) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int threadIdx_x = item.get_local_id(1); + const int blockDim_x = item.get_local_range(1); + const auto& c_bpcache = s_bpcache.get(); +#else + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; +#endif double* __restrict__ a12 = c_bpcache.a12; double* __restrict__ e12 = c_bpcache.e12; double* __restrict__ x12 = c_bpcache.x12; @@ -486,14 +498,14 @@ static void GINTg0_int3c2e_shared(GINTEnvVars envs, double* __restrict__ g0, const int gsize = envs.g_size; __syncthreads(); - for (int i = threadIdx.x; i < nrys_roots; i += blockDim.x) { + for (int i = threadIdx_x; i < nrys_roots; i += blockDim_x) { g0[i] = envs.fac; g0[i+gsize] = fac; g0[i+2*gsize] = weight; } __syncthreads(); - for (int tx = threadIdx.x; tx < nrys_roots*3; tx += blockDim.x) { + for (int tx = threadIdx_x; tx < nrys_roots*3; tx += blockDim_x) { const int iroot = tx % nrys_roots; const int ix = tx / nrys_roots; double *gx = g0 + ix * envs.g_size + iroot; @@ -603,6 +615,10 @@ static void GINTg0_int3c2e(GINTEnvVars envs, double* __restrict__ g, const double norm, const int ish, const int jsh, const int ksh, const int prim_ij, const int prim_kl) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif + double* __restrict__ a12 = c_bpcache.a12; double* __restrict__ e12 = c_bpcache.e12; double* __restrict__ x12 = c_bpcache.x12; @@ -918,6 +934,10 @@ static void GINTg0_int3c2e(GINTEnvVars envs, double* __restrict__ g, const int ish, const int jsh, const int ksh, const int prim_ij, const int prim_kl) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif + double* __restrict__ a12 = c_bpcache.a12; double* __restrict__ e12 = c_bpcache.e12; double* __restrict__ x12 = c_bpcache.x12; diff --git a/gpu4pyscf/lib/gint/g2e_root1.cu b/gpu4pyscf/lib/gint/g2e_root1.cu index b636ee247..95db847da 100644 --- a/gpu4pyscf/lib/gint/g2e_root1.cu +++ b/gpu4pyscf/lib/gint/g2e_root1.cu @@ -19,8 +19,7 @@ static void GINTfill_int2e_kernel0000(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -71,7 +70,7 @@ static void GINTfill_int2e_kernel0000(GINTEnvVars envs, ERITensor eri, BasisProd } gout0 += fac; } } - + int jstride = eri.stride_j; int kstride = eri.stride_k; int lstride = eri.stride_l; @@ -89,8 +88,7 @@ static void GINTfill_int2e_kernel0010(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -204,8 +202,7 @@ static void GINTfill_int2e_kernel1000(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } diff --git a/gpu4pyscf/lib/gint/g2e_root2.cu b/gpu4pyscf/lib/gint/g2e_root2.cu index 97a3e4c3f..1b87c4988 100644 --- a/gpu4pyscf/lib/gint/g2e_root2.cu +++ b/gpu4pyscf/lib/gint/g2e_root2.cu @@ -19,8 +19,7 @@ static void GINTfill_int2e_kernel0011(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -87,7 +86,7 @@ static void GINTfill_int2e_kernel0011(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -169,8 +168,7 @@ static void GINTfill_int2e_kernel0020(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -231,7 +229,7 @@ static void GINTfill_int2e_kernel0020(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -301,8 +299,7 @@ static void GINTfill_int2e_kernel0021(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -378,7 +375,7 @@ static void GINTfill_int2e_kernel0021(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -493,8 +490,7 @@ static void GINTfill_int2e_kernel0030(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -559,7 +555,7 @@ static void GINTfill_int2e_kernel0030(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -644,8 +640,7 @@ static void GINTfill_int2e_kernel1010(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -712,7 +707,7 @@ static void GINTfill_int2e_kernel1010(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -797,8 +792,7 @@ static void GINTfill_int2e_kernel1011(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -886,7 +880,7 @@ static void GINTfill_int2e_kernel1011(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1038,8 +1032,7 @@ static void GINTfill_int2e_kernel1020(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1115,7 +1108,7 @@ static void GINTfill_int2e_kernel1020(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1234,8 +1227,7 @@ static void GINTfill_int2e_kernel1100(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1302,7 +1294,7 @@ static void GINTfill_int2e_kernel1100(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1384,8 +1376,7 @@ static void GINTfill_int2e_kernel1110(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1473,7 +1464,7 @@ static void GINTfill_int2e_kernel1110(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1625,8 +1616,7 @@ static void GINTfill_int2e_kernel2000(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1687,7 +1677,7 @@ static void GINTfill_int2e_kernel2000(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1757,8 +1747,7 @@ static void GINTfill_int2e_kernel2010(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1834,7 +1823,7 @@ static void GINTfill_int2e_kernel2010(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1953,8 +1942,7 @@ static void GINTfill_int2e_kernel2100(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -2030,7 +2018,7 @@ static void GINTfill_int2e_kernel2100(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -2145,8 +2133,7 @@ static void GINTfill_int2e_kernel3000(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -2211,7 +2198,7 @@ static void GINTfill_int2e_kernel3000(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -2290,4 +2277,3 @@ static void GINTfill_int2e_kernel3000(GINTEnvVars envs, ERITensor eri, BasisProd eri_ij[9] = gout9; //eri_ji[9*jstride] = gout9; } - diff --git a/gpu4pyscf/lib/gint/g2e_root3.cu b/gpu4pyscf/lib/gint/g2e_root3.cu index 351c99768..7543ebc6c 100644 --- a/gpu4pyscf/lib/gint/g2e_root3.cu +++ b/gpu4pyscf/lib/gint/g2e_root3.cu @@ -19,8 +19,7 @@ static void GINTfill_int2e_kernel0022(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -114,7 +113,7 @@ static void GINTfill_int2e_kernel0022(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -292,8 +291,7 @@ static void GINTfill_int2e_kernel0031(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -381,7 +379,7 @@ static void GINTfill_int2e_kernel0031(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -538,8 +536,7 @@ static void GINTfill_int2e_kernel0032(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -657,7 +654,7 @@ static void GINTfill_int2e_kernel0032(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -916,8 +913,7 @@ static void GINTfill_int2e_kernel1021(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1032,7 +1028,7 @@ static void GINTfill_int2e_kernel1021(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1277,8 +1273,7 @@ static void GINTfill_int2e_kernel1022(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1447,7 +1442,7 @@ static void GINTfill_int2e_kernel1022(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -1872,8 +1867,7 @@ static void GINTfill_int2e_kernel1030(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -1961,7 +1955,7 @@ static void GINTfill_int2e_kernel1030(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -2122,8 +2116,7 @@ static void GINTfill_int2e_kernel1031(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -2274,7 +2267,7 @@ static void GINTfill_int2e_kernel1031(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -2639,8 +2632,7 @@ static void GINTfill_int2e_kernel1111(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -2785,7 +2777,7 @@ static void GINTfill_int2e_kernel1111(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -3124,8 +3116,7 @@ static void GINTfill_int2e_kernel1120(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -3240,7 +3231,7 @@ static void GINTfill_int2e_kernel1120(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -3486,8 +3477,7 @@ static void GINTfill_int2e_kernel1121(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -3713,7 +3703,7 @@ static void GINTfill_int2e_kernel1121(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -4319,8 +4309,7 @@ static void GINTfill_int2e_kernel1130(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -4471,7 +4460,7 @@ static void GINTfill_int2e_kernel1130(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -4837,8 +4826,7 @@ static void GINTfill_int2e_kernel2011(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -4953,7 +4941,7 @@ static void GINTfill_int2e_kernel2011(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -5199,8 +5187,7 @@ static void GINTfill_int2e_kernel2020(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -5294,7 +5281,7 @@ static void GINTfill_int2e_kernel2020(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -5477,8 +5464,7 @@ static void GINTfill_int2e_kernel2021(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -5647,7 +5633,7 @@ static void GINTfill_int2e_kernel2021(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -6073,8 +6059,7 @@ static void GINTfill_int2e_kernel2030(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -6192,7 +6177,7 @@ static void GINTfill_int2e_kernel2030(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -6456,8 +6441,7 @@ static void GINTfill_int2e_kernel2110(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -6572,7 +6556,7 @@ static void GINTfill_int2e_kernel2110(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -6817,8 +6801,7 @@ static void GINTfill_int2e_kernel2111(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -7044,7 +7027,7 @@ static void GINTfill_int2e_kernel2111(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -7650,8 +7633,7 @@ static void GINTfill_int2e_kernel2120(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -7820,7 +7802,7 @@ static void GINTfill_int2e_kernel2120(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -8246,8 +8228,7 @@ static void GINTfill_int2e_kernel2200(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -8341,7 +8322,7 @@ static void GINTfill_int2e_kernel2200(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -8519,8 +8500,7 @@ static void GINTfill_int2e_kernel2210(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -8689,7 +8669,7 @@ static void GINTfill_int2e_kernel2210(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -9114,8 +9094,7 @@ static void GINTfill_int2e_kernel3010(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -9203,7 +9182,7 @@ static void GINTfill_int2e_kernel3010(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -9364,8 +9343,7 @@ static void GINTfill_int2e_kernel3011(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -9516,7 +9494,7 @@ static void GINTfill_int2e_kernel3011(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -9882,8 +9860,7 @@ static void GINTfill_int2e_kernel3020(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -10001,7 +9978,7 @@ static void GINTfill_int2e_kernel3020(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -10265,8 +10242,7 @@ static void GINTfill_int2e_kernel3100(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -10354,7 +10330,7 @@ static void GINTfill_int2e_kernel3100(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -10439,7 +10415,7 @@ static void GINTfill_int2e_kernel3100(GINTEnvVars envs, ERITensor eri, BasisProd int lstride = eri.stride_l; int *ao_loc = c_bpcache.ao_loc; int i0 = ao_loc[ish] - eri.ao_offsets_i; - int j0 = ao_loc[jsh] - eri.ao_offsets_j; + int j0 = ao_loc[jsh] - eri.ao_offsets_j; int k0 = ao_loc[ksh] - eri.ao_offsets_k; int l0 = ao_loc[lsh] - eri.ao_offsets_l; double* __restrict__ eri_ij = eri.data + l0*lstride+k0*kstride+j0*jstride+i0; @@ -10511,8 +10487,7 @@ static void GINTfill_int2e_kernel3110(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -10663,7 +10638,7 @@ static void GINTfill_int2e_kernel3110(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -11028,8 +11003,7 @@ static void GINTfill_int2e_kernel3200(GINTEnvVars envs, ERITensor eri, BasisProd { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -11147,7 +11121,7 @@ static void GINTfill_int2e_kernel3200(GINTEnvVars envs, ERITensor eri, BasisProd double aijkl = aij + akl; double a1 = aij * akl; double a0 = a1 / aijkl; - double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -11400,4 +11374,3 @@ static void GINTfill_int2e_kernel3200(GINTEnvVars envs, ERITensor eri, BasisProd eri_ij[9+5*jstride] = gout59; //eri_ji[9*jstride+5] = gout59; } - diff --git a/gpu4pyscf/lib/gint/g2e_root_n.cu b/gpu4pyscf/lib/gint/g2e_root_n.cu index 21c6c13ba..d6238ca9e 100644 --- a/gpu4pyscf/lib/gint/g2e_root_n.cu +++ b/gpu4pyscf/lib/gint/g2e_root_n.cu @@ -20,8 +20,14 @@ static void GINTfill_int2e_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffs { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_kl = item.get_global_id(0); + #else + const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -82,8 +88,7 @@ void GINTfill_int2e_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets off { int ntasks_ij = offsets.ntasks_ij; int ntasks_kl = offsets.ntasks_kl; - int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } diff --git a/gpu4pyscf/lib/gint/g3c1e.cu b/gpu4pyscf/lib/gint/g3c1e.cu index 75e1d3dc7..51925e15b 100644 --- a/gpu4pyscf/lib/gint/g3c1e.cu +++ b/gpu4pyscf/lib/gint/g3c1e.cu @@ -21,6 +21,9 @@ __device__ static void GINTwrite_int3c1e(const double* g, double* output, const int ish, const int jsh, const int i_grid, const int i_l, const int j_l, const int stride_j, const int stride_ij, const int ao_offsets_i, const int ao_offsets_j) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif const int* ao_loc = c_bpcache.ao_loc; const int i0 = ao_loc[ish ] - ao_offsets_i; @@ -68,8 +71,15 @@ static void GINTfill_int3c1e_kernel_general(double* output, const BasisProdOffse { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; @@ -141,7 +151,17 @@ static void GINTfill_int3c1e_charge_contracted_kernel_expanded(double* output, c const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -157,7 +177,7 @@ static void GINTfill_int3c1e_charge_contracted_kernel_expanded(double* output, c constexpr int n_density_elements_j = (LJ + 1) * (LJ + 2) / 2; double output_cache[n_density_elements_i * n_density_elements_j] { 0.0 }; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double charge = grid_point[3]; const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; @@ -224,7 +244,17 @@ static void GINTfill_int3c1e_charge_contracted_kernel_general(double* output, co { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -242,7 +272,7 @@ static void GINTfill_int3c1e_charge_contracted_kernel_general(double* output, co double output_cache[(l_i_max_density_elements + 1) * (l_i_max_density_elements + 2) / 2 * (l_j_max_density_elements + 1) * (l_j_max_density_elements + 2) / 2] { 0.0 }; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double charge = grid_point[3]; const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; @@ -277,7 +307,17 @@ static void GINTfill_int3c1e_density_contracted_kernel_general(double* output, c const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_grid = item.get_global_id(0); + const int thread_x_id = item.get_global_id(1); + const int total_threads_x = item.get_global_range(1); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + const int thread_x_id = blockIdx.x * blockDim.x + threadIdx.x; + const int total_threads_x = gridDim.x * blockDim.x; + #endif if (task_grid >= ngrids) { return; } @@ -289,7 +329,7 @@ static void GINTfill_int3c1e_density_contracted_kernel_general(double* output, c const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; double eri_with_density_pair_sum = 0.0; - for (int task_ij = blockIdx.x * blockDim.x + threadIdx.x; task_ij < ntasks_ij; task_ij += gridDim.x * blockDim.x) { + for (int task_ij = thread_x_id; task_ij < ntasks_ij; task_ij += total_threads_x) { const int bas_ij = offsets.bas_ij + task_ij; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; const int* bas_pair2bra = c_bpcache.bas_pair2bra; diff --git a/gpu4pyscf/lib/gint/g3c1e_ip.cu b/gpu4pyscf/lib/gint/g3c1e_ip.cu index 9a806befa..e6e050822 100644 --- a/gpu4pyscf/lib/gint/g3c1e_ip.cu +++ b/gpu4pyscf/lib/gint/g3c1e_ip.cu @@ -21,6 +21,9 @@ __device__ static void GINTwrite_int3c1e_ip(const double* g, double* output, const double minus_two_a, const double* u2, const double* AC, const int ish, const int jsh, const int i_grid, const int i_l, const int j_l, const int stride_j, const int stride_ij, const int ao_offsets_i, const int ao_offsets_j, const int ngrids) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif const int* ao_loc = c_bpcache.ao_loc; const int i0 = ao_loc[ish ] - ao_offsets_i; @@ -97,8 +100,15 @@ static void GINTfill_int3c1e_ip_kernel_general(double* output, const BasisProdOf { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; @@ -203,7 +213,17 @@ static void GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded(double* outpu const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -220,7 +240,7 @@ static void GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded(double* outpu constexpr int n_density_elements_j = (LJ + 1) * (LJ + 2) / 2; double output_cache[n_density_elements_i * n_density_elements_j * 3] { 0.0 }; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double charge = grid_point[3]; const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; @@ -313,7 +333,17 @@ static void GINTfill_int3c1e_ip1_charge_contracted_kernel_general(double* output { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -333,7 +363,7 @@ static void GINTfill_int3c1e_ip1_charge_contracted_kernel_general(double* output * (l_j_max_density_elements + 1) * (l_j_max_density_elements + 2) / 2 * 3] { 0.0 }; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double charge = grid_point[3]; const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; @@ -371,6 +401,9 @@ __device__ static void GINTwrite_int3c1e_ip1_density_contracted(const double* g, double* output, const double minus_two_a, const double* density, const int* aoslice, const int nao, const int ish, const int jsh, const int i_grid, const int i_l, const int j_l, const int ngrids) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif const int* ao_loc = c_bpcache.ao_loc; const int i0 = ao_loc[ish]; @@ -437,8 +470,15 @@ static void GINTfill_int3c1e_ip1_density_contracted_kernel_general(double* outpu { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; @@ -474,7 +514,17 @@ static void GINTfill_int3c1e_ip2_density_contracted_kernel_general(double* outpu const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_grid = item.get_global_id(0); + const int thread_x_id = item.get_global_id(1); + const int total_threads_x = item.get_global_range(1); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + const int thread_x_id = blockIdx.x * blockDim.x + threadIdx.x; + const int total_threads_x = gridDim.x * blockDim.x; + #endif if (task_grid >= ngrids) { return; } @@ -488,7 +538,7 @@ static void GINTfill_int3c1e_ip2_density_contracted_kernel_general(double* outpu double deri_dCx_pair_sum = 0.0; double deri_dCy_pair_sum = 0.0; double deri_dCz_pair_sum = 0.0; - for (int task_ij = blockIdx.x * blockDim.x + threadIdx.x; task_ij < ntasks_ij; task_ij += gridDim.x * blockDim.x) { + for (int task_ij = thread_x_id; task_ij < ntasks_ij; task_ij += total_threads_x) { const int bas_ij = offsets.bas_ij + task_ij; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; const int* bas_pair2bra = c_bpcache.bas_pair2bra; @@ -569,6 +619,9 @@ static void GINTwrite_int3c1e_ip2_charge_contracted(const double* g, double* out const int ish, const int jsh, const int i_grid, const int i_l, const int j_l, const int stride_j, const int stride_ij, const int ao_offsets_i, const int ao_offsets_j, const int* gridslice, const int ngrids) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif const int* ao_loc = c_bpcache.ao_loc; const int i0 = ao_loc[ish] - ao_offsets_i; @@ -639,8 +692,15 @@ static void GINTfill_int3c1e_ip2_charge_contracted_kernel_general(double* output { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int task_grid = item.get_global_id(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (task_ij >= ntasks_ij || task_grid >= ngrids) { return; diff --git a/gpu4pyscf/lib/gint/g3c1e_ipip.cu b/gpu4pyscf/lib/gint/g3c1e_ipip.cu index 87ebb2701..f6520211f 100644 --- a/gpu4pyscf/lib/gint/g3c1e_ipip.cu +++ b/gpu4pyscf/lib/gint/g3c1e_ipip.cu @@ -20,6 +20,9 @@ template __device__ static void GINTwrite_int3c1e_ipip1_charge_contracted(const double* g, double* local_output, const double minus_two_a, const double prefactor, const int i_l, const int j_l) { + #ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); + #endif const int *idx = c_idx; const int *idy = c_idx + TOT_NF; const int *idz = c_idx + TOT_NF * 2; @@ -100,7 +103,17 @@ static void GINTfill_int3c1e_ipip1_charge_contracted_kernel_general(double* outp { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -113,14 +126,21 @@ static void GINTfill_int3c1e_ipip1_charge_contracted_kernel_general(double* outp const int jsh = bas_pair2ket[bas_ij]; const double* __restrict__ a_exponents = c_bpcache.a1; - constexpr int l_sum_max = (NROOTS - 1) * 2 + 1; + // The launcher picks NROOTS = (i_l + j_l + 2) / 2 + 1, so i_l + j_l is at most + // 2 * NROOTS - 3 for this instantiation (it was over-estimated as 2 * NROOTS - 1, + // which made output_cache ~2x larger than it can ever need to be). That matters + // a lot on Intel GPUs: private arrays this large live in the scratch surface, whose + // size is per-thread-cost x all HW threads, so a per-work-item over-allocation is + // multiplied by ~10^5. Measured on PVC (1 tile, def2-tzvpp water): the int3c1e + // ipip kernels reserved 26.5 GB of scratch out of 64 GB. + constexpr int l_sum_max = 2 * NROOTS - 3; constexpr int l_i_max_density_elements = (l_sum_max + 1) / 2; constexpr int l_j_max_density_elements = l_sum_max - l_i_max_density_elements; double output_cache[(l_i_max_density_elements + 1) * (l_i_max_density_elements + 2) / 2 * (l_j_max_density_elements + 1) * (l_j_max_density_elements + 2) / 2 * 6] { 0.0 }; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double charge = grid_point[3]; const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; @@ -257,7 +277,17 @@ static void GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general(double* out { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -271,14 +301,21 @@ static void GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general(double* out const double* __restrict__ a_exponents = c_bpcache.a1; const double* __restrict__ b_exponents = c_bpcache.a2; - constexpr int l_sum_max = (NROOTS - 1) * 2 + 1; + // The launcher picks NROOTS = (i_l + j_l + 2) / 2 + 1, so i_l + j_l is at most + // 2 * NROOTS - 3 for this instantiation (it was over-estimated as 2 * NROOTS - 1, + // which made output_cache ~2x larger than it can ever need to be). That matters + // a lot on Intel GPUs: private arrays this large live in the scratch surface, whose + // size is per-thread-cost x all HW threads, so a per-work-item over-allocation is + // multiplied by ~10^5. Measured on PVC (1 tile, def2-tzvpp water): the int3c1e + // ipip kernels reserved 26.5 GB of scratch out of 64 GB. + constexpr int l_sum_max = 2 * NROOTS - 3; constexpr int l_i_max_density_elements = (l_sum_max + 1) / 2; constexpr int l_j_max_density_elements = l_sum_max - l_i_max_density_elements; double output_cache[(l_i_max_density_elements + 1) * (l_i_max_density_elements + 2) / 2 * (l_j_max_density_elements + 1) * (l_j_max_density_elements + 2) / 2 * 9] { 0.0 }; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double charge = grid_point[3]; const double charge_exponent = (charge_exponents != NULL) ? charge_exponents[task_grid] : 0.0; @@ -424,7 +461,17 @@ static void GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general(double* out { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_ij = item.get_global_id(1); + const int thread_y_id = item.get_global_id(0); + const int total_threads_y = item.get_global_range(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; + const int thread_y_id = blockIdx.y * blockDim.y + threadIdx.y; + const int total_threads_y = gridDim.y * blockDim.y; + #endif if (task_ij >= ntasks_ij) { return; } @@ -445,14 +492,21 @@ static void GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general(double* out const double Ay = bas_y[ish]; const double Az = bas_z[ish]; - constexpr int l_sum_max = (NROOTS - 1) * 2 + 1; + // The launcher picks NROOTS = (i_l + j_l + 2) / 2 + 1, so i_l + j_l is at most + // 2 * NROOTS - 3 for this instantiation (it was over-estimated as 2 * NROOTS - 1, + // which made output_cache ~2x larger than it can ever need to be). That matters + // a lot on Intel GPUs: private arrays this large live in the scratch surface, whose + // size is per-thread-cost x all HW threads, so a per-work-item over-allocation is + // multiplied by ~10^5. Measured on PVC (1 tile, def2-tzvpp water): the int3c1e + // ipip kernels reserved 26.5 GB of scratch out of 64 GB. + constexpr int l_sum_max = 2 * NROOTS - 3; constexpr int l_i_max_density_elements = (l_sum_max + 1) / 2; constexpr int l_j_max_density_elements = l_sum_max - l_i_max_density_elements; double output_cache[(l_i_max_density_elements + 1) * (l_i_max_density_elements + 2) / 2 * (l_j_max_density_elements + 1) * (l_j_max_density_elements + 2) / 2 * 9] { 0.0 }; - for (int task_grid = blockIdx.y * blockDim.y + threadIdx.y; task_grid < ngrids; task_grid += gridDim.y * blockDim.y) { + for (int task_grid = thread_y_id; task_grid < ngrids; task_grid += total_threads_y) { const double* grid_point = grid_points + task_grid * 4; const double Cx = grid_point[0]; const double Cy = grid_point[1]; @@ -513,7 +567,17 @@ static void GINTfill_int3c1e_ipip2_density_contracted_kernel_general(double* out const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int task_grid = item.get_global_id(0); + const int thread_x_id = item.get_global_id(1); + const int total_threads_x = item.get_global_range(1); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_grid = blockIdx.y * blockDim.y + threadIdx.y; + const int thread_x_id = blockIdx.x * blockDim.x + threadIdx.x; + const int total_threads_x = gridDim.x * blockDim.x; + #endif if (task_grid >= ngrids) { return; } @@ -530,7 +594,7 @@ static void GINTfill_int3c1e_ipip2_density_contracted_kernel_general(double* out double d2eri_dCydCy_pair_sum = 0.0; double d2eri_dCydCz_pair_sum = 0.0; double d2eri_dCzdCz_pair_sum = 0.0; - for (int task_ij = blockIdx.x * blockDim.x + threadIdx.x; task_ij < ntasks_ij; task_ij += gridDim.x * blockDim.x) { + for (int task_ij = thread_x_id; task_ij < ntasks_ij; task_ij += total_threads_x) { const int bas_ij = offsets.bas_ij + task_ij; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; diff --git a/gpu4pyscf/lib/gint/g3c2e.cu b/gpu4pyscf/lib/gint/g3c2e.cu index b59633dd1..75d95773f 100644 --- a/gpu4pyscf/lib/gint/g3c2e.cu +++ b/gpu4pyscf/lib/gint/g3c2e.cu @@ -16,9 +16,10 @@ __device__ -static void GINTwrite_int3c2e_direct(GINTEnvVars envs, ERITensor eri, double* g, +static void GINTwrite_int3c2e_direct(GINTEnvVars envs, ERITensor eri, const double* g, const int ish, const int jsh, const int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; size_t jstride = eri.stride_j; size_t kstride = eri.stride_k; @@ -29,9 +30,9 @@ static void GINTwrite_int3c2e_direct(GINTEnvVars envs, ERITensor eri, double* g, int k0 = ao_loc[ksh ] - eri.ao_offsets_k; int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - int *idx = c_idx; - int *idy = c_idx + TOT_NF; - int *idz = c_idx + TOT_NF * 2; + const int *idx = c_idx; + const int *idy = c_idx + TOT_NF; + const int *idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -46,7 +47,7 @@ static void GINTwrite_int3c2e_direct(GINTEnvVars envs, ERITensor eri, double* g, //for (int k = k0; k < k1; ++k) { //for (int j = j0; j < j1; ++j) { //for (int i = i0; i < i1; ++i) { - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -64,14 +65,16 @@ static void GINTwrite_int3c2e_direct(GINTEnvVars envs, ERITensor eri, double* g, for (int ir = 0; ir < nrys_roots; ++ir){ eri_data += g[ix + ir] * g[iy + ir] * g[iz + ir]; } - const int idx = (i+i0) + jstride*(j+j0) + kstride*(k+k0); - eri.data[idx] += eri_data; + const int out_idx = (i+i0) + jstride*(j+j0) + kstride*(k+k0); + //atomicAdd(&eri.data[out_idx], eri_data); + eri.data[out_idx] += eri_data; } } __device__ -static void GINTmemset_int3c2e(GINTEnvVars envs, ERITensor eri, int ish, int jsh, int ksh) +static void GINTmemset_int3c2e(ERITensor eri, int ish, int jsh, int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; size_t jstride = eri.stride_j; size_t kstride = eri.stride_k; @@ -82,7 +85,7 @@ static void GINTmemset_int3c2e(GINTEnvVars envs, ERITensor eri, int ish, int jsh int k0 = ao_loc[ksh ] - eri.ao_offsets_k; int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -93,12 +96,23 @@ static void GINTmemset_int3c2e(GINTEnvVars envs, ERITensor eri, int ish, int jsh } __global__ -void GINTfill_int3c2e_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets) +void GINTfill_int3c2e_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets + #ifdef USE_SYCL + , sycl::nd_item<2> item, double* g + #endif + ) { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; + #ifdef USE_SYCL + const int task_ij = item.get_group(1); + const int task_kl = item.get_group(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x;// * blockDim.x + threadIdx.x; const int task_kl = blockIdx.y;// * blockDim.y + threadIdx.y; + extern __shared__ double g[]; + #endif if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; @@ -114,12 +128,11 @@ void GINTfill_int3c2e_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets o const int ish = bas_pair2bra[bas_ij]; const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - extern __shared__ double g[]; const int as_ish = envs.ibase ? ish: jsh; const int as_jsh = envs.ibase ? jsh: ish; - GINTmemset_int3c2e(envs, eri, ish, jsh, ksh); + GINTmemset_int3c2e(eri, ish, jsh, ksh); for (int ij = prim_ij; ij < prim_ij+nprim_ij; ++ij) { for (int kl = prim_kl; kl < prim_kl+nprim_kl; ++kl) { GINTg0_int3c2e_shared(envs, g, as_ish, as_jsh, ksh, ij, kl); @@ -132,8 +145,7 @@ static void GINTfill_int3c2e_kernel0000(GINTEnvVars envs, ERITensor eri, BasisPr { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -156,7 +168,7 @@ static void GINTfill_int3c2e_kernel0000(GINTEnvVars envs, ERITensor eri, BasisPr double* __restrict__ x12 = c_bpcache.x12; double* __restrict__ y12 = c_bpcache.y12; double* __restrict__ z12 = c_bpcache.z12; - + double gout0 = 0; for (int ij = prim_ij; ij < prim_ij+nprim_ij; ++ij) { for (int kl = prim_kl; kl < prim_kl+nprim_kl; ++kl) { @@ -176,7 +188,7 @@ static void GINTfill_int3c2e_kernel0000(GINTEnvVars envs, ERITensor eri, BasisPr const double aijkl = aij + akl; const double a1 = aij * akl; double a0 = a1 / aijkl; - const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; const double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -187,7 +199,7 @@ static void GINTfill_int3c2e_kernel0000(GINTEnvVars envs, ERITensor eri, BasisPr } gout0 += fac; } } - + const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; @@ -203,8 +215,7 @@ static void GINTfill_int3c2e_kernel0010(GINTEnvVars envs, ERITensor eri, BasisPr { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -254,7 +265,7 @@ static void GINTfill_int3c2e_kernel0010(GINTEnvVars envs, ERITensor eri, BasisPr const double aijkl = aij + akl; const double a1 = aij * akl; double a0 = a1 / aijkl; - const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; const double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); const double fac = norm * eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -293,7 +304,7 @@ static void GINTfill_int3c2e_kernel0010(GINTEnvVars envs, ERITensor eri, BasisPr const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; - + int *ao_loc = c_bpcache.ao_loc; const int i0 = ao_loc[ish] - eri.ao_offsets_i; const int j0 = ao_loc[jsh] - eri.ao_offsets_j; @@ -310,8 +321,7 @@ static void GINTfill_int3c2e_kernel1000(GINTEnvVars envs, ERITensor eri, BasisPr { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -324,7 +334,7 @@ static void GINTfill_int3c2e_kernel1000(GINTEnvVars envs, ERITensor eri, BasisPr const int ish = bas_pair2bra[bas_ij]; const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - + const int nprim_ij = envs.nprim_ij; const int nprim_kl = envs.nprim_kl; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; @@ -368,7 +378,7 @@ static void GINTfill_int3c2e_kernel1000(GINTEnvVars envs, ERITensor eri, BasisPr const double aijkl = aij + akl; const double a1 = aij * akl; double a0 = a1 / aijkl; - const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; const double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); const double fac = eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -404,12 +414,12 @@ static void GINTfill_int3c2e_kernel1000(GINTEnvVars envs, ERITensor eri, BasisPr const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; - + int *ao_loc = c_bpcache.ao_loc; const int i0 = ao_loc[ish] - eri.ao_offsets_i; const int j0 = ao_loc[jsh] - eri.ao_offsets_j; const int k0 = ao_loc[ksh] - eri.ao_offsets_k; - + double* __restrict__ eri_ij = eri.data + k0*kstride+j0*jstride+i0; eri_ij[0] = gout0; @@ -422,8 +432,7 @@ static void GINTfill_int3c2e_kernel0100(GINTEnvVars envs, ERITensor eri, BasisPr { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -436,7 +445,7 @@ static void GINTfill_int3c2e_kernel0100(GINTEnvVars envs, ERITensor eri, BasisPr const int ish = bas_pair2bra[bas_ij]; const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - + const int nprim_ij = envs.nprim_ij; const int nprim_kl = envs.nprim_kl; const int prim_ij = offsets.primitive_ij + task_ij * nprim_ij; @@ -480,7 +489,7 @@ static void GINTfill_int3c2e_kernel0100(GINTEnvVars envs, ERITensor eri, BasisPr const double aijkl = aij + akl; const double a1 = aij * akl; double a0 = a1 / aijkl; - const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; + const double theta = omega > 0.0 ? omega * omega / (omega * omega + a0) : 1.0; a0 *= theta; const double x = a0 * (xijxkl * xijxkl + yijykl * yijykl + zijzkl * zijzkl); const double fac = eij * ekl * sqrt(a0 / (a1 * a1 * a1)); @@ -516,12 +525,12 @@ static void GINTfill_int3c2e_kernel0100(GINTEnvVars envs, ERITensor eri, BasisPr const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; - + int *ao_loc = c_bpcache.ao_loc; const int i0 = ao_loc[ish] - eri.ao_offsets_i; const int j0 = ao_loc[jsh] - eri.ao_offsets_j; const int k0 = ao_loc[ksh] - eri.ao_offsets_k; - + double* __restrict__ eri_ij = eri.data + k0*kstride+j0*jstride+i0; eri_ij[0] = gout0; diff --git a/gpu4pyscf/lib/gint/g3c2e_ip1.cu b/gpu4pyscf/lib/gint/g3c2e_ip1.cu index d1e476b4e..16e8261ea 100644 --- a/gpu4pyscf/lib/gint/g3c2e_ip1.cu +++ b/gpu4pyscf/lib/gint/g3c2e_ip1.cu @@ -17,9 +17,9 @@ template __device__ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* __restrict__ g, const double ai2) { - int *idx = c_idx; - int *idy = c_idx + TOT_NF; - int *idz = c_idx + TOT_NF * 2; + const int *idx = c_idx; + const int *idy = c_idx + TOT_NF; + const int *idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -44,7 +44,7 @@ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* int ix = dk * idx[loc_k] + dj * idx[loc_j] + di * i_idx; int iy = dk * idy[loc_k] + dj * idy[loc_j] + di * i_idy + g_size; int iz = dk * idz[loc_k] + dj * idz[loc_j] + di * i_idz + g_size * 2; - + #pragma unroll for (int n = 0; n < NROOTS; ++n, ++ix, ++iy, ++iz) { const double gx = g[ix]; @@ -73,8 +73,7 @@ void GINTfill_int3c2e_ip1_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffse { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; @@ -117,6 +116,7 @@ void GINTfill_int3c2e_ip1_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffse __device__ static void GINTwrite_int3c2e_ip1_direct(GINTEnvVars envs, ERITensor eri, double* g, double ai2, const int ish, const int jsh, const int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; @@ -128,8 +128,8 @@ static void GINTwrite_int3c2e_ip1_direct(GINTEnvVars envs, ERITensor eri, double const int k0 = ao_loc[ksh ] - eri.ao_offsets_k; const int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -141,7 +141,7 @@ static void GINTwrite_int3c2e_ip1_direct(GINTEnvVars envs, ERITensor eri, double const int lk = envs.k_l; const int nrys_roots = envs.nrys_roots; - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -149,7 +149,7 @@ static void GINTwrite_int3c2e_ip1_direct(GINTEnvVars envs, ERITensor eri, double const int loc_k = c_l_locs[lk] + k; const int loc_j = c_l_locs[lj] + j; const int loc_i = c_l_locs[li] + i; - + const int i_idx = c_idx[loc_i]; const int i_idy = c_idy[loc_i]; const int i_idz = c_idz[loc_i]; @@ -189,10 +189,21 @@ static void GINTwrite_int3c2e_ip1_direct(GINTEnvVars envs, ERITensor eri, double // General version __global__ -void GINTfill_int3c2e_ip1_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets) +void GINTfill_int3c2e_ip1_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double* g + #endif + ) { + #ifdef USE_SYCL + const int task_ij = item.get_group(1); + const int task_kl = item.get_group(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x;// * blockDim.x + threadIdx.x; const int task_kl = blockIdx.y;// * blockDim.y + threadIdx.y; + extern __shared__ double g[]; + #endif const int bas_ij = offsets.bas_ij + task_ij; const int bas_kl = offsets.bas_kl + task_kl; const int nprim_ij = envs.nprim_ij; @@ -205,8 +216,6 @@ void GINTfill_int3c2e_ip1_general_kernel(GINTEnvVars envs, ERITensor eri, BasisP const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - extern __shared__ double g[]; - const int as_ish = envs.ibase ? ish: jsh; const int as_jsh = envs.ibase ? jsh: ish; @@ -223,8 +232,7 @@ static void GINTfill_int3c2e_ip1_kernel000(GINTEnvVars envs, ERITensor eri, Basi { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } diff --git a/gpu4pyscf/lib/gint/g3c2e_ip1ip2.cu b/gpu4pyscf/lib/gint/g3c2e_ip1ip2.cu index 17ad9fb46..0bfe2169c 100644 --- a/gpu4pyscf/lib/gint/g3c2e_ip1ip2.cu +++ b/gpu4pyscf/lib/gint/g3c2e_ip1ip2.cu @@ -17,8 +17,8 @@ template __device__ static void GINTgout3c2e_ip1ip2(GINTEnvVars envs, double* __restrict__ gout, double *g0, double ai2, double ak2) { - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -105,8 +105,7 @@ void GINTfill_int3c2e_ip1ip2_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOf { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; @@ -153,6 +152,7 @@ static void GINTwrite_int3c2e_ip1ip2_direct(GINTEnvVars envs, ERITensor eri, double* __restrict__ g0, const double ai2, const double ak2, const int ish, const int jsh, const int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; @@ -164,8 +164,8 @@ static void GINTwrite_int3c2e_ip1ip2_direct(GINTEnvVars envs, ERITensor eri, const int k0 = ao_loc[ksh ] - eri.ao_offsets_k; const int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -177,7 +177,7 @@ static void GINTwrite_int3c2e_ip1ip2_direct(GINTEnvVars envs, ERITensor eri, const int lk = envs.k_l; const int nrys_roots = envs.nrys_roots; - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -272,10 +272,21 @@ static void GINTwrite_int3c2e_ip1ip2_direct(GINTEnvVars envs, ERITensor eri, // General version __global__ -void GINTfill_int3c2e_ip1ip2_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets) +void GINTfill_int3c2e_ip1ip2_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets +#ifdef USE_SYCL + , sycl::nd_item<2> item, double* g0 +#endif + ) { +#ifdef USE_SYCL + const int task_ij = item.get_group(1); + const int task_kl = item.get_group(0); + const auto& c_bpcache = s_bpcache.get(); +#else const int task_ij = blockIdx.x;// * blockDim.x + threadIdx.x; const int task_kl = blockIdx.y;// * blockDim.y + threadIdx.y; + extern __shared__ double g0[]; +#endif const int bas_ij = offsets.bas_ij + task_ij; const int bas_kl = offsets.bas_kl + task_kl; const int nprim_ij = envs.nprim_ij; @@ -288,8 +299,6 @@ void GINTfill_int3c2e_ip1ip2_general_kernel(GINTEnvVars envs, ERITensor eri, Bas const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - extern __shared__ double g0[]; - const int as_ish = envs.ibase ? ish: jsh; const int as_jsh = envs.ibase ? jsh: ish; @@ -308,8 +317,7 @@ static void GINTfill_int3c2e_ip1ip2_kernel000(GINTEnvVars envs, ERITensor eri, B { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } diff --git a/gpu4pyscf/lib/gint/g3c2e_ip2.cu b/gpu4pyscf/lib/gint/g3c2e_ip2.cu index fb1b85995..220a7f34c 100644 --- a/gpu4pyscf/lib/gint/g3c2e_ip2.cu +++ b/gpu4pyscf/lib/gint/g3c2e_ip2.cu @@ -17,9 +17,9 @@ template __device__ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* __restrict__ g, const double ak2) { - int *idx = c_idx; - int *idy = c_idx + TOT_NF; - int *idz = c_idx + TOT_NF * 2; + const int *idx = c_idx; + const int *idy = c_idx + TOT_NF; + const int *idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -36,7 +36,7 @@ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* const int loc_k = c_l_locs[LK] + ik; const int loc_j = c_l_locs[LJ] + ij; const int loc_i = c_l_locs[LI] + ii; - + const int k_idx = idx[loc_k]; const int k_idy = idy[loc_k]; const int k_idz = idz[loc_k]; @@ -72,8 +72,7 @@ void GINTfill_int3c2e_ip2_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffse { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; @@ -116,6 +115,7 @@ void GINTfill_int3c2e_ip2_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffse __device__ static void GINTwrite_int3c2e_ip2_direct(GINTEnvVars envs, ERITensor eri, double* g, double ak2, int ish, int jsh, int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; @@ -127,8 +127,8 @@ static void GINTwrite_int3c2e_ip2_direct(GINTEnvVars envs, ERITensor eri, double const int k0 = ao_loc[ksh ] - eri.ao_offsets_k; const int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -140,7 +140,7 @@ static void GINTwrite_int3c2e_ip2_direct(GINTEnvVars envs, ERITensor eri, double const int lk = envs.k_l; const int nrys_roots = envs.nrys_roots; - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -190,10 +190,21 @@ static void GINTwrite_int3c2e_ip2_direct(GINTEnvVars envs, ERITensor eri, double // General version __global__ -void GINTfill_int3c2e_ip2_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets) +void GINTfill_int3c2e_ip2_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets +#ifdef USE_SYCL + , sycl::nd_item<2> item, double* g +#endif + ) { + #ifdef USE_SYCL + const int task_ij = item.get_group(1); + const int task_kl = item.get_group(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x;// * blockDim.x + threadIdx.x; const int task_kl = blockIdx.y;// * blockDim.y + threadIdx.y; + extern __shared__ double g[]; + #endif const int bas_ij = offsets.bas_ij + task_ij; const int bas_kl = offsets.bas_kl + task_kl; const int nprim_ij = envs.nprim_ij; @@ -206,8 +217,6 @@ void GINTfill_int3c2e_ip2_general_kernel(GINTEnvVars envs, ERITensor eri, BasisP const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - extern __shared__ double g[]; - const int as_ish = envs.ibase ? ish: jsh; const int as_jsh = envs.ibase ? jsh: ish; @@ -224,8 +233,7 @@ static void GINTfill_int3c2e_ip2_kernel000(GINTEnvVars envs, ERITensor eri, Basi { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } @@ -331,4 +339,3 @@ static void GINTfill_int3c2e_ip2_kernel000(GINTEnvVars envs, ERITensor eri, Basi eri_ij[1*lstride] = gout1; eri_ij[2*lstride] = gout2; } - diff --git a/gpu4pyscf/lib/gint/g3c2e_ipip1.cu b/gpu4pyscf/lib/gint/g3c2e_ipip1.cu index 5b783c0f9..d97dcdb46 100644 --- a/gpu4pyscf/lib/gint/g3c2e_ipip1.cu +++ b/gpu4pyscf/lib/gint/g3c2e_ipip1.cu @@ -17,8 +17,8 @@ template __device__ static void GINTgout3c2e_ipip1(GINTEnvVars envs, double* __restrict__ gout, double *g0, double ai2) { - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -35,7 +35,7 @@ static void GINTgout3c2e_ipip1(GINTEnvVars envs, double* __restrict__ gout, doub const int loc_k = c_l_locs[LK] + ik; const int loc_j = c_l_locs[LJ] + ij; const int loc_i = c_l_locs[LI] + ii; - + const int i_idx = c_idx[loc_i]; const int i_idy = c_idy[loc_i]; const int i_idz = c_idz[loc_i]; @@ -93,8 +93,7 @@ void GINTfill_int3c2e_ipip1_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOff { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; @@ -138,6 +137,7 @@ static void GINTwrite_int3c2e_ipip1_direct(GINTEnvVars envs, ERITensor eri, double* __restrict__ g0, double ai2, const int ish, const int jsh, const int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; @@ -149,8 +149,8 @@ static void GINTwrite_int3c2e_ipip1_direct(GINTEnvVars envs, ERITensor eri, const int k0 = ao_loc[ksh ] - eri.ao_offsets_k; const int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -162,7 +162,7 @@ static void GINTwrite_int3c2e_ipip1_direct(GINTEnvVars envs, ERITensor eri, const int lk = envs.k_l; const int nrys_roots = envs.nrys_roots; - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -244,10 +244,21 @@ static void GINTwrite_int3c2e_ipip1_direct(GINTEnvVars envs, ERITensor eri, // General version __global__ -void GINTfill_int3c2e_ipip1_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets) +void GINTfill_int3c2e_ipip1_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets +#ifdef USE_SYCL + , sycl::nd_item<2> item, double* g0 +#endif + ) { + #ifdef USE_SYCL + const int task_ij = item.get_group(1); + const int task_kl = item.get_group(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x;// * blockDim.x + threadIdx.x; const int task_kl = blockIdx.y;// * blockDim.y + threadIdx.y; + extern __shared__ double g0[]; + #endif const int bas_ij = offsets.bas_ij + task_ij; const int bas_kl = offsets.bas_kl + task_kl; const int nprim_ij = envs.nprim_ij; @@ -260,8 +271,6 @@ void GINTfill_int3c2e_ipip1_general_kernel(GINTEnvVars envs, ERITensor eri, Basi const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - extern __shared__ double g0[]; - const int as_ish = envs.ibase ? ish: jsh; const int as_jsh = envs.ibase ? jsh: ish; @@ -278,8 +287,7 @@ static void GINTfill_int3c2e_ipip1_kernel000(GINTEnvVars envs, ERITensor eri, Ba { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } diff --git a/gpu4pyscf/lib/gint/g3c2e_ipip2.cu b/gpu4pyscf/lib/gint/g3c2e_ipip2.cu index d3b14d3fd..f319b4509 100644 --- a/gpu4pyscf/lib/gint/g3c2e_ipip2.cu +++ b/gpu4pyscf/lib/gint/g3c2e_ipip2.cu @@ -18,8 +18,8 @@ template __device__ static void GINTgout3c2e_ipip2(GINTEnvVars envs, double* __restrict__ gout, double *g0, double ak2) { - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -36,7 +36,7 @@ static void GINTgout3c2e_ipip2(GINTEnvVars envs, double* __restrict__ gout, doub const int loc_k = c_l_locs[LK] + ik; const int loc_j = c_l_locs[LJ] + ij; const int loc_i = c_l_locs[LI] + ii; - + const int k_idx = c_idx[loc_k]; const int k_idy = c_idy[loc_k]; const int k_idz = c_idz[loc_k]; @@ -96,8 +96,7 @@ void GINTfill_int3c2e_ipip2_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOff { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; @@ -142,6 +141,7 @@ static void GINTwrite_int3c2e_ipip2_direct(GINTEnvVars envs, ERITensor eri, double* __restrict__ g0, const double ak2, const int ish, const int jsh, const int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; @@ -153,8 +153,8 @@ static void GINTwrite_int3c2e_ipip2_direct(GINTEnvVars envs, ERITensor eri, const int k0 = ao_loc[ksh ] - eri.ao_offsets_k; const int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -166,7 +166,7 @@ static void GINTwrite_int3c2e_ipip2_direct(GINTEnvVars envs, ERITensor eri, const int lk = envs.k_l; const int nrys_roots = envs.nrys_roots; - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -248,10 +248,21 @@ static void GINTwrite_int3c2e_ipip2_direct(GINTEnvVars envs, ERITensor eri, // General version __global__ -void GINTfill_int3c2e_ipip2_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets) +void GINTfill_int3c2e_ipip2_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets +#ifdef USE_SYCL + , sycl::nd_item<2> item, double* g0 +#endif + ) { + #ifdef USE_SYCL + const int task_ij = item.get_group(1); + const int task_kl = item.get_group(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x;// * blockDim.x + threadIdx.x; const int task_kl = blockIdx.y;// * blockDim.y + threadIdx.y; + extern __shared__ double g0[]; + #endif const int bas_ij = offsets.bas_ij + task_ij; const int bas_kl = offsets.bas_kl + task_kl; const int nprim_ij = envs.nprim_ij; @@ -264,8 +275,6 @@ void GINTfill_int3c2e_ipip2_general_kernel(GINTEnvVars envs, ERITensor eri, Basi const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - extern __shared__ double g0[]; - const int as_ish = envs.ibase ? ish: jsh; const int as_jsh = envs.ibase ? jsh: ish; @@ -283,8 +292,7 @@ static void GINTfill_int3c2e_ipip2_kernel000(GINTEnvVars envs, ERITensor eri, Ba { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } diff --git a/gpu4pyscf/lib/gint/g3c2e_ipvip1.cu b/gpu4pyscf/lib/gint/g3c2e_ipvip1.cu index 2f2e51f29..663bdfdb6 100644 --- a/gpu4pyscf/lib/gint/g3c2e_ipvip1.cu +++ b/gpu4pyscf/lib/gint/g3c2e_ipvip1.cu @@ -17,8 +17,8 @@ template __device__ static void GINTgout3c2e_ipvip1(GINTEnvVars envs, double* __restrict__ gout, double *g0, double ai2, double aj2) { - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -35,7 +35,7 @@ static void GINTgout3c2e_ipvip1(GINTEnvVars envs, double* __restrict__ gout, dou const int loc_k = c_l_locs[LK] + ik; const int loc_j = c_l_locs[LJ] + ij; const int loc_i = c_l_locs[LI] + ii; - + const int i_idx = c_idx[loc_i]; const int i_idy = c_idy[loc_i]; const int i_idz = c_idz[loc_i]; @@ -105,8 +105,7 @@ void GINTfill_int3c2e_ipvip1_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOf { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; @@ -153,6 +152,7 @@ static void GINTwrite_int3c2e_ipvip1_direct(GINTEnvVars envs, ERITensor eri, double* __restrict__ g0, double ai2, double aj2, const int ish, const int jsh, const int ksh) { + KERNEL_SETUP_LOCAL(); int *ao_loc = c_bpcache.ao_loc; const size_t jstride = eri.stride_j; const size_t kstride = eri.stride_k; @@ -164,8 +164,8 @@ static void GINTwrite_int3c2e_ipvip1_direct(GINTEnvVars envs, ERITensor eri, const int k0 = ao_loc[ksh ] - eri.ao_offsets_k; const int k1 = ao_loc[ksh+1] - eri.ao_offsets_k; - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -177,7 +177,7 @@ static void GINTwrite_int3c2e_ipvip1_direct(GINTEnvVars envs, ERITensor eri, const int lk = envs.k_l; const int nrys_roots = envs.nrys_roots; - for (int tx = threadIdx.x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim.x) { + for (int tx = threadIdx_x; tx < (k1-k0)*(j1-j0)*(i1-i0); tx += blockDim_x) { const int k = tx / ((j1-j0)*(i1-i0)); const int j = (tx / (i1-i0)) % (j1-j0); const int i = tx % (i1-i0); @@ -271,10 +271,21 @@ static void GINTwrite_int3c2e_ipvip1_direct(GINTEnvVars envs, ERITensor eri, // General version __global__ -void GINTfill_int3c2e_ipvip1_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets) +void GINTfill_int3c2e_ipvip1_general_kernel(GINTEnvVars envs, ERITensor eri, BasisProdOffsets offsets +#ifdef USE_SYCL + , sycl::nd_item<2> item, double* g0 +#endif + ) { + #ifdef USE_SYCL + const int task_ij = item.get_group(1); + const int task_kl = item.get_group(0); + const auto& c_bpcache = s_bpcache.get(); + #else const int task_ij = blockIdx.x;// * blockDim.x + threadIdx.x; const int task_kl = blockIdx.y;// * blockDim.y + threadIdx.y; + extern __shared__ double g0[]; + #endif const int bas_ij = offsets.bas_ij + task_ij; const int bas_kl = offsets.bas_kl + task_kl; const int nprim_ij = envs.nprim_ij; @@ -287,8 +298,6 @@ void GINTfill_int3c2e_ipvip1_general_kernel(GINTEnvVars envs, ERITensor eri, Bas const int jsh = bas_pair2ket[bas_ij]; const int ksh = bas_pair2bra[bas_kl]; - extern __shared__ double g0[]; - const int as_ish = envs.ibase ? ish: jsh; const int as_jsh = envs.ibase ? jsh: ish; @@ -307,8 +316,7 @@ static void GINTfill_int3c2e_ipvip1_kernel000(GINTEnvVars envs, ERITensor eri, B { const int ntasks_ij = offsets.ntasks_ij; const int ntasks_kl = offsets.ntasks_kl; - const int task_ij = blockIdx.x * blockDim.x + threadIdx.x; - const int task_kl = blockIdx.y * blockDim.y + threadIdx.y; + KERNEL_SETUP(); if (task_ij >= ntasks_ij || task_kl >= ntasks_kl) { return; } diff --git a/gpu4pyscf/lib/gint/gint.h b/gpu4pyscf/lib/gint/gint.h index 60be35332..8bbe30888 100644 --- a/gpu4pyscf/lib/gint/gint.h +++ b/gpu4pyscf/lib/gint/gint.h @@ -17,6 +17,10 @@ #include #include +#ifdef USE_SYCL +#include +#endif + // #include // global parameters in env // Overall cutoff for integral prescreening, value needs to be ~ln(threshold) diff --git a/gpu4pyscf/lib/gint/gout3c2e.cu b/gpu4pyscf/lib/gint/gout3c2e.cu index 7036852ba..840093062 100644 --- a/gpu4pyscf/lib/gint/gout3c2e.cu +++ b/gpu4pyscf/lib/gint/gout3c2e.cu @@ -20,15 +20,16 @@ #include #include #include + #include "g2e.h" #include "cint2e.cuh" template __device__ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* __restrict__ f, double* __restrict__ g) { - int *idx = c_idx; - int *idy = c_idx + TOT_NF; - int *idz = c_idx + TOT_NF * 2; + const int *idx = c_idx; + const int *idy = c_idx + TOT_NF; + const int *idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -73,9 +74,9 @@ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* template __device__ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* __restrict__ f, double* __restrict__ g) { - int *idx = c_idx; - int *idy = c_idx + TOT_NF; - int *idz = c_idx + TOT_NF * 2; + const int *idx = c_idx; + const int *idy = c_idx + TOT_NF; + const int *idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -116,8 +117,8 @@ static void GINTgout3c2e_ip(GINTEnvVars envs, double* __restrict__ gout, double* template __device__ static void GINTgout3c2e(GINTEnvVars envs, double* __restrict__ gout, double* __restrict__ g) { - int * __restrict__ c_idy = c_idx + TOT_NF; - int * __restrict__ c_idz = c_idx + TOT_NF * 2; + const int * __restrict__ c_idy = c_idx + TOT_NF; + const int * __restrict__ c_idz = c_idx + TOT_NF * 2; const int di = envs.stride_i; const int dj = envs.stride_j; @@ -157,6 +158,9 @@ __device__ static void GINTwrite_int3c2e(ERITensor eri, double* __restrict__ gout, int ish, int jsh, int ksh) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif int *ao_loc = c_bpcache.ao_loc; size_t jstride = eri.stride_j; size_t kstride = eri.stride_k; @@ -186,6 +190,9 @@ __device__ static void GINTwrite_int3c2e_ip(ERITensor eri, double* __restrict__ gout, int ish, int jsh, int ksh) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif int *ao_loc = c_bpcache.ao_loc; size_t jstride = eri.stride_j; size_t kstride = eri.stride_k; @@ -225,6 +232,9 @@ __device__ static void GINTwrite_int3c2e_ipip(ERITensor eri, double* __restrict__ gout, int ish, int jsh, int ksh) { +#ifdef USE_SYCL + const auto& c_bpcache = s_bpcache.get(); +#endif int *ao_loc = c_bpcache.ao_loc; size_t jstride = eri.stride_j; size_t kstride = eri.stride_k; diff --git a/gpu4pyscf/lib/gint/j_engine_matrix_reorder.c b/gpu4pyscf/lib/gint/j_engine_matrix_reorder.c index 0185bbed4..ca1975c14 100644 --- a/gpu4pyscf/lib/gint/j_engine_matrix_reorder.c +++ b/gpu4pyscf/lib/gint/j_engine_matrix_reorder.c @@ -102,6 +102,9 @@ int hermite_xyz_to_t_index(const int x, const int y, const int z, const int l) return ((l + 1) * (l + 2) * (l + 3) - (l + 1 - x) * (l + 2 - x) * (l + 3 - x)) / 6 + ((l - x + 1) * (l - x + 2) - (l - x + 1 - y) * (l - x + 2 - y)) / 2 + z; } +#ifdef __cplusplus +extern "C" +#endif void GINTinit_J_density_rys_preprocess(const double* D_matrix, double* D_pair_ordered, const int n_dm, const int n_ao, const int n_pair_type, const int* bas_pair2shls, const int* bas_pairs_locs, const int* l_ij, const int* density_offset, const int* ao_loc, const double* bas_coords, const bool symmetric) diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e.cu index 70da7ea4b..ac064ba79 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "gint1e.h" #include "cuda_alloc.cuh" @@ -30,6 +31,37 @@ #include "g1e_root_1.cu" #include "g3c1e.cu" +// Abstracts 2D thread/block config (THREADSX/Y swapped between SYCL and CUDA). +// Used 3x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ngrids+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + const dim3 threads(THREADSX, THREADSY); \ + const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function with template args if needed +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// ARGS: parenthesized runtime args. Kernel-id (with any template args) is the +// trailing __VA_ARGS__ so its commas survive macro expansion. SYCL kernel name +// is generated inline per source line (unique within this translation unit). +#define LAUNCH_KERNEL(ARGS, ...) \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { __VA_ARGS__ ARGS; }); +#else +#define LAUNCH_KERNEL(ARGS, ...) \ + __VA_ARGS__ <<>> ARGS; +#endif + static int GINTfill_int3c1e_tasks(double* output, const BasisProdOffsets offsets, const int i_l, const int j_l, const int nprim_ij, const int stride_j, const int stride_ij, const int ao_offsets_i, const int ao_offsets_j, const double omega, const double* grid_points, const double* charge_exponents, const cudaStream_t stream) @@ -38,28 +70,26 @@ static int GINTfill_int3c1e_tasks(double* output, const BasisProdOffsets offsets const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); int type_ijkl; switch (nrys_roots) { case 1: type_ijkl = (i_l << 2) | j_l; switch (type_ijkl) { - case (0<<2)|0: GINTfill_int3c1e_kernel00<<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case (1<<2)|0: GINTfill_int3c1e_kernel10<<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case (0<<2)|0: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_kernel00) break; + case (1<<2)|0: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_kernel10) break; default: fprintf(stderr, "roots=1 type_ijkl %d\n", type_ijkl); } break; - case 2: GINTfill_int3c1e_kernel_general<2, GSIZE2_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_kernel_general<3, GSIZE3_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 2: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_kernel_general<2, GSIZE2_INT3C_1E>) break; + case 3: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_kernel_general<3, GSIZE3_INT3C_1E>) break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_kernel_general<5, GSIZE5_INT3C_1E>) break; default: fprintf(stderr, "rys roots %d\n", nrys_roots); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -75,33 +105,31 @@ static int GINTfill_int3c1e_charge_contracted_tasks(double* output, const BasisP { const int ntasks_ij = offsets.ntasks_ij; const int ngrids = (offsets.ntasks_kl + n_charge_sum_per_thread - 1) / n_charge_sum_per_thread; - - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); const int type_ij = i_l * 10 + j_l; + + LAUNCH_CONFIG(); switch (type_ij) { - case 00: GINTfill_int3c1e_charge_contracted_kernel00<<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 10: GINTfill_int3c1e_charge_contracted_kernel10<<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 11: GINTfill_int3c1e_charge_contracted_kernel_expanded<1, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 20: GINTfill_int3c1e_charge_contracted_kernel_expanded<2, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 21: GINTfill_int3c1e_charge_contracted_kernel_expanded<2, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 22: GINTfill_int3c1e_charge_contracted_kernel_expanded<2, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 30: GINTfill_int3c1e_charge_contracted_kernel_expanded<3, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 31: GINTfill_int3c1e_charge_contracted_kernel_expanded<3, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 32: GINTfill_int3c1e_charge_contracted_kernel_expanded<3, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 40: GINTfill_int3c1e_charge_contracted_kernel_expanded<4, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 41: GINTfill_int3c1e_charge_contracted_kernel_expanded<4, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 00: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel00) break; + case 10: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel10) break; + case 11: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<1, 1>) break; + case 20: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<2, 0>) break; + case 21: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<2, 1>) break; + case 22: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<2, 2>) break; + case 30: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<3, 0>) break; + case 31: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<3, 1>) break; + case 32: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<3, 2>) break; + case 40: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<4, 0>) break; + case 41: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_expanded<4, 1>) break; default: const int nrys_roots = (i_l + j_l) / 2 + 1; switch (nrys_roots) { - case 4: GINTfill_int3c1e_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E>) break; default: fprintf(stderr, "type_ij = %d, nrys_roots = %d out of range\n", type_ij, nrys_roots); return 1; } } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -118,24 +146,22 @@ static int GINTfill_int3c1e_density_contracted_tasks(double* output, const doubl const int ntasks_ij = (offsets.ntasks_ij + n_pair_sum_per_thread - 1) / n_pair_sum_per_thread; const int ngrids = offsets.ntasks_kl; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); switch (i_l + j_l) { - case 0: GINTfill_int3c1e_density_contracted_kernel00<<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 1: GINTfill_int3c1e_density_contracted_kernel10<<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 2: GINTfill_int3c1e_density_contracted_kernel_general< 2> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_density_contracted_kernel_general< 3> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_density_contracted_kernel_general< 4> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_density_contracted_kernel_general< 5> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 6: GINTfill_int3c1e_density_contracted_kernel_general< 6> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 7: GINTfill_int3c1e_density_contracted_kernel_general< 7> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 8: GINTfill_int3c1e_density_contracted_kernel_general< 8> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; + case 0: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel00) break; + case 1: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel10) break; + case 2: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel_general< 2>) break; + case 3: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel_general< 3>) break; + case 4: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel_general< 4>) break; + case 5: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel_general< 5>) break; + case 6: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel_general< 6>) break; + case 7: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel_general< 7>) break; + case 8: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_density_contracted_kernel_general< 8>) break; // Up to g + g = 8 now default: fprintf(stderr, "i_l + j_l = %d out of range\n", i_l + j_l); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -163,7 +189,11 @@ int GINTfill_int3c1e(const cudaStream_t stream, const BasisProdCache* bpcache, return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -213,7 +243,11 @@ int GINTfill_int3c1e_charge_contracted(const cudaStream_t stream, const BasisPro return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -263,7 +297,11 @@ int GINTfill_int3c1e_density_contracted(const cudaStream_t stream, const BasisPr return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -300,3 +338,6 @@ int GINTfill_int3c1e_density_contracted(const cudaStream_t stream, const BasisPr return 0; } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ip.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ip.cu index 3ee7c4236..0c0432ca4 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ip.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ip.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "gint1e.h" #include "cuda_alloc.cuh" @@ -30,6 +31,37 @@ #include "g1e_ip_root_1.cu" #include "g3c1e_ip.cu" +// Abstracts 2D thread/block config (THREADSX/Y are swapped between SYCL and CUDA). +// Used 3x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ngrids+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + const dim3 threads(THREADSX, THREADSY); \ + const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. blocks/threads must be in scope. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function with template args if needed +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// ARGS: parenthesized runtime args. Kernel-id (with any template args) is the +// trailing __VA_ARGS__ so its commas survive macro expansion. SYCL kernel name +// is generated inline per source line (unique within this translation unit). +#define LAUNCH_KERNEL(ARGS, ...) \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { __VA_ARGS__ ARGS; }); +#else +#define LAUNCH_KERNEL(ARGS, ...) \ + __VA_ARGS__ <<>> ARGS; +#endif + static int GINTfill_int3c1e_ip_tasks(double* output, const BasisProdOffsets offsets, const int i_l, const int j_l, const int nprim_ij, const int stride_j, const int stride_ij, const int ao_offsets_i, const int ao_offsets_j, const double omega, const double* grid_points, const double* charge_exponents, const cudaStream_t stream) @@ -38,27 +70,25 @@ static int GINTfill_int3c1e_ip_tasks(double* output, const BasisProdOffsets offs const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); int type_ijkl; switch (nrys_roots) { case 1: type_ijkl = (i_l + 1) * 10 + j_l; switch (type_ijkl) { - case 10: GINTfill_int3c1e_ip_kernel00<<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 10: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip_kernel00) break; default: fprintf(stderr, "roots=1 type_ijkl %d\n", type_ijkl); } break; - case 2: GINTfill_int3c1e_ip_kernel_general<2, GSIZE2_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ip_kernel_general<3, GSIZE3_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ip_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ip_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 2: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip_kernel_general<2, GSIZE2_INT3C_1E>) break; + case 3: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip_kernel_general<3, GSIZE3_INT3C_1E>) break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip_kernel_general<5, GSIZE5_INT3C_1E>) break; default: fprintf(stderr, "rys roots %d\n", nrys_roots); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -75,36 +105,34 @@ static int GINTfill_int3c1e_ip1_charge_contracted_tasks(double* output, const Ba const int ntasks_ij = offsets.ntasks_ij; const int ngrids = (offsets.ntasks_kl + n_charge_sum_per_thread - 1) / n_charge_sum_per_thread; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); const int type_ij = i_l * 10 + j_l; switch (type_ij) { - case 00: GINTfill_int3c1e_ip1_charge_contracted_kernel00<<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 01: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 02: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 03: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 3> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 04: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 4> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 10: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 11: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 12: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 13: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 3> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 20: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<2, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 21: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<2, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 22: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<2, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 30: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<3, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 31: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<3, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 40: GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<4, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 00: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel00) break; + case 01: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 1>) break; + case 02: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 2>) break; + case 03: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 3>) break; + case 04: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<0, 4>) break; + case 10: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 0>) break; + case 11: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 1>) break; + case 12: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 2>) break; + case 13: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<1, 3>) break; + case 20: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<2, 0>) break; + case 21: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<2, 1>) break; + case 22: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<2, 2>) break; + case 30: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<3, 0>) break; + case 31: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<3, 1>) break; + case 40: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_expanded<4, 0>) break; default: const int nrys_roots = (i_l + j_l + 1) / 2 + 1; switch (nrys_roots) { - case 4: GINTfill_int3c1e_ip1_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ip1_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E>) break; default: fprintf(stderr, "type_ij = %d, nrys_roots = %d out of range\n", type_ij, nrys_roots); return 1; } } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -121,39 +149,37 @@ static int GINTfill_int3c1e_ip1_density_contracted_tasks(double* output, const B const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); const int type_ij = i_l * 10 + j_l; switch (type_ij) { - case 00: GINTfill_int3c1e_ip1_density_contracted_kernel00<<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 01: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 1> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 02: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 2> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 03: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 3> <<>>(output, offsets, nprim_ij, density, shell, nao, omega, grid_points, charge_exponents); break; - // case 04: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 4> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 10: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 0> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 11: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 1> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 12: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 2> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 13: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 3> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 20: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<2, 0> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 21: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<2, 1> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 22: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<2, 2> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 30: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<3, 0> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 31: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<3, 1> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - // case 40: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<4, 0> <<>>(output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; + case 00: LAUNCH_KERNEL((output, offsets, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_density_contracted_kernel00) break; + // case 01: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 1> + // case 02: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 2> + // case 03: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 3> + // case 04: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<0, 4> + // case 10: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 0> + // case 11: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 1> + // case 12: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 2> + // case 13: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<1, 3> + // case 20: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<2, 0> + // case 21: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<2, 1> + // case 22: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<2, 2> + // case 30: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<3, 0> + // case 31: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<3, 1> + // case 40: GINTfill_int3c1e_ip1_density_contracted_kernel_expanded<4, 0> default: const int nrys_roots = (i_l + j_l + 1) / 2 + 1; switch (nrys_roots) { - case 1: GINTfill_int3c1e_ip1_density_contracted_kernel_general<1, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - case 2: GINTfill_int3c1e_ip1_density_contracted_kernel_general<2, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ip1_density_contracted_kernel_general<3, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ip1_density_contracted_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ip1_density_contracted_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents); break; + case 1: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_density_contracted_kernel_general<1, GSIZE5_INT3C_1E>) break; + case 2: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_density_contracted_kernel_general<2, GSIZE4_INT3C_1E>) break; + case 3: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_density_contracted_kernel_general<3, GSIZE5_INT3C_1E>) break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_density_contracted_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, density, aoslice, nao, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1_density_contracted_kernel_general<5, GSIZE5_INT3C_1E>) break; default: fprintf(stderr, "type_ij = %d, nrys_roots = %d out of range\n", type_ij, nrys_roots); return 1; } } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -170,24 +196,22 @@ static int GINTfill_int3c1e_ip2_density_contracted_tasks(double* output, const d const int ntasks_ij = (offsets.ntasks_ij + n_pair_sum_per_thread - 1) / n_pair_sum_per_thread; const int ngrids = offsets.ntasks_kl; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); switch (i_l + j_l) { - case 0: GINTfill_int3c1e_ip2_density_contracted_kernel00<<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 1: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 1> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 2: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 2> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 3> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 4> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 5> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 6: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 6> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 7: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 7> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 8: GINTfill_int3c1e_ip2_density_contracted_kernel_general< 8> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; + case 0: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel00) break; + case 1: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 1>) break; + case 2: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 2>) break; + case 3: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 3>) break; + case 4: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 4>) break; + case 5: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 5>) break; + case 6: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 6>) break; + case 7: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 7>) break; + case 8: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_density_contracted_kernel_general< 8>) break; // Up to g + g = 8 now default: fprintf(stderr, "i_l + j_l = %d out of range\n", i_l + j_l); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -205,39 +229,37 @@ static int GINTfill_int3c1e_ip2_charge_contracted_tasks(double* output, const Ba const int ntasks_ij = offsets.ntasks_ij; const int ngrids = offsets.ntasks_kl; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); const int type_ij = i_l * 10 + j_l; switch (type_ij) { - case 00: GINTfill_int3c1e_ip2_charge_contracted_kernel00<<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 01: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 02: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 03: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 3> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 04: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 4> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 10: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 11: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 12: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 13: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 3> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 20: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<2, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 21: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<2, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 22: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<2, 2> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 30: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<3, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 31: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<3, 1> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - // case 40: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<4, 0> <<>>(output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; + case 00: LAUNCH_KERNEL((output, offsets, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_charge_contracted_kernel00) break; + // case 01: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 1> + // case 02: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 2> + // case 03: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 3> + // case 04: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<0, 4> + // case 10: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 0> + // case 11: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 1> + // case 12: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 2> + // case 13: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<1, 3> + // case 20: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<2, 0> + // case 21: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<2, 1> + // case 22: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<2, 2> + // case 30: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<3, 0> + // case 31: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<3, 1> + // case 40: GINTfill_int3c1e_ip2_charge_contracted_kernel_expanded<4, 0> default: const int nrys_roots = (i_l + j_l + 1) / 2 + 1; switch (nrys_roots) { - case 1: GINTfill_int3c1e_ip2_charge_contracted_kernel_general<1, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - case 2: GINTfill_int3c1e_ip2_charge_contracted_kernel_general<2, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ip2_charge_contracted_kernel_general<3, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ip2_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ip2_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents); break; + case 1: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_charge_contracted_kernel_general<1, GSIZE5_INT3C_1E>) break; + case 2: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_charge_contracted_kernel_general<2, GSIZE4_INT3C_1E>) break; + case 3: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_charge_contracted_kernel_general<3, GSIZE5_INT3C_1E>) break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, gridslice, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip2_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E>) break; default: fprintf(stderr, "type_ij = %d, nrys_roots = %d out of range\n", type_ij, nrys_roots); return 1; } } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -265,7 +287,11 @@ int GINTfill_int3c1e_ip(const cudaStream_t stream, const BasisProdCache* bpcache return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else // USE_SYCL checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif // USE_SYCL const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -316,7 +342,11 @@ int GINTfill_int3c1e_ip1_density_contracted(const cudaStream_t stream, const Bas return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else // USE_SYCL checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -367,7 +397,11 @@ int GINTfill_int3c1e_ip1_charge_contracted(const cudaStream_t stream, const Basi return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else // USE_SYCL checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -417,7 +451,11 @@ int GINTfill_int3c1e_ip2_density_contracted(const cudaStream_t stream, const Bas return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else // USE_SYCL checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -474,7 +512,11 @@ int GINTfill_int3c1e_ip2_charge_contracted(const cudaStream_t stream, const Basi return 2; } +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else // USE_SYCL checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -506,3 +548,6 @@ int GINTfill_int3c1e_ip2_charge_contracted(const cudaStream_t stream, const Basi return 0; } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ipip.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ipip.cu index 4f3a3dee5..348006a4e 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ipip.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c1e_ipip.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "gint1e.h" #include "cuda_alloc.cuh" @@ -29,6 +30,37 @@ #include "g1e.cu" #include "g3c1e_ipip.cu" +// Abstracts 2D thread/block config (THREADSX/Y are swapped between SYCL and CUDA). +// Used 4x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ngrids+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + const dim3 threads(THREADSX, THREADSY); \ + const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. blocks/threads must be in scope. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function with template args if needed (e.g. func) +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// ARGS: parenthesized runtime args. Kernel-id (with any template args) is the +// trailing __VA_ARGS__ so its commas survive macro expansion. SYCL kernel name +// is generated inline per source line (unique within this translation unit). +#define LAUNCH_KERNEL(ARGS, ...) \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { __VA_ARGS__ ARGS; }); +#else +#define LAUNCH_KERNEL(ARGS, ...) \ + __VA_ARGS__ <<>> ARGS; +#endif + static int GINTfill_int3c1e_ipip1_charge_contracted_tasks(double* output, const BasisProdOffsets offsets, const int i_l, const int j_l, const int nprim_ij, const int stride_j, const int stride_ij, const int ao_offsets_i, const int ao_offsets_j, const double omega, const double* grid_points, const double* charge_exponents, @@ -37,20 +69,23 @@ static int GINTfill_int3c1e_ipip1_charge_contracted_tasks(double* output, const const int ntasks_ij = offsets.ntasks_ij; const int ngrids = (offsets.ntasks_kl + n_charge_sum_per_thread - 1) / n_charge_sum_per_thread; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); const int nrys_roots = (i_l + j_l + 2) / 2 + 1; + // GSIZE must cover 3 * NROOTS * (i_l + 2 + 1) * (j_l + 1) doubles, whose maximum + // over the (i_l, j_l) reaching this case is exactly GSIZE. Cases 2 and 3 + // used to ask for GSIZE5 / GSIZE6 (12x / 7x too much); on Intel GPUs that private + // array lands in the scratch surface, which is sized per HW thread for the whole + // device, so the waste costs gigabytes of device memory. switch (nrys_roots) { - case 2: GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<2, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<3, GSIZE6_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 6: GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<6, GSIZE6_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 2: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<2, GSIZE2_INT3C_1E>) break; + case 3: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<3, GSIZE3_INT3C_1E>) break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E>) break; + case 6: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip1_charge_contracted_kernel_general<6, GSIZE6_INT3C_1E>) break; default: fprintf(stderr, "nrys_roots = %d out of range\n", nrys_roots); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -67,20 +102,23 @@ static int GINTfill_int3c1e_ipvip1_charge_contracted_tasks(double* output, const const int ntasks_ij = offsets.ntasks_ij; const int ngrids = (offsets.ntasks_kl + n_charge_sum_per_thread - 1) / n_charge_sum_per_thread; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); const int nrys_roots = (i_l + j_l + 2) / 2 + 1; + // GSIZE must cover 3 * NROOTS * (i_l + 2 + 1) * (j_l + 1) doubles, whose maximum + // over the (i_l, j_l) reaching this case is exactly GSIZE. Cases 2 and 3 + // used to ask for GSIZE5 / GSIZE6 (12x / 7x too much); on Intel GPUs that private + // array lands in the scratch surface, which is sized per HW thread for the whole + // device, so the waste costs gigabytes of device memory. switch (nrys_roots) { - case 2: GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<2, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<3, GSIZE6_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 6: GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<6, GSIZE6_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 2: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<2, GSIZE2_INT3C_1E>) break; + case 3: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<3, GSIZE3_INT3C_1E>) break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E>) break; + case 6: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipvip1_charge_contracted_kernel_general<6, GSIZE6_INT3C_1E>) break; default: fprintf(stderr, "nrys_roots = %d out of range\n", nrys_roots); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -97,20 +135,23 @@ static int GINTfill_int3c1e_ip1ip2_charge_contracted_tasks(double* output, const const int ntasks_ij = offsets.ntasks_ij; const int ngrids = (offsets.ntasks_kl + n_charge_sum_per_thread - 1) / n_charge_sum_per_thread; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); const int nrys_roots = (i_l + j_l + 2) / 2 + 1; + // GSIZE must cover 3 * NROOTS * (i_l + 2 + 1) * (j_l + 1) doubles, whose maximum + // over the (i_l, j_l) reaching this case is exactly GSIZE. Cases 2 and 3 + // used to ask for GSIZE5 / GSIZE6 (12x / 7x too much); on Intel GPUs that private + // array lands in the scratch surface, which is sized per HW thread for the whole + // device, so the waste costs gigabytes of device memory. switch (nrys_roots) { - case 2: GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<2, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<3, GSIZE6_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; - case 6: GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<6, GSIZE6_INT3C_1E> <<>>(output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents); break; + case 2: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<2, GSIZE2_INT3C_1E>) break; + case 3: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<3, GSIZE3_INT3C_1E>) break; + case 4: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<4, GSIZE4_INT3C_1E>) break; + case 5: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<5, GSIZE5_INT3C_1E>) break; + case 6: LAUNCH_KERNEL((output, offsets, i_l, j_l, nprim_ij, stride_j, stride_ij, ao_offsets_i, ao_offsets_j, omega, grid_points, charge_exponents), GINTfill_int3c1e_ip1ip2_charge_contracted_kernel_general<6, GSIZE6_INT3C_1E>) break; default: fprintf(stderr, "nrys_roots = %d out of range\n", nrys_roots); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -127,24 +168,22 @@ static int GINTfill_int3c1e_ipip2_density_contracted_tasks(double* output, const const int ntasks_ij = (offsets.ntasks_ij + n_pair_sum_per_thread - 1) / n_pair_sum_per_thread; const int ngrids = offsets.ntasks_kl; - const dim3 threads(THREADSX, THREADSY); - const dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ngrids+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); switch (i_l + j_l) { - case 0: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 0> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 1: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 1> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 2: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 2> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 3: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 3> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 4: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 4> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 5: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 5> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 6: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 6> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 7: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 7> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; - case 8: GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 8> <<>>(output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents); break; + case 0: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 0>) break; + case 1: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 1>) break; + case 2: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 2>) break; + case 3: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 3>) break; + case 4: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 4>) break; + case 5: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 5>) break; + case 6: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 6>) break; + case 7: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 7>) break; + case 8: LAUNCH_KERNEL((output, density, hermite_density_offsets, offsets, nprim_ij, omega, grid_points, charge_exponents), GINTfill_int3c1e_ipip2_density_contracted_kernel_general< 8>) break; // Up to g + g = 8 now default: fprintf(stderr, "i_l + j_l = %d out of range\n", i_l + j_l); return 1; } - cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in %s: %s\n", __func__, cudaGetErrorString(err)); @@ -172,7 +211,11 @@ int GINTfill_int3c1e_ipip1_charge_contracted(const cudaStream_t stream, const Ba return 2; } + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -222,7 +265,11 @@ int GINTfill_int3c1e_ipvip1_charge_contracted(const cudaStream_t stream, const B return 2; } + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -272,7 +319,11 @@ int GINTfill_int3c1e_ip1ip2_charge_contracted(const cudaStream_t stream, const B return 2; } + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -322,7 +373,11 @@ int GINTfill_int3c1e_ipip2_density_contracted(const cudaStream_t stream, const B return 2; } + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif const int* bas_pairs_locs = bpcache->bas_pairs_locs; const int* primitive_pairs_locs = bpcache->primitive_pairs_locs; @@ -359,3 +414,6 @@ int GINTfill_int3c1e_ipip2_density_contracted(const cudaStream_t stream, const B return 0; } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e.cu index 8cf8eadaa..621a116ad 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "config.h" #include "cuda_alloc.cuh" @@ -33,6 +34,24 @@ #include "g2e_root3.cu" #include "g3c2e.cu" +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel name generated per source line for unique identification. +// dev_envs/dev_eri/dev_offsets are on-host value copies for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } + +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + static int GINTfill_int3c2e_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { int nrys_roots = envs->nrys_roots; @@ -40,58 +59,75 @@ static int GINTfill_int3c2e_tasks(ERITensor *eri, BasisProdOffsets *offsets, GIN int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); + #ifdef USE_SYCL + sycl::range<2> threads(THREADSY, THREADSX); + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); + #else dim3 threads(THREADSX, THREADSY); dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + #endif int li = envs->i_l; int lj = envs->j_l; int lk = envs->k_l; const int type_ijkl = li * 100 + lj * 10 + lk; switch (type_ijkl) { // nroots = 1 - case 0: GINTfill_int3c2e_kernel0000<<>>(*envs, *eri, *offsets); break; - case 1: GINTfill_int3c2e_kernel0010<<>>(*envs, *eri, *offsets); break; - case 10: GINTfill_int3c2e_kernel0100<<>>(*envs, *eri, *offsets); break; - case 100: GINTfill_int3c2e_kernel1000<<>>(*envs, *eri, *offsets); break; + case 0: LAUNCH_KERNEL(GINTfill_int3c2e_kernel0000) break; + case 1: LAUNCH_KERNEL(GINTfill_int3c2e_kernel0010) break; + case 10: LAUNCH_KERNEL(GINTfill_int3c2e_kernel0100) break; + case 100: LAUNCH_KERNEL(GINTfill_int3c2e_kernel1000) break; // nroots = 2 - case 2: GINTfill_int2e_kernel0020<<>>(*envs, *eri, *offsets); break; - case 3: GINTfill_int2e_kernel0030<<>>(*envs, *eri, *offsets); break; - case 101: GINTfill_int2e_kernel1010<<>>(*envs, *eri, *offsets); break; - case 102: GINTfill_int2e_kernel1020<<>>(*envs, *eri, *offsets); break; - case 110: GINTfill_int2e_kernel1100<<>>(*envs, *eri, *offsets); break; - case 111: GINTfill_int2e_kernel1110<<>>(*envs, *eri, *offsets); break; - case 200: GINTfill_int2e_kernel2000<<>>(*envs, *eri, *offsets); break; - case 201: GINTfill_int2e_kernel2010<<>>(*envs, *eri, *offsets); break; - case 210: GINTfill_int2e_kernel2100<<>>(*envs, *eri, *offsets); break; - case 300: GINTfill_int2e_kernel3000<<>>(*envs, *eri, *offsets); break; + case 2: LAUNCH_KERNEL(GINTfill_int2e_kernel0020) break; + case 3: LAUNCH_KERNEL(GINTfill_int2e_kernel0030) break; + case 101: LAUNCH_KERNEL(GINTfill_int2e_kernel1010) break; + case 102: LAUNCH_KERNEL(GINTfill_int2e_kernel1020) break; + case 110: LAUNCH_KERNEL(GINTfill_int2e_kernel1100) break; + case 111: LAUNCH_KERNEL(GINTfill_int2e_kernel1110) break; + case 200: LAUNCH_KERNEL(GINTfill_int2e_kernel2000) break; + case 201: LAUNCH_KERNEL(GINTfill_int2e_kernel2010) break; + case 210: LAUNCH_KERNEL(GINTfill_int2e_kernel2100) break; + case 300: LAUNCH_KERNEL(GINTfill_int2e_kernel3000) break; // nroots = 3 - case 103: GINTfill_int2e_kernel1030<<>>(*envs, *eri, *offsets); break; - case 112: GINTfill_int2e_kernel1120<<>>(*envs, *eri, *offsets); break; - case 113: GINTfill_int2e_kernel1130<<>>(*envs, *eri, *offsets); break; - case 202: GINTfill_int2e_kernel2020<<>>(*envs, *eri, *offsets); break; - case 203: GINTfill_int2e_kernel2030<<>>(*envs, *eri, *offsets); break; - case 211: GINTfill_int2e_kernel2110<<>>(*envs, *eri, *offsets); break; - case 212: GINTfill_int2e_kernel2120<<>>(*envs, *eri, *offsets); break; - case 220: GINTfill_int2e_kernel2200<<>>(*envs, *eri, *offsets); break; - case 221: GINTfill_int2e_kernel2210<<>>(*envs, *eri, *offsets); break; - case 301: GINTfill_int2e_kernel3010<<>>(*envs, *eri, *offsets); break; - case 302: GINTfill_int2e_kernel3020<<>>(*envs, *eri, *offsets); break; - case 310: GINTfill_int2e_kernel3100<<>>(*envs, *eri, *offsets); break; - case 311: GINTfill_int2e_kernel3110<<>>(*envs, *eri, *offsets); break; - case 320: GINTfill_int2e_kernel3200<<>>(*envs, *eri, *offsets); break; + case 103: LAUNCH_KERNEL(GINTfill_int2e_kernel1030) break; + case 112: LAUNCH_KERNEL(GINTfill_int2e_kernel1120) break; + case 113: LAUNCH_KERNEL(GINTfill_int2e_kernel1130) break; + case 202: LAUNCH_KERNEL(GINTfill_int2e_kernel2020) break; + case 203: LAUNCH_KERNEL(GINTfill_int2e_kernel2030) break; + case 211: LAUNCH_KERNEL(GINTfill_int2e_kernel2110) break; + case 212: LAUNCH_KERNEL(GINTfill_int2e_kernel2120) break; + case 220: LAUNCH_KERNEL(GINTfill_int2e_kernel2200) break; + case 221: LAUNCH_KERNEL(GINTfill_int2e_kernel2210) break; + case 301: LAUNCH_KERNEL(GINTfill_int2e_kernel3010) break; + case 302: LAUNCH_KERNEL(GINTfill_int2e_kernel3020) break; + case 310: LAUNCH_KERNEL(GINTfill_int2e_kernel3100) break; + case 311: LAUNCH_KERNEL(GINTfill_int2e_kernel3110) break; + case 320: LAUNCH_KERNEL(GINTfill_int2e_kernel3200) break; default: { + const int gsize = 3*nrys_roots*(li+1)*(lj+1)*(lk+1); +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADSX*THREADSY); + sycl::range<2> blocks(ntasks_kl, ntasks_ij); + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(gsize+16), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GINTfill_int3c2e_kernel(dev_envs, dev_eri, dev_offsets, item, + GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); }); +#else dim3 threads(THREADSX*THREADSY); dim3 blocks(ntasks_ij, ntasks_kl); - const int gsize = 3*nrys_roots*(li+1)*(lj+1)*(lk+1); - cudaError_t err = cudaFuncSetAttribute( + cudaError_t ferr = cudaFuncSetAttribute( GINTfill_int3c2e_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (gsize+16)*sizeof(double)); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error of GINTfill_int3c2e_kernel: %s\n", cudaGetErrorString(err)); + if (ferr != cudaSuccess) { + fprintf(stderr, "CUDA Error of GINTfill_int3c2e_kernel: %s\n", cudaGetErrorString(ferr)); return 1; } const int shm_size = gsize*sizeof(double); GINTfill_int3c2e_kernel<<>>(*envs, *eri, *offsets); +#endif } } @@ -123,7 +159,11 @@ int GINTfill_int3c2e(cudaStream_t stream, BasisProdCache *bpcache, double *eri, //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif ERITensor eritensor; eritensor.stride_j = strides[1]; diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1.cu index ab5875ffb..992fdd4cc 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "config.h" #include "cuda_alloc.cuh" @@ -31,6 +32,42 @@ #include "gout3c2e.cu" #include "g3c2e_ip1.cu" +// Abstracts 2D thread/block config (THREADSX/Y swapped between SYCL and CUDA). +// Used 1x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + dim3 threads(THREADSX, THREADSY); \ + dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. dev_envs/dev_eri/dev_offsets are value copies +// hoisted unconditionally so both branches use identical argument names. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function (with template args if needed) +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel-id (with any template args) is the trailing __VA_ARGS__ so its commas +// survive macro expansion. SYCL kernel name is generated inline per source line. +// dev_envs/dev_eri/dev_offsets are on-host value copies made just before launch +// for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + __host__ static int GINTfill_int3c2e_ip1_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { @@ -38,82 +75,93 @@ static int GINTfill_int3c2e_ip1_tasks(ERITensor *eri, BasisProdOffsets *offsets, int ntasks_ij = offsets->ntasks_ij; int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); - dim3 threads(THREADSX, THREADSY); - dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); int li = envs->i_l; int lj = envs->j_l; int lk = envs->k_l; int type_ijk = li * 100 + lj * 10 + lk; switch (type_ijk) { - case 0: GINTfill_int3c2e_ip1_kernel000<<>>(*envs, *eri, *offsets); break; + case 0: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel000) break; // li+lj+lk=1 - case 1: GINTfill_int3c2e_ip1_kernel<0,0,1><<>>(*envs, *eri, *offsets); break; - case 10: GINTfill_int3c2e_ip1_kernel<0,1,0><<>>(*envs, *eri, *offsets); break; - case 100: GINTfill_int3c2e_ip1_kernel<1,0,0><<>>(*envs, *eri, *offsets); break; + case 1: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,0,1>) break; + case 10: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,1,0>) break; + case 100: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,0,0>) break; // li+lj+lk=2 - case 2: GINTfill_int3c2e_ip1_kernel<0,0,2><<>>(*envs, *eri, *offsets); break; - case 11: GINTfill_int3c2e_ip1_kernel<0,1,1><<>>(*envs, *eri, *offsets); break; - case 20: GINTfill_int3c2e_ip1_kernel<0,2,0><<>>(*envs, *eri, *offsets); break; - case 101: GINTfill_int3c2e_ip1_kernel<1,0,1><<>>(*envs, *eri, *offsets); break; - case 110: GINTfill_int3c2e_ip1_kernel<1,1,0><<>>(*envs, *eri, *offsets); break; - case 200: GINTfill_int3c2e_ip1_kernel<2,0,0><<>>(*envs, *eri, *offsets); break; + case 2: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,0,2>) break; + case 11: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,1,1>) break; + case 20: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,2,0>) break; + case 101: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,0,1>) break; + case 110: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,1,0>) break; + case 200: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,0,0>) break; // li+lj+lk=3 - case 3: GINTfill_int3c2e_ip1_kernel<0,0,3><<>>(*envs, *eri, *offsets); break; - case 12: GINTfill_int3c2e_ip1_kernel<0,1,2><<>>(*envs, *eri, *offsets); break; - case 21: GINTfill_int3c2e_ip1_kernel<0,2,1><<>>(*envs, *eri, *offsets); break; - case 30: GINTfill_int3c2e_ip1_kernel<0,3,0><<>>(*envs, *eri, *offsets); break; - case 102: GINTfill_int3c2e_ip1_kernel<1,0,2><<>>(*envs, *eri, *offsets); break; - case 111: GINTfill_int3c2e_ip1_kernel<1,1,1><<>>(*envs, *eri, *offsets); break; - case 120: GINTfill_int3c2e_ip1_kernel<1,2,0><<>>(*envs, *eri, *offsets); break; - case 201: GINTfill_int3c2e_ip1_kernel<2,0,1><<>>(*envs, *eri, *offsets); break; - case 210: GINTfill_int3c2e_ip1_kernel<2,1,0><<>>(*envs, *eri, *offsets); break; - case 300: GINTfill_int3c2e_ip1_kernel<3,0,0><<>>(*envs, *eri, *offsets); break; + case 3: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,0,3>) break; + case 12: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,1,2>) break; + case 21: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,2,1>) break; + case 30: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,3,0>) break; + case 102: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,0,2>) break; + case 111: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,1,1>) break; + case 120: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,2,0>) break; + case 201: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,0,1>) break; + case 210: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,1,0>) break; + case 300: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<3,0,0>) break; // li+lj+lk=4 - case 4: GINTfill_int3c2e_ip1_kernel<0,0,4><<>>(*envs, *eri, *offsets); break; - case 13: GINTfill_int3c2e_ip1_kernel<0,1,3><<>>(*envs, *eri, *offsets); break; - case 22: GINTfill_int3c2e_ip1_kernel<0,2,2><<>>(*envs, *eri, *offsets); break; - case 31: GINTfill_int3c2e_ip1_kernel<0,3,1><<>>(*envs, *eri, *offsets); break; - case 40: GINTfill_int3c2e_ip1_kernel<0,4,0><<>>(*envs, *eri, *offsets); break; - case 103: GINTfill_int3c2e_ip1_kernel<1,0,3><<>>(*envs, *eri, *offsets); break; - case 112: GINTfill_int3c2e_ip1_kernel<1,1,2><<>>(*envs, *eri, *offsets); break; - case 121: GINTfill_int3c2e_ip1_kernel<1,2,1><<>>(*envs, *eri, *offsets); break; - case 130: GINTfill_int3c2e_ip1_kernel<1,3,0><<>>(*envs, *eri, *offsets); break; - case 202: GINTfill_int3c2e_ip1_kernel<2,0,2><<>>(*envs, *eri, *offsets); break; - case 211: GINTfill_int3c2e_ip1_kernel<2,1,1><<>>(*envs, *eri, *offsets); break; - case 220: GINTfill_int3c2e_ip1_kernel<2,2,0><<>>(*envs, *eri, *offsets); break; - case 301: GINTfill_int3c2e_ip1_kernel<3,0,1><<>>(*envs, *eri, *offsets); break; - case 310: GINTfill_int3c2e_ip1_kernel<3,1,0><<>>(*envs, *eri, *offsets); break; - case 400: GINTfill_int3c2e_ip1_kernel<4,0,0><<>>(*envs, *eri, *offsets); break; + case 4: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,0,4>) break; + case 13: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,1,3>) break; + case 22: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,2,2>) break; + case 31: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,3,1>) break; + case 40: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,4,0>) break; + case 103: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,0,3>) break; + case 112: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,1,2>) break; + case 121: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,2,1>) break; + case 130: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,3,0>) break; + case 202: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,0,2>) break; + case 211: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,1,1>) break; + case 220: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,2,0>) break; + case 301: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<3,0,1>) break; + case 310: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<3,1,0>) break; + case 400: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<4,0,0>) break; // li+lj+lk=5 - //case 5: GINTfill_int3c2e_ip1_kernel<0,0,5><<>>(*envs, *eri, *offsets); break; - case 14: GINTfill_int3c2e_ip1_kernel<0,1,4><<>>(*envs, *eri, *offsets); break; - case 23: GINTfill_int3c2e_ip1_kernel<0,2,3><<>>(*envs, *eri, *offsets); break; - case 32: GINTfill_int3c2e_ip1_kernel<0,3,2><<>>(*envs, *eri, *offsets); break; - case 41: GINTfill_int3c2e_ip1_kernel<0,4,1><<>>(*envs, *eri, *offsets); break; - //case 50: GINTfill_int3c2e_ip1_kernel<0,5,0><<>>(*envs, *eri, *offsets); break; - case 104: GINTfill_int3c2e_ip1_kernel<1,0,4><<>>(*envs, *eri, *offsets); break; - case 113: GINTfill_int3c2e_ip1_kernel<1,1,3><<>>(*envs, *eri, *offsets); break; - case 122: GINTfill_int3c2e_ip1_kernel<1,2,2><<>>(*envs, *eri, *offsets); break; - case 131: GINTfill_int3c2e_ip1_kernel<1,3,1><<>>(*envs, *eri, *offsets); break; - case 140: GINTfill_int3c2e_ip1_kernel<1,4,0><<>>(*envs, *eri, *offsets); break; - case 203: GINTfill_int3c2e_ip1_kernel<2,0,3><<>>(*envs, *eri, *offsets); break; - case 212: GINTfill_int3c2e_ip1_kernel<2,1,2><<>>(*envs, *eri, *offsets); break; - case 221: GINTfill_int3c2e_ip1_kernel<2,2,1><<>>(*envs, *eri, *offsets); break; - case 230: GINTfill_int3c2e_ip1_kernel<2,3,0><<>>(*envs, *eri, *offsets); break; - case 302: GINTfill_int3c2e_ip1_kernel<3,0,2><<>>(*envs, *eri, *offsets); break; - case 311: GINTfill_int3c2e_ip1_kernel<3,1,1><<>>(*envs, *eri, *offsets); break; - case 320: GINTfill_int3c2e_ip1_kernel<3,2,0><<>>(*envs, *eri, *offsets); break; - case 401: GINTfill_int3c2e_ip1_kernel<4,0,1><<>>(*envs, *eri, *offsets); break; - case 410: GINTfill_int3c2e_ip1_kernel<4,1,0><<>>(*envs, *eri, *offsets); break; + //case 5: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,0,5>) break; + case 14: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,1,4>) break; + case 23: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,2,3>) break; + case 32: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,3,2>) break; + case 41: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,4,1>) break; + //case 50: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<0,5,0>) break; + case 104: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,0,4>) break; + case 113: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,1,3>) break; + case 122: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,2,2>) break; + case 131: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,3,1>) break; + case 140: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<1,4,0>) break; + case 203: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,0,3>) break; + case 212: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,1,2>) break; + case 221: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,2,1>) break; + case 230: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<2,3,0>) break; + case 302: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<3,0,2>) break; + case 311: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<3,1,1>) break; + case 320: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<3,2,0>) break; + case 401: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<4,0,1>) break; + case 410: LAUNCH_KERNEL(GINTfill_int3c2e_ip1_kernel<4,1,0>) break; //case 500: GINTfill_int3c2e_ip1_kernel<5,0,0><<>>(*envs, *eri, *offsets); break; #ifdef UNROLL_INT3C2E #endif default: { - dim3 threads(THREADSX*THREADSY); - dim3 blocks(ntasks_ij, ntasks_kl); const int li_ceil = li + 1; const int gsize = 3*nrys_roots*(li_ceil+1)*(lj+1)*(lk+1); + + #ifdef USE_SYCL + sycl::range<2> threads(1, THREADSX*THREADSY); + sycl::range<2> blocks(ntasks_kl, ntasks_ij); + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(gsize+16), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GINTfill_int3c2e_ip1_general_kernel(dev_envs, dev_eri, dev_offsets, item, + GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); }); + #else + dim3 threads(THREADSX*THREADSY); + dim3 blocks(ntasks_ij, ntasks_kl); cudaError_t err = cudaFuncSetAttribute( GINTfill_int3c2e_ip1_general_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -124,14 +172,17 @@ static int GINTfill_int3c2e_ip1_tasks(ERITensor *eri, BasisProdOffsets *offsets, } const int shm_size = gsize*sizeof(double); GINTfill_int3c2e_ip1_general_kernel<<>>(*envs, *eri, *offsets); + #endif } } + #ifndef USE_SYCL cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ip1_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -154,7 +205,11 @@ int GINTfill_int3c2e_ip1(cudaStream_t stream, BasisProdCache *bpcache, double *e //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif ERITensor eritensor; eritensor.stride_j = strides[1]; @@ -203,3 +258,6 @@ int GINTfill_int3c2e_ip1(cudaStream_t stream, BasisProdCache *bpcache, double *e } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1ip2.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1ip2.cu index 937c69a10..6e2590214 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1ip2.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip1ip2.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "config.h" #include "cuda_alloc.cuh" @@ -31,6 +32,42 @@ #include "gout3c2e.cu" #include "g3c2e_ip1ip2.cu" +// Abstracts 2D thread/block config (THREADSX/Y swapped between SYCL and CUDA). +// Used 1x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + dim3 threads(THREADSX, THREADSY); \ + dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. dev_envs/dev_eri/dev_offsets are value copies +// hoisted unconditionally so both branches use identical argument names. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function (with template args if needed) +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel-id (with any template args) is the trailing __VA_ARGS__ so its commas +// survive macro expansion. SYCL kernel name is generated inline per source line. +// dev_envs/dev_eri/dev_offsets are on-host value copies made just before launch +// for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + __host__ static int GINTfill_int3c2e_ip1ip2_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { @@ -38,8 +75,7 @@ static int GINTfill_int3c2e_ip1ip2_tasks(ERITensor *eri, BasisProdOffsets *offse int ntasks_ij = offsets->ntasks_ij; int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); - dim3 threads(THREADSX, THREADSY); - dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); int li = envs->i_l; int lj = envs->j_l; int lk = envs->k_l; @@ -47,75 +83,86 @@ static int GINTfill_int3c2e_ip1ip2_tasks(ERITensor *eri, BasisProdOffsets *offse switch (type_ijk) { // li+lj+lk=0 - case 0: GINTfill_int3c2e_ip1ip2_kernel000<<>>(*envs, *eri, *offsets); break; + case 0: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel000) break; // li+lj+lk=1 - case 1: GINTfill_int3c2e_ip1ip2_kernel<0,0,1><<>>(*envs, *eri, *offsets); break; - case 10: GINTfill_int3c2e_ip1ip2_kernel<0,1,0><<>>(*envs, *eri, *offsets); break; - case 100: GINTfill_int3c2e_ip1ip2_kernel<1,0,0><<>>(*envs, *eri, *offsets); break; + case 1: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,0,1>) break; + case 10: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,1,0>) break; + case 100: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,0,0>) break; // li+lj+lk=2 - case 2: GINTfill_int3c2e_ip1ip2_kernel<0,0,2><<>>(*envs, *eri, *offsets); break; - case 11: GINTfill_int3c2e_ip1ip2_kernel<0,1,1><<>>(*envs, *eri, *offsets); break; - case 20: GINTfill_int3c2e_ip1ip2_kernel<0,2,0><<>>(*envs, *eri, *offsets); break; - case 101: GINTfill_int3c2e_ip1ip2_kernel<1,0,1><<>>(*envs, *eri, *offsets); break; - case 110: GINTfill_int3c2e_ip1ip2_kernel<1,1,0><<>>(*envs, *eri, *offsets); break; - case 200: GINTfill_int3c2e_ip1ip2_kernel<2,0,0><<>>(*envs, *eri, *offsets); break; + case 2: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,0,2>) break; + case 11: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,1,1>) break; + case 20: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,2,0>) break; + case 101: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,0,1>) break; + case 110: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,1,0>) break; + case 200: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,0,0>) break; // li+lj+lk=3 - case 3: GINTfill_int3c2e_ip1ip2_kernel<0,0,3><<>>(*envs, *eri, *offsets); break; - case 12: GINTfill_int3c2e_ip1ip2_kernel<0,1,2><<>>(*envs, *eri, *offsets); break; - case 21: GINTfill_int3c2e_ip1ip2_kernel<0,2,1><<>>(*envs, *eri, *offsets); break; - case 30: GINTfill_int3c2e_ip1ip2_kernel<0,3,0><<>>(*envs, *eri, *offsets); break; - case 102: GINTfill_int3c2e_ip1ip2_kernel<1,0,2><<>>(*envs, *eri, *offsets); break; - case 111: GINTfill_int3c2e_ip1ip2_kernel<1,1,1><<>>(*envs, *eri, *offsets); break; - case 120: GINTfill_int3c2e_ip1ip2_kernel<1,2,0><<>>(*envs, *eri, *offsets); break; - case 201: GINTfill_int3c2e_ip1ip2_kernel<2,0,1><<>>(*envs, *eri, *offsets); break; - case 210: GINTfill_int3c2e_ip1ip2_kernel<2,1,0><<>>(*envs, *eri, *offsets); break; - case 300: GINTfill_int3c2e_ip1ip2_kernel<3,0,0><<>>(*envs, *eri, *offsets); break; + case 3: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,0,3>) break; + case 12: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,1,2>) break; + case 21: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,2,1>) break; + case 30: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,3,0>) break; + case 102: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,0,2>) break; + case 111: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,1,1>) break; + case 120: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,2,0>) break; + case 201: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,0,1>) break; + case 210: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,1,0>) break; + case 300: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<3,0,0>) break; // li+lj+lk=4 - case 4: GINTfill_int3c2e_ip1ip2_kernel<0,0,4><<>>(*envs, *eri, *offsets); break; - case 13: GINTfill_int3c2e_ip1ip2_kernel<0,1,3><<>>(*envs, *eri, *offsets); break; - case 22: GINTfill_int3c2e_ip1ip2_kernel<0,2,2><<>>(*envs, *eri, *offsets); break; - case 31: GINTfill_int3c2e_ip1ip2_kernel<0,3,1><<>>(*envs, *eri, *offsets); break; - case 40: GINTfill_int3c2e_ip1ip2_kernel<0,4,0><<>>(*envs, *eri, *offsets); break; - case 103: GINTfill_int3c2e_ip1ip2_kernel<1,0,3><<>>(*envs, *eri, *offsets); break; - case 112: GINTfill_int3c2e_ip1ip2_kernel<1,1,2><<>>(*envs, *eri, *offsets); break; - case 121: GINTfill_int3c2e_ip1ip2_kernel<1,2,1><<>>(*envs, *eri, *offsets); break; - case 130: GINTfill_int3c2e_ip1ip2_kernel<1,3,0><<>>(*envs, *eri, *offsets); break; - case 202: GINTfill_int3c2e_ip1ip2_kernel<2,0,2><<>>(*envs, *eri, *offsets); break; - case 211: GINTfill_int3c2e_ip1ip2_kernel<2,1,1><<>>(*envs, *eri, *offsets); break; - case 220: GINTfill_int3c2e_ip1ip2_kernel<2,2,0><<>>(*envs, *eri, *offsets); break; - case 301: GINTfill_int3c2e_ip1ip2_kernel<3,0,1><<>>(*envs, *eri, *offsets); break; - case 310: GINTfill_int3c2e_ip1ip2_kernel<3,1,0><<>>(*envs, *eri, *offsets); break; - case 400: GINTfill_int3c2e_ip1ip2_kernel<4,0,0><<>>(*envs, *eri, *offsets); break; + case 4: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,0,4>) break; + case 13: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,1,3>) break; + case 22: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,2,2>) break; + case 31: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,3,1>) break; + case 40: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,4,0>) break; + case 103: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,0,3>) break; + case 112: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,1,2>) break; + case 121: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,2,1>) break; + case 130: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,3,0>) break; + case 202: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,0,2>) break; + case 211: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,1,1>) break; + case 220: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,2,0>) break; + case 301: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<3,0,1>) break; + case 310: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<3,1,0>) break; + case 400: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<4,0,0>) break; // li+lj+lk=5 - //case 5: GINTfill_int3c2e_ip1ip2_kernel<0,0,5><<>>(*envs, *eri, *offsets); break; - case 14: GINTfill_int3c2e_ip1ip2_kernel<0,1,4><<>>(*envs, *eri, *offsets); break; - case 23: GINTfill_int3c2e_ip1ip2_kernel<0,2,3><<>>(*envs, *eri, *offsets); break; - case 32: GINTfill_int3c2e_ip1ip2_kernel<0,3,2><<>>(*envs, *eri, *offsets); break; - case 41: GINTfill_int3c2e_ip1ip2_kernel<0,4,1><<>>(*envs, *eri, *offsets); break; - //case 50: GINTfill_int3c2e_ip1ip2_kernel<0,5,0><<>>(*envs, *eri, *offsets); break; - case 104: GINTfill_int3c2e_ip1ip2_kernel<1,0,4><<>>(*envs, *eri, *offsets); break; - case 113: GINTfill_int3c2e_ip1ip2_kernel<1,1,3><<>>(*envs, *eri, *offsets); break; - case 122: GINTfill_int3c2e_ip1ip2_kernel<1,2,2><<>>(*envs, *eri, *offsets); break; - case 131: GINTfill_int3c2e_ip1ip2_kernel<1,3,1><<>>(*envs, *eri, *offsets); break; - case 140: GINTfill_int3c2e_ip1ip2_kernel<1,4,0><<>>(*envs, *eri, *offsets); break; - case 203: GINTfill_int3c2e_ip1ip2_kernel<2,0,3><<>>(*envs, *eri, *offsets); break; - case 212: GINTfill_int3c2e_ip1ip2_kernel<2,1,2><<>>(*envs, *eri, *offsets); break; - case 221: GINTfill_int3c2e_ip1ip2_kernel<2,2,1><<>>(*envs, *eri, *offsets); break; - case 230: GINTfill_int3c2e_ip1ip2_kernel<2,3,0><<>>(*envs, *eri, *offsets); break; - case 302: GINTfill_int3c2e_ip1ip2_kernel<3,0,2><<>>(*envs, *eri, *offsets); break; - case 311: GINTfill_int3c2e_ip1ip2_kernel<3,1,1><<>>(*envs, *eri, *offsets); break; - case 320: GINTfill_int3c2e_ip1ip2_kernel<3,2,0><<>>(*envs, *eri, *offsets); break; - case 401: GINTfill_int3c2e_ip1ip2_kernel<4,0,1><<>>(*envs, *eri, *offsets); break; - case 410: GINTfill_int3c2e_ip1ip2_kernel<4,1,0><<>>(*envs, *eri, *offsets); break; - //case 500: GINTfill_int3c2e_ip1ip2_kernel<5,0,0><<>>(*envs, *eri, *offsets); break; + //case 5: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,0,5>) break; + case 14: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,1,4>) break; + case 23: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,2,3>) break; + case 32: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,3,2>) break; + case 41: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,4,1>) break; + //case 50: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<0,5,0>) break; + case 104: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,0,4>) break; + case 113: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,1,3>) break; + case 122: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,2,2>) break; + case 131: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,3,1>) break; + case 140: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<1,4,0>) break; + case 203: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,0,3>) break; + case 212: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,1,2>) break; + case 221: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,2,1>) break; + case 230: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<2,3,0>) break; + case 302: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<3,0,2>) break; + case 311: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<3,1,1>) break; + case 320: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<3,2,0>) break; + case 401: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<4,0,1>) break; + case 410: LAUNCH_KERNEL(GINTfill_int3c2e_ip1ip2_kernel<4,1,0>) break; + //case 500: GINTfill_int3c2e_ip1ip2_kernel<5,0,0>(dev_envs, dev_eri, dev_offsets); }); break; #ifdef UNROLL_INT3C2E #endif default: { - dim3 threads(THREADSX*THREADSY); - dim3 blocks(ntasks_ij, ntasks_kl); const int li_ceil = li + 1; const int lk_ceil = lk + 1; const int gsize = 3*nrys_roots*(li_ceil+1)*(lj+1)*(lk_ceil+1); + #ifdef USE_SYCL + sycl::range<2> threads(1, THREADSX*THREADSY); + sycl::range<2> blocks(ntasks_kl, ntasks_ij); + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(gsize+16), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GINTfill_int3c2e_ip1ip2_general_kernel(dev_envs, dev_eri, dev_offsets, item, + GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); }); + #else + dim3 threads(THREADSX*THREADSY); + dim3 blocks(ntasks_ij, ntasks_kl); cudaError_t err = cudaFuncSetAttribute( GINTfill_int3c2e_ip1ip2_general_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -123,15 +170,17 @@ static int GINTfill_int3c2e_ip1ip2_tasks(ERITensor *eri, BasisProdOffsets *offse const int shm_size = gsize*sizeof(double); GINTfill_int3c2e_ip1ip2_general_kernel<<>>( *envs, *eri, *offsets); - + #endif } } +#ifndef USE_SYCL cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ip1ip2_kernel: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -156,7 +205,11 @@ int GINTfill_int3c2e_ip1ip2(cudaStream_t stream, BasisProdCache *bpcache, double //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory +#ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); +#else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); +#endif ERITensor eritensor; eritensor.stride_j = strides[1]; @@ -204,3 +257,6 @@ int GINTfill_int3c2e_ip1ip2(cudaStream_t stream, BasisProdCache *bpcache, double } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip2.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip2.cu index 012ace164..ca30b80e7 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip2.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ip2.cu @@ -20,9 +20,9 @@ #include #include +#include "cuda_alloc.cuh" #include "gint.h" #include "config.h" -#include "cuda_alloc.cuh" #include "g2e.h" #include "rys_roots.cu" @@ -31,6 +31,42 @@ #include "gout3c2e.cu" #include "g3c2e_ip2.cu" +// Abstracts 2D thread/block config (THREADSX/Y swapped between SYCL and CUDA). +// Used 1x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + dim3 threads(THREADSX, THREADSY); \ + dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. dev_envs/dev_eri/dev_offsets are value copies +// hoisted unconditionally so both branches use identical argument names. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function (with template args if needed) +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel-id (with any template args) is the trailing __VA_ARGS__ so its commas +// survive macro expansion. SYCL kernel name is generated inline per source line. +// dev_envs/dev_eri/dev_offsets are on-host value copies made just before launch +// for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + __host__ static int GINTfill_int3c2e_ip2_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { @@ -38,96 +74,109 @@ static int GINTfill_int3c2e_ip2_tasks(ERITensor *eri, BasisProdOffsets *offsets, int ntasks_ij = offsets->ntasks_ij; int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); - dim3 threads(THREADSX, THREADSY); - dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); int li = envs->i_l; int lj = envs->j_l; int lk = envs->k_l; int type_ijk = li * 100 + lj * 10 + lk; switch (type_ijk) { - case 0: GINTfill_int3c2e_ip2_kernel000<<>>(*envs, *eri, *offsets); break; + case 0: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel000) break; // li+lj+lk=1 - case 1: GINTfill_int3c2e_ip2_kernel<0,0,1><<>>(*envs, *eri, *offsets); break; - case 10: GINTfill_int3c2e_ip2_kernel<0,1,0><<>>(*envs, *eri, *offsets); break; - case 100: GINTfill_int3c2e_ip2_kernel<1,0,0><<>>(*envs, *eri, *offsets); break; + case 1: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,0,1>) break; + case 10: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,1,0>) break; + case 100: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,0,0>) break; // li+lj+lk=2 - case 2: GINTfill_int3c2e_ip2_kernel<0,0,2><<>>(*envs, *eri, *offsets); break; - case 11: GINTfill_int3c2e_ip2_kernel<0,1,1><<>>(*envs, *eri, *offsets); break; - case 20: GINTfill_int3c2e_ip2_kernel<0,2,0><<>>(*envs, *eri, *offsets); break; - case 101: GINTfill_int3c2e_ip2_kernel<1,0,1><<>>(*envs, *eri, *offsets); break; - case 110: GINTfill_int3c2e_ip2_kernel<1,1,0><<>>(*envs, *eri, *offsets); break; - case 200: GINTfill_int3c2e_ip2_kernel<2,0,0><<>>(*envs, *eri, *offsets); break; + case 2: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,0,2>) break; + case 11: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,1,1>) break; + case 20: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,2,0>) break; + case 101: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,0,1>) break; + case 110: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,1,0>) break; + case 200: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,0,0>) break; // li+lj+lk=3 - case 3: GINTfill_int3c2e_ip2_kernel<0,0,3><<>>(*envs, *eri, *offsets); break; - case 12: GINTfill_int3c2e_ip2_kernel<0,1,2><<>>(*envs, *eri, *offsets); break; - case 21: GINTfill_int3c2e_ip2_kernel<0,2,1><<>>(*envs, *eri, *offsets); break; - case 30: GINTfill_int3c2e_ip2_kernel<0,3,0><<>>(*envs, *eri, *offsets); break; - case 102: GINTfill_int3c2e_ip2_kernel<1,0,2><<>>(*envs, *eri, *offsets); break; - case 111: GINTfill_int3c2e_ip2_kernel<1,1,1><<>>(*envs, *eri, *offsets); break; - case 120: GINTfill_int3c2e_ip2_kernel<1,2,0><<>>(*envs, *eri, *offsets); break; - case 201: GINTfill_int3c2e_ip2_kernel<2,0,1><<>>(*envs, *eri, *offsets); break; - case 210: GINTfill_int3c2e_ip2_kernel<2,1,0><<>>(*envs, *eri, *offsets); break; - case 300: GINTfill_int3c2e_ip2_kernel<3,0,0><<>>(*envs, *eri, *offsets); break; + case 3: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,0,3>) break; + case 12: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,1,2>) break; + case 21: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,2,1>) break; + case 30: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,3,0>) break; + case 102: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,0,2>) break; + case 111: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,1,1>) break; + case 120: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,2,0>) break; + case 201: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,0,1>) break; + case 210: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,1,0>) break; + case 300: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<3,0,0>) break; // li+lj+lk=4 - case 4: GINTfill_int3c2e_ip2_kernel<0,0,4><<>>(*envs, *eri, *offsets); break; - case 13: GINTfill_int3c2e_ip2_kernel<0,1,3><<>>(*envs, *eri, *offsets); break; - case 22: GINTfill_int3c2e_ip2_kernel<0,2,2><<>>(*envs, *eri, *offsets); break; - case 31: GINTfill_int3c2e_ip2_kernel<0,3,1><<>>(*envs, *eri, *offsets); break; - case 40: GINTfill_int3c2e_ip2_kernel<0,4,0><<>>(*envs, *eri, *offsets); break; - case 103: GINTfill_int3c2e_ip2_kernel<1,0,3><<>>(*envs, *eri, *offsets); break; - case 112: GINTfill_int3c2e_ip2_kernel<1,1,2><<>>(*envs, *eri, *offsets); break; - case 121: GINTfill_int3c2e_ip2_kernel<1,2,1><<>>(*envs, *eri, *offsets); break; - case 130: GINTfill_int3c2e_ip2_kernel<1,3,0><<>>(*envs, *eri, *offsets); break; - case 202: GINTfill_int3c2e_ip2_kernel<2,0,2><<>>(*envs, *eri, *offsets); break; - case 211: GINTfill_int3c2e_ip2_kernel<2,1,1><<>>(*envs, *eri, *offsets); break; - case 220: GINTfill_int3c2e_ip2_kernel<2,2,0><<>>(*envs, *eri, *offsets); break; - case 301: GINTfill_int3c2e_ip2_kernel<3,0,1><<>>(*envs, *eri, *offsets); break; - case 310: GINTfill_int3c2e_ip2_kernel<3,1,0><<>>(*envs, *eri, *offsets); break; - case 400: GINTfill_int3c2e_ip2_kernel<4,0,0><<>>(*envs, *eri, *offsets); break; + case 4: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,0,4>) break; + case 13: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,1,3>) break; + case 22: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,2,2>) break; + case 31: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,3,1>) break; + case 40: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,4,0>) break; + case 103: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,0,3>) break; + case 112: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,1,2>) break; + case 121: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,2,1>) break; + case 130: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,3,0>) break; + case 202: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,0,2>) break; + case 211: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,1,1>) break; + case 220: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,2,0>) break; + case 301: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<3,0,1>) break; + case 310: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<3,1,0>) break; + case 400: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<4,0,0>) break; // li+lj+lk=5 - //case 5: GINTfill_int3c2e_ip2_kernel<0,0,5><<>>(*envs, *eri, *offsets); break; - case 14: GINTfill_int3c2e_ip2_kernel<0,1,4><<>>(*envs, *eri, *offsets); break; - case 23: GINTfill_int3c2e_ip2_kernel<0,2,3><<>>(*envs, *eri, *offsets); break; - case 32: GINTfill_int3c2e_ip2_kernel<0,3,2><<>>(*envs, *eri, *offsets); break; - case 41: GINTfill_int3c2e_ip2_kernel<0,4,1><<>>(*envs, *eri, *offsets); break; - //case 50: GINTfill_int3c2e_ip2_kernel<0,5,0><<>>(*envs, *eri, *offsets); break; - case 104: GINTfill_int3c2e_ip2_kernel<1,0,4><<>>(*envs, *eri, *offsets); break; - case 113: GINTfill_int3c2e_ip2_kernel<1,1,3><<>>(*envs, *eri, *offsets); break; - case 122: GINTfill_int3c2e_ip2_kernel<1,2,2><<>>(*envs, *eri, *offsets); break; - case 131: GINTfill_int3c2e_ip2_kernel<1,3,1><<>>(*envs, *eri, *offsets); break; - case 140: GINTfill_int3c2e_ip2_kernel<1,4,0><<>>(*envs, *eri, *offsets); break; - case 203: GINTfill_int3c2e_ip2_kernel<2,0,3><<>>(*envs, *eri, *offsets); break; - case 212: GINTfill_int3c2e_ip2_kernel<2,1,2><<>>(*envs, *eri, *offsets); break; - case 221: GINTfill_int3c2e_ip2_kernel<2,2,1><<>>(*envs, *eri, *offsets); break; - case 230: GINTfill_int3c2e_ip2_kernel<2,3,0><<>>(*envs, *eri, *offsets); break; - case 302: GINTfill_int3c2e_ip2_kernel<3,0,2><<>>(*envs, *eri, *offsets); break; - case 311: GINTfill_int3c2e_ip2_kernel<3,1,1><<>>(*envs, *eri, *offsets); break; - case 320: GINTfill_int3c2e_ip2_kernel<3,2,0><<>>(*envs, *eri, *offsets); break; - case 401: GINTfill_int3c2e_ip2_kernel<4,0,1><<>>(*envs, *eri, *offsets); break; - case 410: GINTfill_int3c2e_ip2_kernel<4,1,0><<>>(*envs, *eri, *offsets); break; - //case 500: GINTfill_int3c2e_ip2_kernel<5,0,0><<>>(*envs, *eri, *offsets); break; + //case 5: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,0,5>) break; + case 14: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,1,4>) break; + case 23: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,2,3>) break; + case 32: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,3,2>) break; + case 41: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,4,1>) break; + //case 50: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<0,5,0>) break; + case 104: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,0,4>) break; + case 113: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,1,3>) break; + case 122: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,2,2>) break; + case 131: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,3,1>) break; + case 140: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<1,4,0>) break; + case 203: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,0,3>) break; + case 212: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,1,2>) break; + case 221: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,2,1>) break; + case 230: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<2,3,0>) break; + case 302: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<3,0,2>) break; + case 311: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<3,1,1>) break; + case 320: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<3,2,0>) break; + case 401: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<4,0,1>) break; + case 410: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<4,1,0>) break; + //case 500: LAUNCH_KERNEL(GINTfill_int3c2e_ip2_kernel<5,0,0>) break; #ifdef UNROLL_INT3C2E #endif default: { - dim3 threads(THREADSX*THREADSY); - dim3 blocks(ntasks_ij, ntasks_kl); const int lk_ceil = lk + 1; const int gsize = 3*nrys_roots*(li+1)*(lj+1)*(lk_ceil+1); + #ifdef USE_SYCL + sycl::range<2> threads(1, THREADSX*THREADSY); + sycl::range<2> blocks(ntasks_kl, ntasks_ij); + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(gsize+16), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GINTfill_int3c2e_ip2_general_kernel(dev_envs, dev_eri, dev_offsets, item, + GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); }); + #else + dim3 threads(THREADSX*THREADSY); + dim3 blocks(ntasks_ij, ntasks_kl); cudaError_t err = cudaFuncSetAttribute( GINTfill_int3c2e_ip2_general_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (gsize+16)*sizeof(double)); const int shm_size = gsize*sizeof(double); GINTfill_int3c2e_ip2_general_kernel<<>>(*envs, *eri, *offsets); + #endif } } + #ifndef USE_SYCL cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ip2_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -150,7 +199,11 @@ int GINTfill_int3c2e_ip2(cudaStream_t stream, BasisProdCache *bpcache, double *e //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif ERITensor eritensor; eritensor.stride_j = strides[1]; @@ -198,3 +251,6 @@ int GINTfill_int3c2e_ip2(cudaStream_t stream, BasisProdCache *bpcache, double *e return 0; } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip1.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip1.cu index 0ba1bfa77..18afdfc89 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip1.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip1.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "config.h" #include "cuda_alloc.cuh" @@ -31,6 +32,24 @@ #include "gout3c2e.cu" #include "g3c2e_ipip1.cu" +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel-id (with any template args) is the trailing __VA_ARGS__ so its commas +// survive macro expansion. SYCL kernel name generated inline per source line. +// dev_envs/dev_eri/dev_offsets are on-host value copies for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + __host__ static int GINTfill_int3c2e_ipip1_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { @@ -38,8 +57,13 @@ static int GINTfill_int3c2e_ipip1_tasks(ERITensor *eri, BasisProdOffsets *offset int ntasks_ij = offsets->ntasks_ij; int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); + #ifdef USE_SYCL + sycl::range<2> threads(THREADSY, THREADSX); + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); + #else dim3 threads(THREADSX, THREADSY); dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + #endif int li = envs->i_l; int lj = envs->j_l; int lk = envs->k_l; @@ -47,92 +71,106 @@ static int GINTfill_int3c2e_ipip1_tasks(ERITensor *eri, BasisProdOffsets *offset switch (type_ijk) { // li+lj+lk=0 - case 0: GINTfill_int3c2e_ipip1_kernel000<<>>(*envs, *eri, *offsets); break; + case 0: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel000) break; // li+lj+lk=1 - case 1: GINTfill_int3c2e_ipip1_kernel<0,0,1><<>>(*envs, *eri, *offsets); break; - case 10: GINTfill_int3c2e_ipip1_kernel<0,1,0><<>>(*envs, *eri, *offsets); break; - case 100: GINTfill_int3c2e_ipip1_kernel<1,0,0><<>>(*envs, *eri, *offsets); break; + case 1: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,0,1>) break; + case 10: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,1,0>) break; + case 100: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,0,0>) break; // li+lj+lk=2 - case 2: GINTfill_int3c2e_ipip1_kernel<0,0,2><<>>(*envs, *eri, *offsets); break; - case 11: GINTfill_int3c2e_ipip1_kernel<0,1,1><<>>(*envs, *eri, *offsets); break; - case 20: GINTfill_int3c2e_ipip1_kernel<0,2,0><<>>(*envs, *eri, *offsets); break; - case 101: GINTfill_int3c2e_ipip1_kernel<1,0,1><<>>(*envs, *eri, *offsets); break; - case 110: GINTfill_int3c2e_ipip1_kernel<1,1,0><<>>(*envs, *eri, *offsets); break; - case 200: GINTfill_int3c2e_ipip1_kernel<2,0,0><<>>(*envs, *eri, *offsets); break; + case 2: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,0,2>) break; + case 11: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,1,1>) break; + case 20: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,2,0>) break; + case 101: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,0,1>) break; + case 110: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,1,0>) break; + case 200: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,0,0>) break; // li+lj+lk=3 - case 3: GINTfill_int3c2e_ipip1_kernel<0,0,3><<>>(*envs, *eri, *offsets); break; - case 12: GINTfill_int3c2e_ipip1_kernel<0,1,2><<>>(*envs, *eri, *offsets); break; - case 21: GINTfill_int3c2e_ipip1_kernel<0,2,1><<>>(*envs, *eri, *offsets); break; - case 30: GINTfill_int3c2e_ipip1_kernel<0,3,0><<>>(*envs, *eri, *offsets); break; - case 102: GINTfill_int3c2e_ipip1_kernel<1,0,2><<>>(*envs, *eri, *offsets); break; - case 111: GINTfill_int3c2e_ipip1_kernel<1,1,1><<>>(*envs, *eri, *offsets); break; - case 120: GINTfill_int3c2e_ipip1_kernel<1,2,0><<>>(*envs, *eri, *offsets); break; - case 201: GINTfill_int3c2e_ipip1_kernel<2,0,1><<>>(*envs, *eri, *offsets); break; - case 210: GINTfill_int3c2e_ipip1_kernel<2,1,0><<>>(*envs, *eri, *offsets); break; - case 300: GINTfill_int3c2e_ipip1_kernel<3,0,0><<>>(*envs, *eri, *offsets); break; + case 3: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,0,3>) break; + case 12: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,1,2>) break; + case 21: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,2,1>) break; + case 30: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,3,0>) break; + case 102: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,0,2>) break; + case 111: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,1,1>) break; + case 120: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,2,0>) break; + case 201: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,0,1>) break; + case 210: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,1,0>) break; + case 300: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<3,0,0>) break; // li+lj+lk=4 - case 4: GINTfill_int3c2e_ipip1_kernel<0,0,4><<>>(*envs, *eri, *offsets); break; - case 13: GINTfill_int3c2e_ipip1_kernel<0,1,3><<>>(*envs, *eri, *offsets); break; - case 22: GINTfill_int3c2e_ipip1_kernel<0,2,2><<>>(*envs, *eri, *offsets); break; - case 31: GINTfill_int3c2e_ipip1_kernel<0,3,1><<>>(*envs, *eri, *offsets); break; - case 40: GINTfill_int3c2e_ipip1_kernel<0,4,0><<>>(*envs, *eri, *offsets); break; - case 103: GINTfill_int3c2e_ipip1_kernel<1,0,3><<>>(*envs, *eri, *offsets); break; - case 112: GINTfill_int3c2e_ipip1_kernel<1,1,2><<>>(*envs, *eri, *offsets); break; - case 121: GINTfill_int3c2e_ipip1_kernel<1,2,1><<>>(*envs, *eri, *offsets); break; - case 130: GINTfill_int3c2e_ipip1_kernel<1,3,0><<>>(*envs, *eri, *offsets); break; - case 202: GINTfill_int3c2e_ipip1_kernel<2,0,2><<>>(*envs, *eri, *offsets); break; - case 211: GINTfill_int3c2e_ipip1_kernel<2,1,1><<>>(*envs, *eri, *offsets); break; - case 220: GINTfill_int3c2e_ipip1_kernel<2,2,0><<>>(*envs, *eri, *offsets); break; - case 301: GINTfill_int3c2e_ipip1_kernel<3,0,1><<>>(*envs, *eri, *offsets); break; - case 310: GINTfill_int3c2e_ipip1_kernel<3,1,0><<>>(*envs, *eri, *offsets); break; - case 400: GINTfill_int3c2e_ipip1_kernel<4,0,0><<>>(*envs, *eri, *offsets); break; + case 4: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,0,4>) break; + case 13: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,1,3>) break; + case 22: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,2,2>) break; + case 31: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,3,1>) break; + case 40: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,4,0>) break; + case 103: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,0,3>) break; + case 112: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,1,2>) break; + case 121: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,2,1>) break; + case 130: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,3,0>) break; + case 202: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,0,2>) break; + case 211: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,1,1>) break; + case 220: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,2,0>) break; + case 301: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<3,0,1>) break; + case 310: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<3,1,0>) break; + case 400: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<4,0,0>) break; // li+lj+lk=5 - //case 5: GINTfill_int3c2e_ipip1_kernel<0,0,5><<>>(*envs, *eri, *offsets); break; - case 14: GINTfill_int3c2e_ipip1_kernel<0,1,4><<>>(*envs, *eri, *offsets); break; - case 23: GINTfill_int3c2e_ipip1_kernel<0,2,3><<>>(*envs, *eri, *offsets); break; - case 32: GINTfill_int3c2e_ipip1_kernel<0,3,2><<>>(*envs, *eri, *offsets); break; - case 41: GINTfill_int3c2e_ipip1_kernel<0,4,1><<>>(*envs, *eri, *offsets); break; - //case 50: GINTfill_int3c2e_ipip1_kernel<0,5,0><<>>(*envs, *eri, *offsets); break; - case 104: GINTfill_int3c2e_ipip1_kernel<1,0,4><<>>(*envs, *eri, *offsets); break; - case 113: GINTfill_int3c2e_ipip1_kernel<1,1,3><<>>(*envs, *eri, *offsets); break; - case 122: GINTfill_int3c2e_ipip1_kernel<1,2,2><<>>(*envs, *eri, *offsets); break; - case 131: GINTfill_int3c2e_ipip1_kernel<1,3,1><<>>(*envs, *eri, *offsets); break; - case 140: GINTfill_int3c2e_ipip1_kernel<1,4,0><<>>(*envs, *eri, *offsets); break; - case 203: GINTfill_int3c2e_ipip1_kernel<2,0,3><<>>(*envs, *eri, *offsets); break; - case 212: GINTfill_int3c2e_ipip1_kernel<2,1,2><<>>(*envs, *eri, *offsets); break; - case 221: GINTfill_int3c2e_ipip1_kernel<2,2,1><<>>(*envs, *eri, *offsets); break; - case 230: GINTfill_int3c2e_ipip1_kernel<2,3,0><<>>(*envs, *eri, *offsets); break; - case 302: GINTfill_int3c2e_ipip1_kernel<3,0,2><<>>(*envs, *eri, *offsets); break; - case 311: GINTfill_int3c2e_ipip1_kernel<3,1,1><<>>(*envs, *eri, *offsets); break; - case 320: GINTfill_int3c2e_ipip1_kernel<3,2,0><<>>(*envs, *eri, *offsets); break; - case 401: GINTfill_int3c2e_ipip1_kernel<4,0,1><<>>(*envs, *eri, *offsets); break; - case 410: GINTfill_int3c2e_ipip1_kernel<4,1,0><<>>(*envs, *eri, *offsets); break; - //case 500: GINTfill_int3c2e_ipip1_kernel<5,0,0><<>>(*envs, *eri, *offsets); break; + //case 5: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,0,5>) break; + case 14: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,1,4>) break; + case 23: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,2,3>) break; + case 32: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,3,2>) break; + case 41: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,4,1>) break; + //case 50: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<0,5,0>) break; + case 104: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,0,4>) break; + case 113: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,1,3>) break; + case 122: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,2,2>) break; + case 131: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,3,1>) break; + case 140: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<1,4,0>) break; + case 203: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,0,3>) break; + case 212: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,1,2>) break; + case 221: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,2,1>) break; + case 230: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<2,3,0>) break; + case 302: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<3,0,2>) break; + case 311: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<3,1,1>) break; + case 320: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<3,2,0>) break; + case 401: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<4,0,1>) break; + case 410: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<4,1,0>) break; + //case 500: LAUNCH_KERNEL(GINTfill_int3c2e_ipip1_kernel<5,0,0>) break; #ifdef UNROLL_INT3C2E #endif default: { - dim3 threads(THREADSX*THREADSY); - dim3 blocks(ntasks_ij, ntasks_kl); const int li_ceil = li + 2; const int gsize = 3*nrys_roots*(li_ceil+1)*(lj+1)*(lk+1); - cudaError_t err = cudaFuncSetAttribute( +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADSX*THREADSY); + sycl::range<2> blocks(ntasks_kl, ntasks_ij); + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(gsize+16), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GINTfill_int3c2e_ipip1_general_kernel(dev_envs, dev_eri, dev_offsets, item, + GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); }); +#else + dim3 threads(THREADSX*THREADSY); + dim3 blocks(ntasks_ij, ntasks_kl); + cudaError_t ferr = cudaFuncSetAttribute( GINTfill_int3c2e_ipip1_general_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (gsize+16)*sizeof(double)); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ipip1_kernel: %s\n", cudaGetErrorString(err)); + if (ferr != cudaSuccess) { + fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ipip1_kernel: %s\n", cudaGetErrorString(ferr)); return 1; } const int shm_size = gsize*sizeof(double); GINTfill_int3c2e_ipip1_general_kernel<<>>(*envs, *eri, *offsets); +#endif } } +#ifndef USE_SYCL cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ipip1_kernel: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -156,7 +194,11 @@ int GINTfill_int3c2e_ipip1(cudaStream_t stream, BasisProdCache *bpcache, double //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif ERITensor eritensor; eritensor.stride_j = strides[1]; diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip2.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip2.cu index 1b5a688b6..93e0d45dc 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip2.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipip2.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "config.h" #include "cuda_alloc.cuh" @@ -31,6 +32,42 @@ #include "gout3c2e.cu" #include "g3c2e_ipip2.cu" +// Abstracts 2D thread/block config (THREADSX/Y swapped between SYCL and CUDA). +// Used 1x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + dim3 threads(THREADSX, THREADSY); \ + dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. dev_envs/dev_eri/dev_offsets are value copies +// hoisted unconditionally so both branches use identical argument names. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function (with template args if needed) +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel-id (with any template args) is the trailing __VA_ARGS__ so its commas +// survive macro expansion. SYCL kernel name is generated inline per source line. +// dev_envs/dev_eri/dev_offsets are on-host value copies made just before launch +// for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + __host__ static int GINTfill_int3c2e_ipip2(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { @@ -38,8 +75,7 @@ static int GINTfill_int3c2e_ipip2(ERITensor *eri, BasisProdOffsets *offsets, GIN int ntasks_ij = offsets->ntasks_ij; int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); - dim3 threads(THREADSX, THREADSY); - dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); int li = envs->i_l; int lj = envs->j_l; int lk = envs->k_l; @@ -47,74 +83,85 @@ static int GINTfill_int3c2e_ipip2(ERITensor *eri, BasisProdOffsets *offsets, GIN switch (type_ijk) { // li+lj+lk=0 - case 0: GINTfill_int3c2e_ipip2_kernel000<<>>(*envs, *eri, *offsets); break; + case 0: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel000) break; // li+lj+lk=1 - case 1: GINTfill_int3c2e_ipip2_kernel<0,0,1><<>>(*envs, *eri, *offsets); break; - case 10: GINTfill_int3c2e_ipip2_kernel<0,1,0><<>>(*envs, *eri, *offsets); break; - case 100: GINTfill_int3c2e_ipip2_kernel<1,0,0><<>>(*envs, *eri, *offsets); break; + case 1: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,0,1>) break; + case 10: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,1,0>) break; + case 100: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,0,0>) break; // li+lj+lk=2 - case 2: GINTfill_int3c2e_ipip2_kernel<0,0,2><<>>(*envs, *eri, *offsets); break; - case 11: GINTfill_int3c2e_ipip2_kernel<0,1,1><<>>(*envs, *eri, *offsets); break; - case 20: GINTfill_int3c2e_ipip2_kernel<0,2,0><<>>(*envs, *eri, *offsets); break; - case 101: GINTfill_int3c2e_ipip2_kernel<1,0,1><<>>(*envs, *eri, *offsets); break; - case 110: GINTfill_int3c2e_ipip2_kernel<1,1,0><<>>(*envs, *eri, *offsets); break; - case 200: GINTfill_int3c2e_ipip2_kernel<2,0,0><<>>(*envs, *eri, *offsets); break; + case 2: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,0,2>) break; + case 11: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,1,1>) break; + case 20: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,2,0>) break; + case 101: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,0,1>) break; + case 110: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,1,0>) break; + case 200: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,0,0>) break; // li+lj+lk=3 - case 3: GINTfill_int3c2e_ipip2_kernel<0,0,3><<>>(*envs, *eri, *offsets); break; - case 12: GINTfill_int3c2e_ipip2_kernel<0,1,2><<>>(*envs, *eri, *offsets); break; - case 21: GINTfill_int3c2e_ipip2_kernel<0,2,1><<>>(*envs, *eri, *offsets); break; - case 30: GINTfill_int3c2e_ipip2_kernel<0,3,0><<>>(*envs, *eri, *offsets); break; - case 102: GINTfill_int3c2e_ipip2_kernel<1,0,2><<>>(*envs, *eri, *offsets); break; - case 111: GINTfill_int3c2e_ipip2_kernel<1,1,1><<>>(*envs, *eri, *offsets); break; - case 120: GINTfill_int3c2e_ipip2_kernel<1,2,0><<>>(*envs, *eri, *offsets); break; - case 201: GINTfill_int3c2e_ipip2_kernel<2,0,1><<>>(*envs, *eri, *offsets); break; - case 210: GINTfill_int3c2e_ipip2_kernel<2,1,0><<>>(*envs, *eri, *offsets); break; - case 300: GINTfill_int3c2e_ipip2_kernel<3,0,0><<>>(*envs, *eri, *offsets); break; + case 3: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,0,3>) break; + case 12: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,1,2>) break; + case 21: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,2,1>) break; + case 30: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,3,0>) break; + case 102: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,0,2>) break; + case 111: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,1,1>) break; + case 120: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,2,0>) break; + case 201: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,0,1>) break; + case 210: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,1,0>) break; + case 300: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<3,0,0>) break; // li+lj+lk=4 - case 4: GINTfill_int3c2e_ipip2_kernel<0,0,4><<>>(*envs, *eri, *offsets); break; - case 13: GINTfill_int3c2e_ipip2_kernel<0,1,3><<>>(*envs, *eri, *offsets); break; - case 22: GINTfill_int3c2e_ipip2_kernel<0,2,2><<>>(*envs, *eri, *offsets); break; - case 31: GINTfill_int3c2e_ipip2_kernel<0,3,1><<>>(*envs, *eri, *offsets); break; - case 40: GINTfill_int3c2e_ipip2_kernel<0,4,0><<>>(*envs, *eri, *offsets); break; - case 103: GINTfill_int3c2e_ipip2_kernel<1,0,3><<>>(*envs, *eri, *offsets); break; - case 112: GINTfill_int3c2e_ipip2_kernel<1,1,2><<>>(*envs, *eri, *offsets); break; - case 121: GINTfill_int3c2e_ipip2_kernel<1,2,1><<>>(*envs, *eri, *offsets); break; - case 130: GINTfill_int3c2e_ipip2_kernel<1,3,0><<>>(*envs, *eri, *offsets); break; - case 202: GINTfill_int3c2e_ipip2_kernel<2,0,2><<>>(*envs, *eri, *offsets); break; - case 211: GINTfill_int3c2e_ipip2_kernel<2,1,1><<>>(*envs, *eri, *offsets); break; - case 220: GINTfill_int3c2e_ipip2_kernel<2,2,0><<>>(*envs, *eri, *offsets); break; - case 301: GINTfill_int3c2e_ipip2_kernel<3,0,1><<>>(*envs, *eri, *offsets); break; - case 310: GINTfill_int3c2e_ipip2_kernel<3,1,0><<>>(*envs, *eri, *offsets); break; - case 400: GINTfill_int3c2e_ipip2_kernel<4,0,0><<>>(*envs, *eri, *offsets); break; + case 4: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,0,4>) break; + case 13: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,1,3>) break; + case 22: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,2,2>) break; + case 31: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,3,1>) break; + case 40: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,4,0>) break; + case 103: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,0,3>) break; + case 112: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,1,2>) break; + case 121: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,2,1>) break; + case 130: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,3,0>) break; + case 202: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,0,2>) break; + case 211: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,1,1>) break; + case 220: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,2,0>) break; + case 301: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<3,0,1>) break; + case 310: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<3,1,0>) break; + case 400: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<4,0,0>) break; // li+lj+lk=5 - //case 5: GINTfill_int3c2e_ipip2_kernel<0,0,5><<>>(*envs, *eri, *offsets); break; - case 14: GINTfill_int3c2e_ipip2_kernel<0,1,4><<>>(*envs, *eri, *offsets); break; - case 23: GINTfill_int3c2e_ipip2_kernel<0,2,3><<>>(*envs, *eri, *offsets); break; - case 32: GINTfill_int3c2e_ipip2_kernel<0,3,2><<>>(*envs, *eri, *offsets); break; - case 41: GINTfill_int3c2e_ipip2_kernel<0,4,1><<>>(*envs, *eri, *offsets); break; - //case 50: GINTfill_int3c2e_ipip2_kernel<0,5,0><<>>(*envs, *eri, *offsets); break; - case 104: GINTfill_int3c2e_ipip2_kernel<1,0,4><<>>(*envs, *eri, *offsets); break; - case 113: GINTfill_int3c2e_ipip2_kernel<1,1,3><<>>(*envs, *eri, *offsets); break; - case 122: GINTfill_int3c2e_ipip2_kernel<1,2,2><<>>(*envs, *eri, *offsets); break; - case 131: GINTfill_int3c2e_ipip2_kernel<1,3,1><<>>(*envs, *eri, *offsets); break; - case 140: GINTfill_int3c2e_ipip2_kernel<1,4,0><<>>(*envs, *eri, *offsets); break; - case 203: GINTfill_int3c2e_ipip2_kernel<2,0,3><<>>(*envs, *eri, *offsets); break; - case 212: GINTfill_int3c2e_ipip2_kernel<2,1,2><<>>(*envs, *eri, *offsets); break; - case 221: GINTfill_int3c2e_ipip2_kernel<2,2,1><<>>(*envs, *eri, *offsets); break; - case 230: GINTfill_int3c2e_ipip2_kernel<2,3,0><<>>(*envs, *eri, *offsets); break; - case 302: GINTfill_int3c2e_ipip2_kernel<3,0,2><<>>(*envs, *eri, *offsets); break; - case 311: GINTfill_int3c2e_ipip2_kernel<3,1,1><<>>(*envs, *eri, *offsets); break; - case 320: GINTfill_int3c2e_ipip2_kernel<3,2,0><<>>(*envs, *eri, *offsets); break; - case 401: GINTfill_int3c2e_ipip2_kernel<4,0,1><<>>(*envs, *eri, *offsets); break; - case 410: GINTfill_int3c2e_ipip2_kernel<4,1,0><<>>(*envs, *eri, *offsets); break; - //case 500: GINTfill_int3c2e_ipip2_kernel<5,0,0><<>>(*envs, *eri, *offsets); break; + //case 5: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,0,5>) break; + case 14: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,1,4>) break; + case 23: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,2,3>) break; + case 32: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,3,2>) break; + case 41: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,4,1>) break; + //case 50: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<0,5,0>) break; + case 104: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,0,4>) break; + case 113: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,1,3>) break; + case 122: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,2,2>) break; + case 131: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,3,1>) break; + case 140: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<1,4,0>) break; + case 203: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,0,3>) break; + case 212: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,1,2>) break; + case 221: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,2,1>) break; + case 230: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<2,3,0>) break; + case 302: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<3,0,2>) break; + case 311: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<3,1,1>) break; + case 320: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<3,2,0>) break; + case 401: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<4,0,1>) break; + case 410: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<4,1,0>) break; + //case 500: LAUNCH_KERNEL(GINTfill_int3c2e_ipip2_kernel<5,0,0>) break; #ifdef UNROLL_INT3C2E #endif default: { - dim3 threads(THREADSX*THREADSY); - dim3 blocks(ntasks_ij, ntasks_kl); const int lk_ceil = lk + 2; const int gsize = 3*nrys_roots*(li+1)*(lj+1)*(lk_ceil+1); + #ifdef USE_SYCL + sycl::range<2> threads(1, THREADSX*THREADSY); + sycl::range<2> blocks(ntasks_kl, ntasks_ij); + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(gsize+16), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GINTfill_int3c2e_ipip2_general_kernel(dev_envs, dev_eri, dev_offsets, item, + GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); }); + #else + dim3 threads(THREADSX*THREADSY); + dim3 blocks(ntasks_ij, ntasks_kl); cudaError_t err = cudaFuncSetAttribute( GINTfill_int3c2e_ipip2_general_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -125,14 +172,17 @@ static int GINTfill_int3c2e_ipip2(ERITensor *eri, BasisProdOffsets *offsets, GIN } const int shm_size = gsize*sizeof(double); GINTfill_int3c2e_ipip2_general_kernel<<>>(*envs, *eri, *offsets); + #endif } } + #ifndef USE_SYCL cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ipip2_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -158,7 +208,11 @@ int GINTfill_int3c2e_ipip2(cudaStream_t stream, BasisProdCache *bpcache, double //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif ERITensor eritensor; eritensor.stride_j = strides[1]; @@ -207,3 +261,6 @@ int GINTfill_int3c2e_ipip2(cudaStream_t stream, BasisProdCache *bpcache, double } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipvip1.cu b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipvip1.cu index 7bbccdb0a..9f4124258 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipvip1.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_int3c2e_ipvip1.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "config.h" #include "cuda_alloc.cuh" @@ -31,6 +32,40 @@ #include "gout3c2e.cu" #include "g3c2e_ipvip1.cu" +// Abstracts 2D thread/block config (THREADSX/Y swapped between SYCL and CUDA). +// Used 1x in this file. +#ifdef USE_SYCL +#define LAUNCH_CONFIG() \ + sycl::range<2> threads(THREADSY, THREADSX); \ + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); +#else +#define LAUNCH_CONFIG() \ + dim3 threads(THREADSX, THREADSY); \ + dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); +#endif + +// Abstracts 2D kernel launch syntax. On SYCL, dev_envs/dev_eri/dev_offsets are +// on-host value copies (needed for lambda capture); on CUDA they alias *envs/ +// *eri/*offsets directly (passed by value at launch, like master). +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel-id (with any template args) is the trailing __VA_ARGS__ so its commas +// survive macro expansion. SYCL kernel name is generated inline per source line. +// dev_envs/dev_eri/dev_offsets are on-host value copies made just before launch +// for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + __host__ static int GINTfill_int3c2e_ipvip1(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { @@ -38,8 +73,7 @@ static int GINTfill_int3c2e_ipvip1(ERITensor *eri, BasisProdOffsets *offsets, GI int ntasks_ij = offsets->ntasks_ij; int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); - dim3 threads(THREADSX, THREADSY); - dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + LAUNCH_CONFIG(); int li = envs->i_l; int lj = envs->j_l; int lk = envs->k_l; @@ -47,75 +81,86 @@ static int GINTfill_int3c2e_ipvip1(ERITensor *eri, BasisProdOffsets *offsets, GI switch (type_ijk) { // li+lj+lk=0 - case 0: GINTfill_int3c2e_ipvip1_kernel000<<>>(*envs, *eri, *offsets); break; + case 0: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel000) break; // li+lj+lk=1 - case 1: GINTfill_int3c2e_ipvip1_kernel<0,0,1><<>>(*envs, *eri, *offsets); break; - case 10: GINTfill_int3c2e_ipvip1_kernel<0,1,0><<>>(*envs, *eri, *offsets); break; - case 100: GINTfill_int3c2e_ipvip1_kernel<1,0,0><<>>(*envs, *eri, *offsets); break; + case 1: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,0,1>) break; + case 10: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,1,0>) break; + case 100: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,0,0>) break; // li+lj+lk=2 - case 2: GINTfill_int3c2e_ipvip1_kernel<0,0,2><<>>(*envs, *eri, *offsets); break; - case 11: GINTfill_int3c2e_ipvip1_kernel<0,1,1><<>>(*envs, *eri, *offsets); break; - case 20: GINTfill_int3c2e_ipvip1_kernel<0,2,0><<>>(*envs, *eri, *offsets); break; - case 101: GINTfill_int3c2e_ipvip1_kernel<1,0,1><<>>(*envs, *eri, *offsets); break; - case 110: GINTfill_int3c2e_ipvip1_kernel<1,1,0><<>>(*envs, *eri, *offsets); break; - case 200: GINTfill_int3c2e_ipvip1_kernel<2,0,0><<>>(*envs, *eri, *offsets); break; + case 2: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,0,2>) break; + case 11: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,1,1>) break; + case 20: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,2,0>) break; + case 101: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,0,1>) break; + case 110: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,1,0>) break; + case 200: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,0,0>) break; // li+lj+lk=3 - case 3: GINTfill_int3c2e_ipvip1_kernel<0,0,3><<>>(*envs, *eri, *offsets); break; - case 12: GINTfill_int3c2e_ipvip1_kernel<0,1,2><<>>(*envs, *eri, *offsets); break; - case 21: GINTfill_int3c2e_ipvip1_kernel<0,2,1><<>>(*envs, *eri, *offsets); break; - case 30: GINTfill_int3c2e_ipvip1_kernel<0,3,0><<>>(*envs, *eri, *offsets); break; - case 102: GINTfill_int3c2e_ipvip1_kernel<1,0,2><<>>(*envs, *eri, *offsets); break; - case 111: GINTfill_int3c2e_ipvip1_kernel<1,1,1><<>>(*envs, *eri, *offsets); break; - case 120: GINTfill_int3c2e_ipvip1_kernel<1,2,0><<>>(*envs, *eri, *offsets); break; - case 201: GINTfill_int3c2e_ipvip1_kernel<2,0,1><<>>(*envs, *eri, *offsets); break; - case 210: GINTfill_int3c2e_ipvip1_kernel<2,1,0><<>>(*envs, *eri, *offsets); break; - case 300: GINTfill_int3c2e_ipvip1_kernel<3,0,0><<>>(*envs, *eri, *offsets); break; + case 3: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,0,3>) break; + case 12: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,1,2>) break; + case 21: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,2,1>) break; + case 30: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,3,0>) break; + case 102: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,0,2>) break; + case 111: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,1,1>) break; + case 120: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,2,0>) break; + case 201: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,0,1>) break; + case 210: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,1,0>) break; + case 300: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<3,0,0>) break; // li+lj+lk=4 - case 4: GINTfill_int3c2e_ipvip1_kernel<0,0,4><<>>(*envs, *eri, *offsets); break; - case 13: GINTfill_int3c2e_ipvip1_kernel<0,1,3><<>>(*envs, *eri, *offsets); break; - case 22: GINTfill_int3c2e_ipvip1_kernel<0,2,2><<>>(*envs, *eri, *offsets); break; - case 31: GINTfill_int3c2e_ipvip1_kernel<0,3,1><<>>(*envs, *eri, *offsets); break; - case 40: GINTfill_int3c2e_ipvip1_kernel<0,4,0><<>>(*envs, *eri, *offsets); break; - case 103: GINTfill_int3c2e_ipvip1_kernel<1,0,3><<>>(*envs, *eri, *offsets); break; - case 112: GINTfill_int3c2e_ipvip1_kernel<1,1,2><<>>(*envs, *eri, *offsets); break; - case 121: GINTfill_int3c2e_ipvip1_kernel<1,2,1><<>>(*envs, *eri, *offsets); break; - case 130: GINTfill_int3c2e_ipvip1_kernel<1,3,0><<>>(*envs, *eri, *offsets); break; - case 202: GINTfill_int3c2e_ipvip1_kernel<2,0,2><<>>(*envs, *eri, *offsets); break; - case 211: GINTfill_int3c2e_ipvip1_kernel<2,1,1><<>>(*envs, *eri, *offsets); break; - case 220: GINTfill_int3c2e_ipvip1_kernel<2,2,0><<>>(*envs, *eri, *offsets); break; - case 301: GINTfill_int3c2e_ipvip1_kernel<3,0,1><<>>(*envs, *eri, *offsets); break; - case 310: GINTfill_int3c2e_ipvip1_kernel<3,1,0><<>>(*envs, *eri, *offsets); break; - case 400: GINTfill_int3c2e_ipvip1_kernel<4,0,0><<>>(*envs, *eri, *offsets); break; + case 4: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,0,4>) break; + case 13: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,1,3>) break; + case 22: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,2,2>) break; + case 31: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,3,1>) break; + case 40: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,4,0>) break; + case 103: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,0,3>) break; + case 112: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,1,2>) break; + case 121: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,2,1>) break; + case 130: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,3,0>) break; + case 202: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,0,2>) break; + case 211: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,1,1>) break; + case 220: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,2,0>) break; + case 301: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<3,0,1>) break; + case 310: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<3,1,0>) break; + case 400: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<4,0,0>) break; // li+lj+lk=5 - //case 5: GINTfill_int3c2e_ipvip1_kernel<0,0,5><<>>(*envs, *eri, *offsets); break; - case 14: GINTfill_int3c2e_ipvip1_kernel<0,1,4><<>>(*envs, *eri, *offsets); break; - case 23: GINTfill_int3c2e_ipvip1_kernel<0,2,3><<>>(*envs, *eri, *offsets); break; - case 32: GINTfill_int3c2e_ipvip1_kernel<0,3,2><<>>(*envs, *eri, *offsets); break; - case 41: GINTfill_int3c2e_ipvip1_kernel<0,4,1><<>>(*envs, *eri, *offsets); break; - //case 50: GINTfill_int3c2e_ipvip1_kernel<0,5,0><<>>(*envs, *eri, *offsets); break; - case 104: GINTfill_int3c2e_ipvip1_kernel<1,0,4><<>>(*envs, *eri, *offsets); break; - case 113: GINTfill_int3c2e_ipvip1_kernel<1,1,3><<>>(*envs, *eri, *offsets); break; - case 122: GINTfill_int3c2e_ipvip1_kernel<1,2,2><<>>(*envs, *eri, *offsets); break; - case 131: GINTfill_int3c2e_ipvip1_kernel<1,3,1><<>>(*envs, *eri, *offsets); break; - case 140: GINTfill_int3c2e_ipvip1_kernel<1,4,0><<>>(*envs, *eri, *offsets); break; - case 203: GINTfill_int3c2e_ipvip1_kernel<2,0,3><<>>(*envs, *eri, *offsets); break; - case 212: GINTfill_int3c2e_ipvip1_kernel<2,1,2><<>>(*envs, *eri, *offsets); break; - case 221: GINTfill_int3c2e_ipvip1_kernel<2,2,1><<>>(*envs, *eri, *offsets); break; - case 230: GINTfill_int3c2e_ipvip1_kernel<2,3,0><<>>(*envs, *eri, *offsets); break; - case 302: GINTfill_int3c2e_ipvip1_kernel<3,0,2><<>>(*envs, *eri, *offsets); break; - case 311: GINTfill_int3c2e_ipvip1_kernel<3,1,1><<>>(*envs, *eri, *offsets); break; - case 320: GINTfill_int3c2e_ipvip1_kernel<3,2,0><<>>(*envs, *eri, *offsets); break; - case 401: GINTfill_int3c2e_ipvip1_kernel<4,0,1><<>>(*envs, *eri, *offsets); break; - case 410: GINTfill_int3c2e_ipvip1_kernel<4,1,0><<>>(*envs, *eri, *offsets); break; - //case 500: GINTfill_int3c2e_ipvip1_kernel<5,0,0><<>>(*envs, *eri, *offsets); break; + //case 5: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,0,5>) break; + case 14: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,1,4>) break; + case 23: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,2,3>) break; + case 32: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,3,2>) break; + case 41: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,4,1>) break; + //case 50: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<0,5,0>) break; + case 104: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,0,4>) break; + case 113: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,1,3>) break; + case 122: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,2,2>) break; + case 131: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,3,1>) break; + case 140: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<1,4,0>) break; + case 203: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,0,3>) break; + case 212: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,1,2>) break; + case 221: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,2,1>) break; + case 230: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<2,3,0>) break; + case 302: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<3,0,2>) break; + case 311: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<3,1,1>) break; + case 320: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<3,2,0>) break; + case 401: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<4,0,1>) break; + case 410: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<4,1,0>) break; + //case 500: LAUNCH_KERNEL(GINTfill_int3c2e_ipvip1_kernel<5,0,0>) break; #ifdef UNROLL_INT3C2E #endif default: { - dim3 threads(THREADSX*THREADSY); - dim3 blocks(ntasks_ij, ntasks_kl); const int li_ceil = li + 1; const int lj_ceil = lj + 1; const int gsize = 3*nrys_roots*(li_ceil+1)*(lj_ceil+1)*(lk+1); + #ifdef USE_SYCL + sycl::range<2> threads(1, THREADSX*THREADSY); + sycl::range<2> blocks(ntasks_kl, ntasks_ij); + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(gsize+16), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + GINTfill_int3c2e_ipvip1_general_kernel(dev_envs, dev_eri, dev_offsets, item, + GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); }); + #else + dim3 threads(THREADSX*THREADSY); + dim3 blocks(ntasks_ij, ntasks_kl); cudaError_t err = cudaFuncSetAttribute( GINTfill_int3c2e_ipvip1_general_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -126,14 +171,17 @@ static int GINTfill_int3c2e_ipvip1(ERITensor *eri, BasisProdOffsets *offsets, GI } const int shm_size = gsize*sizeof(double); GINTfill_int3c2e_ipvip1_general_kernel<<>>(*envs, *eri, *offsets); + #endif } } + #ifndef USE_SYCL cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GINTfill_int3c2e_ipvip1_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -159,7 +207,11 @@ int GINTfill_int3c2e_ipvip1(cudaStream_t stream, BasisProdCache *bpcache, double //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif ERITensor eritensor; eritensor.stride_j = strides[1]; @@ -208,3 +260,6 @@ int GINTfill_int3c2e_ipvip1(cudaStream_t stream, BasisProdCache *bpcache, double } } + +#undef LAUNCH_CONFIG +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/nr_fill_ao_ints.cu b/gpu4pyscf/lib/gint/nr_fill_ao_ints.cu index d75d99c91..3d8729b6c 100644 --- a/gpu4pyscf/lib/gint/nr_fill_ao_ints.cu +++ b/gpu4pyscf/lib/gint/nr_fill_ao_ints.cu @@ -20,6 +20,7 @@ #include #include + #include "gint.h" #include "config.h" #include "cuda_alloc.cuh" @@ -36,6 +37,30 @@ #include "g2e_root3.cu" #include "g2e_root_n.cu" +// Abstracts 2D kernel launch syntax. dev_envs/dev_eri/dev_offsets are value +// copies hoisted unconditionally so both branches use identical argument names. +// TAG: unique SYCL class name (ignored on CUDA) +// KERNEL: kernel function (with template args if needed) +// ...: kernel arguments +#define GINT_CAT_(a, b) a##b +#define GINT_CAT(a, b) GINT_CAT_(a, b) +#ifdef USE_SYCL +// Kernel-id (with any template args) is the trailing __VA_ARGS__ so its commas +// survive macro expansion. SYCL kernel name is generated inline per source line. +// dev_envs/dev_eri/dev_offsets are on-host value copies made just before launch +// for lambda capture. +#define LAUNCH_KERNEL(...) { \ + auto dev_envs = *envs; auto dev_eri = *eri; auto dev_offsets = *offsets; \ + stream.parallel_for( \ + sycl::nd_range<2>(blocks * threads, threads), \ + [=](auto item) [[intel::kernel_args_restrict]] { \ + __VA_ARGS__(dev_envs, dev_eri, dev_offsets); }); } +#else +// CUDA passes the dereferenced structs by value at launch, like master. +#define LAUNCH_KERNEL(...) \ + __VA_ARGS__ <<>>(*envs, *eri, *offsets); +#endif + __host__ static int GINTfill_int2e_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTEnvVars *envs, cudaStream_t stream) { @@ -44,15 +69,20 @@ static int GINTfill_int2e_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTE int ntasks_kl = offsets->ntasks_kl; assert(ntasks_kl < 65536*THREADSY); int type_ijkl; + #ifdef USE_SYCL + sycl::range<2> threads(THREADSY, THREADSX); + sycl::range<2> blocks((ntasks_kl+THREADSY-1)/THREADSY, (ntasks_ij+THREADSX-1)/THREADSX); + #else dim3 threads(THREADSX, THREADSY); dim3 blocks((ntasks_ij+THREADSX-1)/THREADSX, (ntasks_kl+THREADSY-1)/THREADSY); + #endif switch (nrys_roots) { case 1: type_ijkl = (envs->i_l << 3) | (envs->j_l << 2) | (envs->k_l << 1) | envs->l_l; switch (type_ijkl) { - case 0b0000: GINTfill_int2e_kernel0000<<>>(*envs, *eri, *offsets); break; - case 0b0010: GINTfill_int2e_kernel0010<<>>(*envs, *eri, *offsets); break; - case 0b1000: GINTfill_int2e_kernel1000<<>>(*envs, *eri, *offsets); break; + case 0b0000: LAUNCH_KERNEL(GINTfill_int2e_kernel0000) break; + case 0b0010: LAUNCH_KERNEL(GINTfill_int2e_kernel0010) break; + case 0b1000: LAUNCH_KERNEL(GINTfill_int2e_kernel1000) break; default: //GINTfill_int2e_kernel<1, GOUTSIZE1> <<>>(*envs, *eri, *offsets); break; fprintf(stderr, "roots=1 type_ijkl %d\n", type_ijkl); @@ -61,65 +91,65 @@ static int GINTfill_int2e_tasks(ERITensor *eri, BasisProdOffsets *offsets, GINTE case 2: type_ijkl = (envs->i_l << 6) | (envs->j_l << 4) | (envs->k_l << 2) | envs->l_l; switch (type_ijkl) { - case (0<<6)|(0<<4)|(1<<2)|1: GINTfill_int2e_kernel0011<<>>(*envs, *eri, *offsets); break; - case (0<<6)|(0<<4)|(2<<2)|0: GINTfill_int2e_kernel0020<<>>(*envs, *eri, *offsets); break; - case (0<<6)|(0<<4)|(2<<2)|1: GINTfill_int2e_kernel0021<<>>(*envs, *eri, *offsets); break; - case (0<<6)|(0<<4)|(3<<2)|0: GINTfill_int2e_kernel0030<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(0<<4)|(1<<2)|0: GINTfill_int2e_kernel1010<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(0<<4)|(1<<2)|1: GINTfill_int2e_kernel1011<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(0<<4)|(2<<2)|0: GINTfill_int2e_kernel1020<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(1<<4)|(0<<2)|0: GINTfill_int2e_kernel1100<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(1<<4)|(1<<2)|0: GINTfill_int2e_kernel1110<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(0<<4)|(0<<2)|0: GINTfill_int2e_kernel2000<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(0<<4)|(1<<2)|0: GINTfill_int2e_kernel2010<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(1<<4)|(0<<2)|0: GINTfill_int2e_kernel2100<<>>(*envs, *eri, *offsets); break; - case (3<<6)|(0<<4)|(0<<2)|0: GINTfill_int2e_kernel3000<<>>(*envs, *eri, *offsets); break; - default: - GINTfill_int2e_kernel<2, GOUTSIZE2> <<>>(*envs, *eri, *offsets); break; + case (0<<6)|(0<<4)|(1<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel0011) break; + case (0<<6)|(0<<4)|(2<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel0020) break; + case (0<<6)|(0<<4)|(2<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel0021) break; + case (0<<6)|(0<<4)|(3<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel0030) break; + case (1<<6)|(0<<4)|(1<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel1010) break; + case (1<<6)|(0<<4)|(1<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel1011) break; + case (1<<6)|(0<<4)|(2<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel1020) break; + case (1<<6)|(1<<4)|(0<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel1100) break; + case (1<<6)|(1<<4)|(1<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel1110) break; + case (2<<6)|(0<<4)|(0<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2000) break; + case (2<<6)|(0<<4)|(1<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2010) break; + case (2<<6)|(1<<4)|(0<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2100) break; + case (3<<6)|(0<<4)|(0<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel3000) break; + default: LAUNCH_KERNEL(GINTfill_int2e_kernel<2, GOUTSIZE2>) break; } break; case 3: type_ijkl = (envs->i_l << 6) | (envs->j_l << 4) | (envs->k_l << 2) | envs->l_l; switch (type_ijkl) { - case (0<<6)|(0<<4)|(2<<2)|2: GINTfill_int2e_kernel0022<<>>(*envs, *eri, *offsets); break; - case (0<<6)|(0<<4)|(3<<2)|1: GINTfill_int2e_kernel0031<<>>(*envs, *eri, *offsets); break; - case (0<<6)|(0<<4)|(3<<2)|2: GINTfill_int2e_kernel0032<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(0<<4)|(2<<2)|1: GINTfill_int2e_kernel1021<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(0<<4)|(2<<2)|2: GINTfill_int2e_kernel1022<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(0<<4)|(3<<2)|0: GINTfill_int2e_kernel1030<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(0<<4)|(3<<2)|1: GINTfill_int2e_kernel1031<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(1<<4)|(1<<2)|1: GINTfill_int2e_kernel1111<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(1<<4)|(2<<2)|0: GINTfill_int2e_kernel1120<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(1<<4)|(2<<2)|1: GINTfill_int2e_kernel1121<<>>(*envs, *eri, *offsets); break; - case (1<<6)|(1<<4)|(3<<2)|0: GINTfill_int2e_kernel1130<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(0<<4)|(1<<2)|1: GINTfill_int2e_kernel2011<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(0<<4)|(2<<2)|0: GINTfill_int2e_kernel2020<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(0<<4)|(2<<2)|1: GINTfill_int2e_kernel2021<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(0<<4)|(3<<2)|0: GINTfill_int2e_kernel2030<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(1<<4)|(1<<2)|0: GINTfill_int2e_kernel2110<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(1<<4)|(1<<2)|1: GINTfill_int2e_kernel2111<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(1<<4)|(2<<2)|0: GINTfill_int2e_kernel2120<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(2<<4)|(0<<2)|0: GINTfill_int2e_kernel2200<<>>(*envs, *eri, *offsets); break; - case (2<<6)|(2<<4)|(1<<2)|0: GINTfill_int2e_kernel2210<<>>(*envs, *eri, *offsets); break; - case (3<<6)|(0<<4)|(1<<2)|0: GINTfill_int2e_kernel3010<<>>(*envs, *eri, *offsets); break; - case (3<<6)|(0<<4)|(1<<2)|1: GINTfill_int2e_kernel3011<<>>(*envs, *eri, *offsets); break; - case (3<<6)|(0<<4)|(2<<2)|0: GINTfill_int2e_kernel3020<<>>(*envs, *eri, *offsets); break; - case (3<<6)|(1<<4)|(0<<2)|0: GINTfill_int2e_kernel3100<<>>(*envs, *eri, *offsets); break; - case (3<<6)|(1<<4)|(1<<2)|0: GINTfill_int2e_kernel3110<<>>(*envs, *eri, *offsets); break; - case (3<<6)|(2<<4)|(0<<2)|0: GINTfill_int2e_kernel3200<<>>(*envs, *eri, *offsets); break; - default: - GINTfill_int2e_kernel<3, GOUTSIZE3> <<>>(*envs, *eri, *offsets); break; + case (0<<6)|(0<<4)|(2<<2)|2: LAUNCH_KERNEL(GINTfill_int2e_kernel0022) break; + case (0<<6)|(0<<4)|(3<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel0031) break; + case (0<<6)|(0<<4)|(3<<2)|2: LAUNCH_KERNEL(GINTfill_int2e_kernel0032) break; + case (1<<6)|(0<<4)|(2<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel1021) break; + case (1<<6)|(0<<4)|(2<<2)|2: LAUNCH_KERNEL(GINTfill_int2e_kernel1022) break; + case (1<<6)|(0<<4)|(3<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel1030) break; + case (1<<6)|(0<<4)|(3<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel1031) break; + case (1<<6)|(1<<4)|(1<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel1111) break; + case (1<<6)|(1<<4)|(2<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel1120) break; + case (1<<6)|(1<<4)|(2<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel1121) break; + case (1<<6)|(1<<4)|(3<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel1130) break; + case (2<<6)|(0<<4)|(1<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel2011) break; + case (2<<6)|(0<<4)|(2<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2020) break; + case (2<<6)|(0<<4)|(2<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel2021) break; + case (2<<6)|(0<<4)|(3<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2030) break; + case (2<<6)|(1<<4)|(1<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2110) break; + case (2<<6)|(1<<4)|(1<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel2111) break; + case (2<<6)|(1<<4)|(2<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2120) break; + case (2<<6)|(2<<4)|(0<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2200) break; + case (2<<6)|(2<<4)|(1<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel2210) break; + case (3<<6)|(0<<4)|(1<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel3010) break; + case (3<<6)|(0<<4)|(1<<2)|1: LAUNCH_KERNEL(GINTfill_int2e_kernel3011) break; + case (3<<6)|(0<<4)|(2<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel3020) break; + case (3<<6)|(1<<4)|(0<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel3100) break; + case (3<<6)|(1<<4)|(1<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel3110) break; + case (3<<6)|(2<<4)|(0<<2)|0: LAUNCH_KERNEL(GINTfill_int2e_kernel3200) break; + default: LAUNCH_KERNEL(GINTfill_int2e_kernel<3, GOUTSIZE3>) break; } break; - case 4: GINTfill_int2e_kernel<4, GOUTSIZE4> <<>>(*envs, *eri, *offsets); break; - case 5: GINTfill_int2e_kernel<5, GOUTSIZE5> <<>>(*envs, *eri, *offsets); break; - case 6: GINTfill_int2e_kernel<6, GOUTSIZE6> <<>>(*envs, *eri, *offsets); break; - case 7: GINTfill_int2e_kernel<7, GOUTSIZE7> <<>>(*envs, *eri, *offsets); break; - case 8: GINTfill_int2e_kernel<8, GOUTSIZE8> <<>>(*envs, *eri, *offsets); break; + + case 4: LAUNCH_KERNEL(GINTfill_int2e_kernel<4, GOUTSIZE4>) break; + case 5: LAUNCH_KERNEL(GINTfill_int2e_kernel<5, GOUTSIZE5>) break; + case 6: LAUNCH_KERNEL(GINTfill_int2e_kernel<6, GOUTSIZE6>) break; + case 7: LAUNCH_KERNEL(GINTfill_int2e_kernel<7, GOUTSIZE7>) break; + case 8: LAUNCH_KERNEL(GINTfill_int2e_kernel<8, GOUTSIZE8>) break; default: fprintf(stderr, "rys roots %d\n", nrys_roots); return 1; } + cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error of GINTfill_int2e_kernel: %s\n", cudaGetErrorString(err)); @@ -149,7 +179,11 @@ int GINTfill_int2e(cudaStream_t stream, BasisProdCache *bpcache, double *eri, in //checkCudaErrors(cudaMemcpyToSymbol(c_envs, &envs, sizeof(GINTEnvVars))); // move bpcache to constant memory + #ifdef USE_SYCL + stream.memcpy(s_bpcache, bpcache, sizeof(BasisProdCache)).wait(); + #else checkCudaErrors(cudaMemcpyToSymbol(c_bpcache, bpcache, sizeof(BasisProdCache))); + #endif ERITensor eritensor; eritensor.stride_j = strides[1]; eritensor.stride_k = strides[2]; @@ -206,3 +240,5 @@ int GINTfill_int2e(cudaStream_t stream, BasisProdCache *bpcache, double *eri, in return 0; } } + +#undef LAUNCH_KERNEL diff --git a/gpu4pyscf/lib/gint/reduction.cu b/gpu4pyscf/lib/gint/reduction.cu index 8f047cd7c..e8812a7c2 100644 --- a/gpu4pyscf/lib/gint/reduction.cu +++ b/gpu4pyscf/lib/gint/reduction.cu @@ -18,6 +18,10 @@ template __device__ static void block_reduce_x(double val, double *addr, int tx, int ty){ + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + double (&sdata)[blockx*blocky] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + #else __shared__ double sdata[blockx*blocky]; sdata[tx*blocky+ty] = val; __syncthreads(); if (blockx >= 32) if (tx < 16) sdata[tx*blocky+ty] += sdata[(tx+16)*blocky+ty]; __syncthreads(); @@ -26,6 +30,7 @@ __device__ static void block_reduce_x(double val, double *addr, int tx, int ty){ if (blockx >= 4) if (tx < 2) sdata[tx*blocky+ty] += sdata[(tx+2)*blocky+ty]; __syncthreads(); if (blockx >= 2) if (tx < 1) sdata[tx*blocky+ty] += sdata[(tx+1)*blocky+ty]; __syncthreads(); if (tx == 0) atomicAdd(addr, sdata[ty]); + #endif } template @@ -38,7 +43,12 @@ __device__ static void block_reduce_y(double val, double *addr, int tx, int ty){ if(blocky >= 2) sdata[tx*blocky+ty] += sdata[tx*blocky+ty+1]; */ int stride = blocky + 1; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + double (&sdata)[blockx*(blocky+1)] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + #else __shared__ double sdata[blockx*(blocky+1)]; + #endif sdata[tx*stride+ty] = val; __syncthreads(); if (blocky >= 32) if (ty < 16) sdata[tx*stride+ty] += sdata[tx*stride+ty+16]; __syncthreads(); if (blocky >= 16) if (ty < 8) sdata[tx*stride+ty] += sdata[tx*stride+ty+8]; __syncthreads(); @@ -48,11 +58,18 @@ __device__ static void block_reduce_y(double val, double *addr, int tx, int ty){ if (ty == 0) atomicAdd(addr, sdata[tx*stride]); } -template +template __device__ void block_reduce(double *sum, double a){ - const int tx = threadIdx.x; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int tx = item.get_local_id(1); + __syncthreads(); + double (&as)[BLKSIZE] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + #else __syncthreads(); + const int tx = threadIdx.x; __shared__ double as[BLKSIZE]; + #endif as[tx] = a; __syncthreads(); diff --git a/gpu4pyscf/lib/gint/rys_roots.cu b/gpu4pyscf/lib/gint/rys_roots.cu index 9b2b2b21a..92f2fe9ec 100644 --- a/gpu4pyscf/lib/gint/rys_roots.cu +++ b/gpu4pyscf/lib/gint/rys_roots.cu @@ -47,10 +47,10 @@ static void GINTrys_root(double x, double *rw) for (int rt_id = 0; rt_id < NROOTS; ++rt_id) { const int it = (int)(x * .4); - double *datax = ROOT_RW_DATA + DEGREE1*INTERVALS * NROOTS*(NROOTS-1); + const double *datax = ROOT_RW_DATA + DEGREE1*INTERVALS * NROOTS*(NROOTS-1); const double u = (x - it * 2.5) * 0.8 - 1.; const double u2 = u * 2.; - double *c = datax + (2*rt_id) * DEGREE1 * INTERVALS; + const double *c = datax + (2*rt_id) * DEGREE1 * INTERVALS; //for i in range(2, degree + 1): // c0, c1 = c[degree-i] - c1, c0 + c1*u2 double c0 = c[it + DEGREE *INTERVALS]; @@ -123,10 +123,16 @@ inline void GINTscale_u(double *u, double theta) __device__ static void GINTrys_root(int nroots, double x, double *rw) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int threadIdx_x = item.get_local_id(1); +#else + const int threadIdx_x = threadIdx.x; +#endif // roots and weights are distributed in each thread const int off = nroots * (nroots - 1) / 2; const double t = sqrt(PIE4/x); - const int rt_id = threadIdx.x % nroots; + const int rt_id = threadIdx_x % nroots; if (x<3.0e-7){ const double r = ROOT_SMALLX_R0[off+rt_id] + ROOT_SMALLX_R1[off+rt_id] * x; const double w = ROOT_SMALLX_W0[off+rt_id] + ROOT_SMALLX_W1[off+rt_id] * x; @@ -144,10 +150,10 @@ static void GINTrys_root(int nroots, double x, double *rw) } const int it = (int)(x * .4); - double *datax = ROOT_RW_DATA + DEGREE1*INTERVALS * nroots*(nroots-1); + const double *datax = ROOT_RW_DATA + DEGREE1*INTERVALS * nroots*(nroots-1); const double u = (x - it * 2.5) * 0.8 - 1.; const double u2 = u * 2.; - double *c = datax + (2*rt_id) * DEGREE1 * INTERVALS; + const double *c = datax + (2*rt_id) * DEGREE1 * INTERVALS; //for i in range(2, degree + 1): // c0, c1 = c[degree-i] - c1, c0 + c1*u2 double c0 = c[it + DEGREE *INTERVALS]; diff --git a/gpu4pyscf/lib/gsycl/CMakeLists.txt b/gpu4pyscf/lib/gsycl/CMakeLists.txt new file mode 100644 index 000000000..5e4edb7a4 --- /dev/null +++ b/gpu4pyscf/lib/gsycl/CMakeLists.txt @@ -0,0 +1,44 @@ +# Copyright 2021-2024 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# INTERFACE target — carries the include path to any consumer + +# ── gsycl: SYCL compatibility layer ──────────────────────────────────── +# +# Provides two things: +# 1. sycl_compat — INTERFACE target that injects gsycl/ at the front +# of the include path, making and +# resolve to the SYCL shim headers instead +# of the real CUDA toolkit headers. +# 2. gsycl — shared library exposing the Python ABI (extern "C" +# device/queue management symbols) via libgsycl.so. +# All other .so files link against this to resolve +# sycl_get_queue_ptr, sycl_set_queue_ptr, etc. + +# 1. Header-only INTERFACE target — consumed by every other subdirectory +add_library(sycl_compat INTERFACE) +target_include_directories(sycl_compat INTERFACE + "$" + "$" +) + +# 2. Python ABI shared library — ONE translation unit owns all queue state +add_library(gsycl SHARED sycl_api_python.cpp) +target_link_libraries(gsycl PRIVATE sycl_compat) +set_target_properties(gsycl PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + +# 3. Namespaced alias — lets consumers write gsycl::gsycl instead of gsycl, +# making it immediately clear this is an external dependency target +add_library(gsycl::gsycl ALIAS gsycl) diff --git a/gpu4pyscf/lib/gsycl/cub/cub.cuh b/gpu4pyscf/lib/gsycl/cub/cub.cuh new file mode 100644 index 000000000..c10905294 --- /dev/null +++ b/gpu4pyscf/lib/gsycl/cub/cub.cuh @@ -0,0 +1,3 @@ +#pragma once +// Drop-in replacement for when building with SYCL. +#include "sycl_device.hpp" diff --git a/gpu4pyscf/lib/gsycl/cuda.h b/gpu4pyscf/lib/gsycl/cuda.h new file mode 100644 index 000000000..c10905294 --- /dev/null +++ b/gpu4pyscf/lib/gsycl/cuda.h @@ -0,0 +1,3 @@ +#pragma once +// Drop-in replacement for when building with SYCL. +#include "sycl_device.hpp" diff --git a/gpu4pyscf/lib/gsycl/cuda_runtime.h b/gpu4pyscf/lib/gsycl/cuda_runtime.h new file mode 100644 index 000000000..ae8798520 --- /dev/null +++ b/gpu4pyscf/lib/gsycl/cuda_runtime.h @@ -0,0 +1,5 @@ +#pragma once +// Drop-in replacement for when building with SYCL. +// Source files include unconditionally — this file +// is found first on the include path when USE_SYCL is active. +#include "sycl_device.hpp" diff --git a/gpu4pyscf/lib/gsycl/gpu_compat.h b/gpu4pyscf/lib/gsycl/gpu_compat.h new file mode 100644 index 000000000..02864bac4 --- /dev/null +++ b/gpu4pyscf/lib/gsycl/gpu_compat.h @@ -0,0 +1,105 @@ +/* + * Copyright 2021-2024 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Backend definitions and abstractions + +#pragma once + +#include + + +#ifdef USE_SYCL + +#define setup_context() \ + const auto item = syclex::this_work_item::get_nd_item<3>() + +#define threadIdx_x item.get_local_id(2) +#define threadIdx_y item.get_local_id(1) +#define threadIdx_z item.get_local_id(0) + +#define blockIdx_x item.get_group(2) +#define blockIdx_y item.get_group(1) +#define blockIdx_z item.get_group(0) + +#define blockDim_x item.get_local_range(2) +#define blockDim_y item.get_local_range(1) +#define blockDim_z item.get_local_range(0) + +#define gridDim_x item.get_group_range(2) +#define gridDim_y item.get_group_range(1) +#define gridDim_z item.get_group_range(0) + +#else + +#define setup_context() + +#define threadIdx_x threadIdx.x +#define threadIdx_y threadIdx.y +#define threadIdx_z threadIdx.z + +#define blockIdx_x blockIdx.x +#define blockIdx_y blockIdx.y +#define blockIdx_z blockIdx.z + +#define blockDim_x blockDim.x +#define blockDim_y blockDim.y +#define blockDim_z blockDim.z + +#define gridDim_x gridDim.x +#define gridDim_y gridDim.y +#define gridDim_z gridDim.z + +#endif + + +#ifdef USE_SYCL +inline sycl::range<3> make_grid( + size_t x, size_t y = 1, size_t z = 1) +{ + return sycl::range<3>(z, y, x); +} + +inline sycl::range<3> make_block( + size_t x, size_t y = 1, size_t z = 1) +{ + return sycl::range<3>(z, y, x); +} + +#define LAUNCH_KERNEL(kernel, grid, block, shm_size, stream, ...) \ + { \ + (stream).parallel_for( \ + sycl::nd_range<3>(grid * block, block), \ + [=](sycl::nd_item<3>) { kernel(__VA_ARGS__); }); \ + } + +#else +inline dim3 make_grid( + unsigned int x, unsigned int y = 1, unsigned int z = 1) +{ + return dim3(x, y, z); +} + +inline dim3 make_block( + unsigned int x, unsigned int y = 1, unsigned int z = 1) +{ + return dim3(x, y, z); +} + +#define LAUNCH_KERNEL(kernel, grid, block, shm_size, stream, ...) \ + { \ + kernel<<>>(__VA_ARGS__); \ + } +#endif diff --git a/gpu4pyscf/lib/gsycl/sycl_api_python.cpp b/gpu4pyscf/lib/gsycl/sycl_api_python.cpp new file mode 100644 index 000000000..ee61ca10c --- /dev/null +++ b/gpu4pyscf/lib/gsycl/sycl_api_python.cpp @@ -0,0 +1,163 @@ +#include +#include "sycl_device.hpp" + +#define GPU4PYSCF_EXPORT __attribute__((visibility("default"))) + +// ===================================================================== +// Process-global queue storage. +// +// This is the ONE place queues live. All .so files that include +// sycl_device.hpp call the extern "C" functions below (resolved at +// link/load time to THIS translation unit in libgsycl.so). +// +// Python creates in-order dpctl.SyclQueue objects and pushes them +// here via sycl_set_queue_ptr(). No queues are created in C++. +// ===================================================================== + +static std::mutex g_mutex; +static std::vector g_queues; // one per GPU, set by Python +static std::atomic g_current_device{0}; // process-global + +// ===================================================================== +// Event storage (for sycl_record_event / sycl_wait_event) +// ===================================================================== + +static std::unordered_map> g_event_map; +static std::mutex g_event_mutex; + +// ===================================================================== +// Exported C API — called by ALL .so files through the header +// ===================================================================== + +extern "C" { + +GPU4PYSCF_EXPORT int sycl_get_device_id() { + return g_current_device.load(std::memory_order_relaxed); +} + +GPU4PYSCF_EXPORT void* sycl_get_queue_ptr() { + int id = g_current_device.load(std::memory_order_relaxed); + std::lock_guard lock(g_mutex); + if (id < 0 || id >= static_cast(g_queues.size()) || + g_queues[id] == nullptr) { + throw std::runtime_error( + "sycl_get_queue_ptr: queue for device " + std::to_string(id) + + " not set. Ensure 'import cupy' runs before any GPU op."); + } + return static_cast(g_queues[id]); +} +GPU4PYSCF_EXPORT void* sycl_get_queue_ptr_for(int device_id) { + std::lock_guard lock(g_mutex); + if (device_id < 0 || + device_id >= static_cast(g_queues.size()) || + g_queues[device_id] == nullptr) { + throw std::runtime_error( + "sycl_get_queue_ptr_for: no queue registered for device " + + std::to_string(device_id)); + } + return static_cast(g_queues[device_id]); +} + +GPU4PYSCF_EXPORT void sycl_set_queue_ptr(int device_id, void* queue_ptr) { + if (device_id < 0) + throw std::runtime_error("sycl_set_queue_ptr: negative device id"); + std::lock_guard lock(g_mutex); + if (static_cast(device_id) >= g_queues.size()) + g_queues.resize(device_id + 1, nullptr); + g_queues[device_id] = static_cast(queue_ptr); +} + +GPU4PYSCF_EXPORT void sycl_set_device(int device_id) { + std::lock_guard lock(g_mutex); + if (device_id < 0 || + device_id >= static_cast(g_queues.size()) || + g_queues[device_id] == nullptr) { + throw std::runtime_error( + "sycl_set_device: no queue registered for device " + + std::to_string(device_id)); + } + g_current_device.store(device_id, std::memory_order_relaxed); +} + +// GPU4PYSCF_EXPORT void sycl_queue_synchronize(void* queue_ptr) { +// auto* q = static_cast(queue_ptr); +// q->wait(); +// } + +// GPU4PYSCF_EXPORT void* sycl_record_event() { +// auto* q = static_cast(sycl_get_queue_ptr()); +// auto ev = std::make_shared(q->ext_oneapi_submit_barrier()); + +// void* handle = static_cast(ev.get()); +// { +// std::lock_guard lock(g_event_mutex); +// g_event_map[handle] = ev; +// } +// return handle; +// } + +// GPU4PYSCF_EXPORT void sycl_wait_event(void* handle) { +// std::shared_ptr ev; +// { +// std::lock_guard lock(g_event_mutex); +// auto it = g_event_map.find(handle); +// if (it != g_event_map.end()) { +// ev = it->second; +// g_event_map.erase(it); +// } +// } +// if (ev) { +// try { +// ev->wait(); +// } catch (...) { +// // Swallow — may happen during shutdown +// } +// } +// } + +GPU4PYSCF_EXPORT size_t sycl_get_total_memory() { + auto* q = static_cast(sycl_get_queue_ptr()); + return q->get_device().get_info(); +} + +GPU4PYSCF_EXPORT size_t sycl_get_shared_memory() { + auto* q = static_cast(sycl_get_queue_ptr()); + return q->get_device().get_info(); +} + +// Maps to CUDA cudaDeviceProp::multiProcessorCount +GPU4PYSCF_EXPORT int sycl_get_compute_units() { + auto* q = static_cast(sycl_get_queue_ptr()); + return static_cast( + q->get_device().get_info()); +} + +// Maps to CUDA cudaDeviceProp::name. Copies the device name into the +// caller-provided buffer (NUL-terminated, truncated to buf_size-1). +GPU4PYSCF_EXPORT void sycl_get_device_name(char* buf, int buf_size) { + if (buf == nullptr || buf_size <= 0) return; + auto* q = static_cast(sycl_get_queue_ptr()); + std::string name = q->get_device().get_info(); + int n = static_cast(name.size()); + if (n > buf_size - 1) n = buf_size - 1; + for (int i = 0; i < n; ++i) buf[i] = name[i]; + buf[n] = '\0'; +} + +GPU4PYSCF_EXPORT size_t sycl_get_free_memory() { + auto* q = static_cast(sycl_get_queue_ptr()); + auto dev = q->get_device(); + if (!dev.has(sycl::aspect::ext_intel_free_memory)) { + return static_cast( + 0.9 * dev.get_info()); + } + return dev.get_info(); +} + +GPU4PYSCF_EXPORT size_t sycl_memcpy(void* dst, void* src, size_t size) { + auto* q = static_cast(sycl_get_queue_ptr()); + q->memcpy(dst, src, size); + return 0; +} + +} // extern "C" diff --git a/gpu4pyscf/lib/gsycl/sycl_device.hpp b/gpu4pyscf/lib/gsycl/sycl_device.hpp new file mode 100644 index 000000000..a9c6c7a1e --- /dev/null +++ b/gpu4pyscf/lib/gsycl/sycl_device.hpp @@ -0,0 +1,222 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#define warpSize (item.get_sub_group().get_max_local_range()[0]) + +using cudaError_t = int; +constexpr int cudaSuccess = 0; +inline unsigned int __activemask() { return 0; } +inline int cudaGetLastError() { return 0; } +inline int cudaPeekAtLastError() { return 0; } +inline void checkCudaErrors(int) { } +inline const char* cudaGetErrorString(int) { return "No error"; } + +#define cudaGetDevice(ptr) (syclGetDevice(ptr)) + +#if defined(__SYCL_DEVICE_ONLY__) +#define printf(...) sycl::ext::oneapi::experimental::printf(__VA_ARGS__) +#endif + +extern "C" { + SYCL_EXTERNAL unsigned __attribute__((overloadable)) intel_get_slice_id(void); + SYCL_EXTERNAL unsigned __attribute__((overloadable)) intel_get_subslice_id(void); + SYCL_EXTERNAL unsigned __attribute__((overloadable)) intel_get_eu_id(void); +} +inline void __trap() { __builtin_trap(); } + +enum cudaFuncAttribute { + cudaFuncAttributeMaxDynamicSharedMemorySize = 0, + cudaFuncAttributePreferredSharedMemoryCarveout = 1, +}; + +enum cudaFuncCache { + cudaFuncCachePreferNone = 0, + cudaFuncCachePreferShared = 1, + cudaFuncCachePreferL1 = 2, + cudaFuncCachePreferEqual = 3 +}; + +#ifndef cudaFuncSetAttribute +#define cudaFuncSetAttribute(...) (cudaSuccess) +#endif + +#ifndef cudaFuncSetCacheConfig +#define cudaFuncSetCacheConfig(...) (cudaSuccess) +#endif + +#define CUDA_VERSION 12040 +#define __maxnreg__(x) + +#define __global__ __attribute__((always_inline)) +#define __device__ __attribute__((always_inline)) +#define __forceinline__ __attribute__((always_inline)) +#define __host__ __attribute__((always_inline)) +#define __constant__ inline constexpr + +using cudaStream_t = sycl::queue&; +namespace syclex = sycl::ext::oneapi; +using double2 = sycl::double2; + +#define rnorm3d(d1,d2,d3) (1 / sycl::length(sycl::double3((d1), (d2), (d3)))) +#define norm3d(d1,d2,d3) (sycl::length(sycl::double3((d1), (d2), (d3)))) +#define __syncthreads() (sycl::group_barrier(item.get_group())) +#define __syncwarps() (sycl::group_barrier(item.get_sub_group())) +#define __threadfence_block() (sycl::atomic_fence(sycl::memory_order::seq_cst, sycl::memory_scope::work_group)) +#define __shfl_down_sync(mask, val, delta) (sycl::shift_group_left((item.get_sub_group()), (val), (delta))) +#define __ballot_sync(mask, predicate) (sycl::reduce_over_group(item.get_sub_group(), (predicate) ? (0x1 << item.get_sub_group().get_local_linear_id()) : 0, sycl::plus<>())) +#define __popc(x) (sycl::popcount(x)) + +template __attribute__((always_inline)) auto ceil(T x) { return sycl::ceil(x); } +template __attribute__((always_inline)) auto sqrtf(T x) { return sycl::sqrt(x); } +template __attribute__((always_inline)) auto sqrt(T x) { return sycl::sqrt(x); } +template __attribute__((always_inline)) auto rsqrt(T x) { return sycl::rsqrt(x); } +template __attribute__((always_inline)) auto min(T x, T y) { return sycl::min(x, y); } +template __attribute__((always_inline)) auto max(T x, T y) { return sycl::max(x, y); } +template __attribute__((always_inline)) auto exp(T x) { return sycl::exp(x); } +template __attribute__((always_inline)) auto expf(T x) { return sycl::exp(x); } +template __attribute__((always_inline)) auto fabs(T x) -> std::enable_if_t, double> { return sycl::fabs(x); } +template __attribute__((always_inline)) auto fabsf(T x) -> std::enable_if_t, float> { return sycl::fabs(x); } +template __attribute__((always_inline)) auto erf(T x) { return sycl::erf(x); } +template __attribute__((always_inline)) auto floor(T x) { return sycl::floor(x); } +template __attribute__((always_inline)) auto pow(T x, int n) { return sycl::pown(x, n); } +template __attribute__((always_inline)) auto pow(T x, U n) { return sycl::pow(x, n); } +template __attribute__((always_inline)) typename std::enable_if::value, float>::type logf(T x) { return sycl::log(x); } +template __attribute__((always_inline)) auto log(T x) { return sycl::log(x); } +template __attribute__((always_inline)) void sincos(T x, T* sptr, T* cptr) { *sptr = sycl::sincos(x, cptr); } +#define NAN std::numeric_limits::quiet_NaN() + +namespace constants { + constexpr double pi = 3.141592653589793238462643383279502884; +} +#ifndef M_PI +#define M_PI constants::pi +#endif + +namespace compat { + struct double3 { + double x, y, z; + constexpr double3() : x(0.0), y(0.0), z(0.0) {} + constexpr double3(double x_, double y_, double z_) : x(x_), y(y_), z(z_) {} + }; +} +using double3 = compat::double3; + +template +static inline T1 +atomicAdd(T1* addr, const T2 val) { + sycl::atomic_ref atom(*(addr)); + return atom.fetch_add(static_cast((val))); +} +template +static inline T1 +atomicMax(T1* addr, const T2 val) { + sycl::atomic_ref atom(*(addr)); + return atom.fetch_max(static_cast((val))); +} +template +static inline typename std::enable_if::value, T>::type +atomicOr(T* addr, const T val) { + sycl::atomic_ref atom(*(addr)); + return atom.fetch_or((val)); +} + +template +using sycl_device_global = sycl::ext::oneapi::experimental::device_global; + +#if defined(__INTEL_LLVM_COMPILER) && __INTEL_LLVM_COMPILER >= 20230200 +#define GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(accessor) \ + accessor.get_multi_ptr().get() +#else +#define GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(accessor) accessor.get_pointer() +#endif + +// ===================================================================== +// Queue management — DECLARATIONS ONLY +// +// Implementation lives in sycl_api_python.cpp, compiled into ONE .so +// (libgsycl.so). All other .so files (libgdft, libgint, ...) link +// against libgsycl.so so the process has exactly ONE sycl::queue* per +// device. +// +// Ownership model: Python owns the queues. On `import cupy`, +// cupy/cuda.py creates one in-order dpctl.SyclQueue per GPU and calls +// sycl_set_queue_ptr(device_id, q.addressof_ref()) for each. C++ never +// creates queues and never enumerates devices — Python is authoritative. +// +// Contract for C++ callers: sycl_get_queue_ptr() will throw if called +// before Python has registered a queue for the current device. In +// practice this means any C++ path reachable from Python is safe, since +// `import cupy` runs first; standalone C++ test binaries must call +// sycl_set_queue_ptr() themselves. +// +// Thread safety: g_current_device is std::atomic — any thread can read +// the current device id. sycl_set_device() / sycl_set_queue_ptr() are +// serialised by an internal mutex. +// ===================================================================== + +extern "C" { + int sycl_get_device_id(); + void* sycl_get_queue_ptr(); + void* sycl_get_queue_ptr_for(int device_id); + void sycl_set_queue_ptr(int device_id, void* queue_ptr); + void sycl_set_device(int device_id); +} + +// Convenience: return typed sycl::queue* for C++ callers +static inline sycl::queue* sycl_get_queue() { + return static_cast(sycl_get_queue_ptr()); +} + +static inline int syclGetDevice(int* id) { + *id = sycl_get_device_id(); + return 0; +} + +static inline void syclSetDevice(int id) { + sycl_set_device(id); +} + +// --- CUDA-compat device-property helpers --- + +struct cudaDeviceProp { + int multiProcessorCount; +}; + +static inline cudaError_t cudaGetDeviceProperties(cudaDeviceProp* prop, int /*device*/) { + prop->multiProcessorCount = static_cast( + sycl_get_queue()->get_device().get_info()); + return cudaSuccess; +} + +// --- CUDA-compat memory helpers (use current thread's queue) --- + +static inline void cudaMalloc(void** ptr, size_t size) { + (*ptr) = sycl::malloc_device(size, *(sycl_get_queue())); +} + +static inline void cudaFree(void* ptr) { + sycl::free(ptr, *(sycl_get_queue())); +} + +static inline void cudaMemset(void* ptr, int val, size_t size) { + sycl_get_queue()->memset(ptr, static_cast(val), size).wait(); +} diff --git a/gpu4pyscf/lib/gvhf-md/CMakeLists.txt b/gpu4pyscf/lib/gvhf-md/CMakeLists.txt index 30ec559f6..7cc298d73 100644 --- a/gpu4pyscf/lib/gvhf-md/CMakeLists.txt +++ b/gpu4pyscf/lib/gvhf-md/CMakeLists.txt @@ -1,6 +1,4 @@ -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128") - -add_library(gvhf_md SHARED +set(GPU_SRCS md_contract_j.cu unrolled_md_j.cu unrolled_md_j_4dm.cu md_indices.cu md_j_driver.cu md_pairdata.c contract_int3c2e.cu @@ -8,14 +6,27 @@ add_library(gvhf_md SHARED contract_Et.cu ) -#option(BUILD_SHARED_LIBS "build shared libraries" 1) -#option(ENABLE_STATIC "Enforce static library build" 0) -#if(ENABLE_STATIC) -# set(BUILD_SHARED_LIBS 0) -#endif() +add_library(gvhf_md SHARED ${GPU_SRCS}) set_target_properties(gvhf_md PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_SEPARABLE_COMPILATION ON) + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + +if (USE_SYCL) + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + + # Under CUDA, md_indices.cu is a translation unit: it defines the + # __device__/__constant__ index tables that every other TU declares + # extern. SYCL has no cross-TU __device__ linkage, so md_j.cuh + # #includes md_indices.cu directly and the tables become + # `inline constexpr`. Compiling it as a TU as well would be redundant. + set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/md_indices.cu + PROPERTIES HEADER_FILE_ONLY ON) + + target_compile_options(gvhf_md PRIVATE -x c++ -nocudainc -nocudalib) + target_link_libraries(gvhf_md PRIVATE sycl_compat gsycl) +else() + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v") + set_target_properties(gvhf_md PROPERTIES CUDA_SEPARABLE_COMPILATION ON) +endif() -target_link_libraries(gvhf_md OpenMP::OpenMP_C) +target_link_libraries(gvhf_md PRIVATE OpenMP::OpenMP_C) diff --git a/gpu4pyscf/lib/gvhf-md/contract_Et.cu b/gpu4pyscf/lib/gvhf-md/contract_Et.cu index 3ed8c151e..82f5080f8 100644 --- a/gpu4pyscf/lib/gvhf-md/contract_Et.cu +++ b/gpu4pyscf/lib/gvhf-md/contract_Et.cu @@ -33,7 +33,12 @@ void dm_to_Rt_kernel(double *out, double *dm, int n_dm, RysIntEnvVars envs, uint32_t *bas_ij_idx, int *pair_loc, int npairs, int *ao_loc) { - int pair_ij = blockIdx.x * blockDim.x + threadIdx.x; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int pair_ij = item.get_group(0) * item.get_local_range(0) + item.get_local_id(0); +#else + int pair_ij = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (pair_ij >= npairs) { return; } @@ -175,7 +180,12 @@ void Rt_to_dm_kernel(double *dm, double *Rt, int n_dm, RysIntEnvVars envs, uint32_t *bas_ij_idx, int *pair_loc, int npairs, int *ao_loc) { - int pair_ij = blockIdx.x * blockDim.x + threadIdx.x; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int pair_ij = item.get_group(0) * item.get_local_range(0) + item.get_local_id(0); +#else + int pair_ij = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (pair_ij >= npairs) { return; } @@ -316,7 +326,12 @@ __global__ static void aux_to_Rt_kernel(double *out, double *aux, RysIntEnvVars envs, int *aux_loc, int *aux_xyz_loc, int nbas_aux) { - int ksh = blockIdx.x * blockDim.x + threadIdx.x; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int ksh = item.get_group(0) * item.get_local_range(0) + item.get_local_id(0); +#else + int ksh = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (ksh >= nbas_aux) { return; } @@ -487,12 +502,19 @@ int dm_to_Rt(double *out, double *dm, int n_dm, RysIntEnvVars *envs, uint32_t *bas_ij_idx, int *pair_loc, int npairs, int *ao_loc) { int blocks = (npairs + THREADS - 1) / THREADS; + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * THREADS, THREADS), [=](auto item) { + dm_to_Rt_kernel(out, dm, n_dm, dev_envs, bas_ij_idx, pair_loc, npairs, ao_loc); + }); + #else dm_to_Rt_kernel<<>>(out, dm, n_dm, *envs, bas_ij_idx, pair_loc, npairs, ao_loc); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in dm_to_Rt_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -500,12 +522,19 @@ int Rt_to_dm(double *dm, double *Rt, int n_dm, RysIntEnvVars *envs, uint32_t *bas_ij_idx, int *pair_loc, int npairs, int *ao_loc) { int blocks = (npairs + THREADS - 1) / THREADS; + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * THREADS, THREADS), [=](auto item) { + Rt_to_dm_kernel(dm, Rt, n_dm, dev_envs, bas_ij_idx, pair_loc, npairs, ao_loc); + }); + #else Rt_to_dm_kernel<<>>(dm, Rt, n_dm, *envs, bas_ij_idx, pair_loc, npairs, ao_loc); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in Rt_to_dm_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -513,12 +542,19 @@ int aux_to_Rt(double *out, double *aux, RysIntEnvVars *envs, int *aux_loc, int *aux_xyz_loc, int nbas_aux) { int blocks = (nbas_aux + THREADS - 1) / THREADS; + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * THREADS, THREADS), [=](auto item) { + aux_to_Rt_kernel(out, aux, dev_envs, aux_loc, aux_xyz_loc, nbas_aux); + }); + #else aux_to_Rt_kernel<<>>(out, aux, *envs, aux_loc, aux_xyz_loc, nbas_aux); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in aux_to_Rt_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-md/contract_int3c2e.cu b/gpu4pyscf/lib/gvhf-md/contract_int3c2e.cu index a656c4536..cac99dc44 100644 --- a/gpu4pyscf/lib/gvhf-md/contract_int3c2e.cu +++ b/gpu4pyscf/lib/gvhf-md/contract_int3c2e.cu @@ -32,11 +32,14 @@ template __device__ inline void iter_Rt_n(double *Rt, double rx, double ry, double rz, int l, int nsq_per_block, int gout_id, int gout_stride) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); +#endif int nf2 = (l + 1) * (l + 2) / 2; int nf3 = nf2 * (l + 3) / 3; int offsets = nf3 * l / 4 - l; //l*(l+1)*(l+2)*(l+3)/24 - l; - uint16_t *p1 = c_Rt_idx + offsets; - int8_t *tuv_fac = c_Rt_tuv_fac + offsets; + const uint16_t *p1 = c_Rt_idx + offsets; + const int8_t *tuv_fac = c_Rt_tuv_fac + offsets; double Rt_tmp[RT_SIZE]; nf2 -= 1; // Drop the first element in Rt. It is assigned outside nf3 -= 1; @@ -185,9 +188,12 @@ void _dot_Et(double *out, double *Rt, double ai) } } +// Dot the P-tensor of the aux basis (efg_phase-weighted auxvec, already cached +// at "auxvec") into Rt at Hermite index i (out of nf3ij), producing the +// ij-Cartesian component contribution of contract('ijP,P->ij', int3c2e, auxvec). template __device__ inline void _dot_aux(double& out, double *Rt, double *auxvec, - uint16_t *p1_ij, int nf3ij, int i, int nsp_per_block) + const uint16_t *p1_ij, int nf3ij, int i, int nsp_per_block) { if constexpr (L == 0) { out += Rt[p1_ij[0*nf3ij+i]*nsp_per_block] * auxvec[0]; @@ -332,17 +338,70 @@ void _dot_aux(double& out, double *Rt, double *auxvec, template __device__ inline void unrolled_contract_int3c2e(RysIntEnvVars& envs, JKMatrix& jk, int *shl_pair_offsets, uint32_t *bas_ij_idx, - int *pair_ij_loc, int *nsp_lookup) + int *pair_ij_loc, int *nsp_lookup + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int blockDim_x = item.get_local_range(1); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + // Pack small shared vars into a single group_local_memory allocation + // instead of 9 separate ones + struct SharedVars { + int shl_pair0, shl_pair1, order, nf3ij, nf3ijkl, kprim; + int nsp_per_block, Rt_stride; + double rk[3]; + double ak, ck; + double shared[8]; + }; + + auto thread_block = item.get_group(); + double *phase = reinterpret_cast(shm_mem); + auto &sv = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair0 = sv.shl_pair0; + int &shl_pair1 = sv.shl_pair1; + int &order = sv.order; + int &nf3ij = sv.nf3ij; + int &nf3ijkl = sv.nf3ijkl; + int &kprim = sv.kprim; + int &nsp_per_block = sv.nsp_per_block; + int &Rt_stride = sv.Rt_stride; + double (&rk)[3] = sv.rk; + double &ak = sv.ak; + double &ck = sv.ck; + double (&shared)[8] = sv.shared; + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int blockDim_x = blockDim.x; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + + __shared__ int shl_pair0, shl_pair1; + __shared__ int order, nf3ij, nf3ijkl, kprim; + __shared__ int nsp_per_block, Rt_stride; + __shared__ double rk[3]; + extern __shared__ double phase[]; + + __shared__ double ak, ck; + __shared__ double shared[8]; + #endif constexpr int lk = LK; constexpr int nfk = (lk + 1) * (lk + 2) / 2; constexpr int nf3k = nfk * (lk + 3) / 3; - int sp_block_id = gridDim.y - blockIdx.y - 1; - int ksh = gridDim.x - blockIdx.x - 1 + envs.nbas; - int thread_id = threadIdx.x; + int sp_block_id = gridDim_y - blockIdx_y - 1; + int ksh = gridDim_x - blockIdx_x - 1 + envs.nbas; + int thread_id = threadIdx_x; int *bas = envs.bas; double *env = envs.env; - __shared__ int shl_pair0, shl_pair1; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -354,9 +413,6 @@ void unrolled_contract_int3c2e(RysIntEnvVars& envs, JKMatrix& jk, int li = bas[ANG_OF + ish0*BAS_SLOTS]; int lj = bas[ANG_OF + jsh0*BAS_SLOTS]; int lij = li + lj; - __shared__ int order, nf3ij, kprim; - __shared__ int nsp_per_block, Rt_stride; - __shared__ double rk[3]; if (thread_id == 0) { order = lij + lk; nf3ij = (lij+1)*(lij+2)*(lij+3) / 6; @@ -366,20 +422,18 @@ void unrolled_contract_int3c2e(RysIntEnvVars& envs, JKMatrix& jk, rk[1] = env[rk_ptr+1]; rk[2] = env[rk_ptr+2]; nsp_per_block = nsp_lookup[lij*(L_AUX_MAX+1)+lk]; - Rt_stride = blockDim.x / nsp_per_block; + Rt_stride = blockDim_x / nsp_per_block; } __syncthreads(); int sp_id = thread_id % nsp_per_block; int Rt_id = thread_id / nsp_per_block; - extern __shared__ double shared_memory[]; - double *gamma_inc = shared_memory + sp_id; - double *Rt = shared_memory + (order+1) * nsp_per_block + sp_id; - uint16_t *p1_ij = Rt2_kl_ij + Rt2_idx_offsets[lij*RT2_MAX+lk]; - int8_t *efg_phase = c_Rt2_efg_phase + Rt2_idx_offsets[lk]; + double *gamma_inc = phase + sp_id; + double *Rt = phase + (order+1) * nsp_per_block + sp_id; + const uint16_t *p1_ij = Rt2_kl_ij + Rt2_idx_offsets[lij*RT2_MAX+lk]; + const int8_t *efg_phase = c_Rt2_efg_phase + Rt2_idx_offsets[lk]; for (int kp = 0; kp < kprim; ++kp) { __syncthreads(); - __shared__ double ak, ck; if (thread_id == 0) { ck = env[bas[ksh*BAS_SLOTS+PTR_COEFF] + kp] * PI_FAC; ak = env[bas[ksh*BAS_SLOTS+PTR_EXP] + kp]; @@ -517,12 +571,12 @@ void unrolled_contract_int3c2e(RysIntEnvVars& envs, JKMatrix& jk, val += __shfl_down_sync(0xffffffff, val, offset); } if (lane == 0) { - shared_memory[wid] = val; + phase[wid] = val; } __syncthreads(); if (thread_id < 8) { - val = shared_memory[lane]; + val = phase[lane]; } for (int offset = 4; offset > 0; offset >>= 1) { val += __shfl_down_sync(0xff, val, offset); @@ -538,8 +592,25 @@ void unrolled_contract_int3c2e(RysIntEnvVars& envs, JKMatrix& jk, __global__ static void contract_int3c2e_kernel(RysIntEnvVars envs, JKMatrix jk, int *shl_pair_offsets, uint32_t *bas_ij_idx, - int *pair_ij_loc, int *nsp_lookup) + int *pair_ij_loc, int *nsp_lookup + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { + #ifdef USE_SYCL + int ksh = item.get_group_range(1) - item.get_group(1) - 1 + envs.nbas; + int lk = envs.bas[ANG_OF + ksh*BAS_SLOTS]; + switch (lk) { + case 0: unrolled_contract_int3c2e<0,42>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, item, shm_mem); break; + case 1: unrolled_contract_int3c2e<1,42>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, item, shm_mem); break; + case 2: unrolled_contract_int3c2e<2,42>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, item, shm_mem); break; + case 3: unrolled_contract_int3c2e<3,42>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, item, shm_mem); break; + case 4: unrolled_contract_int3c2e<4,42>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, item, shm_mem); break; + case 5: unrolled_contract_int3c2e<5,42>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, item, shm_mem); break; + case 6: unrolled_contract_int3c2e<6,30>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, item, shm_mem); break; + } + #else int ksh = gridDim.x - blockIdx.x - 1 + envs.nbas; int lk = envs.bas[ANG_OF + ksh*BAS_SLOTS]; switch (lk) { @@ -551,22 +622,72 @@ void contract_int3c2e_kernel(RysIntEnvVars envs, JKMatrix jk, case 5: unrolled_contract_int3c2e<5,42>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup); break; case 6: unrolled_contract_int3c2e<6,30>(envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup); break; } + #endif } +// IJ_SIZE bounds the number of (Rt_id-strided) ij Hermite components handled +// by a single thread; picked per LK to match the nsp_per_block/Rt_stride +// combinations produced by the host-side sizing formula in j_engine_3c2e.py. template __device__ inline void unroll_contract_auxvec(RysIntEnvVars& envs, JKMatrix& jk, int *shl_pair_offsets, int *ksh_offsets, uint32_t *bas_ij_idx, int *pair_ij_loc, - int *aux_loc, int *nsp_lookup) + int *aux_loc, int *nsp_lookup + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - int thread_id = threadIdx.x; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int blockDim_x = item.get_local_range(1); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + // Pack small shared vars into a single group_local_memory allocation + // instead of several separate ones + struct SharedVars { + int shl_pair0, shl_pair1, ksh0, ksh1; + int order, nf3ij, nf3ijkl; + int nsp_per_block, Rt_stride; + }; + + auto thread_block = item.get_group(); + double *shared_memory = reinterpret_cast(shm_mem); + auto &sv = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair0 = sv.shl_pair0; + int &shl_pair1 = sv.shl_pair1; + int &ksh0 = sv.ksh0; + int &ksh1 = sv.ksh1; + int &order = sv.order; + int &nf3ij = sv.nf3ij; + int &nf3ijkl = sv.nf3ijkl; + int &nsp_per_block = sv.nsp_per_block; + int &Rt_stride = sv.Rt_stride; + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int blockDim_x = blockDim.x; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + + __shared__ int shl_pair0, shl_pair1, ksh0, ksh1; + __shared__ int order, nf3ij, nf3ijkl; + __shared__ int nsp_per_block, Rt_stride; + extern __shared__ double shared_memory[]; + #endif + constexpr int lk = LK; + constexpr int nfk = (lk + 1) * (lk + 2) / 2; + constexpr int nf3k = nfk * (lk + 3) / 3; + int sp_block_id = gridDim_x - blockIdx_x - 1; + int ksh_block_id = gridDim_y - blockIdx_y - 1; + int thread_id = threadIdx_x; int *bas = envs.bas; double *env = envs.env; - __shared__ int shl_pair0, shl_pair1; - __shared__ int ksh0, ksh1; if (thread_id == 0) { - int sp_block_id = gridDim.x - blockIdx.x - 1; - int ksh_block_id = gridDim.y - blockIdx.y - 1; ksh0 = ksh_offsets[ksh_block_id]; ksh1 = ksh_offsets[ksh_block_id+1]; shl_pair0 = shl_pair_offsets[sp_block_id]; @@ -578,27 +699,22 @@ void unroll_contract_auxvec(RysIntEnvVars& envs, JKMatrix& jk, int jsh0 = bas_ij0 % envs.nbas; int li = bas[ANG_OF + ish0*BAS_SLOTS]; int lj = bas[ANG_OF + jsh0*BAS_SLOTS]; - //int lk = bas[ksh0*BAS_SLOTS+ANG_OF]; - constexpr int lk = LK; int lij = li + lj; - int order = lij + lk; - constexpr int nfk = (lk + 1) * (lk + 2) / 2; - constexpr int nf3k = nfk * (lk + 3) / 3; - int nf3ij = (lij+1)*(lij+2)*(lij+3) / 6; - __shared__ int nsp_per_block, Rt_stride; if (thread_id == 0) { + order = lij + lk; + nf3ij = (lij+1)*(lij+2)*(lij+3) / 6; nsp_per_block = nsp_lookup[lij*(L_AUX_MAX+1)+lk]; - Rt_stride = blockDim.x / nsp_per_block; + Rt_stride = blockDim_x / nsp_per_block; } __syncthreads(); int sp_id = thread_id % nsp_per_block; int Rt_id = thread_id / nsp_per_block; - extern __shared__ double shared_memory[]; + double *gamma_inc = shared_memory + sp_id; double *auxvec_cache = shared_memory + (order+1) * nsp_per_block; double *Rt = auxvec_cache + nf3k + sp_id; - uint16_t *p1_ij = Rt2_kl_ij + Rt2_idx_offsets[lij*RT2_MAX+lk]; - int8_t *efg_phase = c_Rt2_efg_phase + Rt2_idx_offsets[lk]; + const uint16_t *p1_ij = Rt2_kl_ij + Rt2_idx_offsets[lij*RT2_MAX+lk]; + const int8_t *efg_phase = c_Rt2_efg_phase + Rt2_idx_offsets[lk]; double *auxvec = jk.dm; for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { @@ -623,6 +739,7 @@ void unroll_contract_auxvec(RysIntEnvVars& envs, JKMatrix& jk, double xij = (ai * ri[0] + aj * rj[0]) / aij; double yij = (ai * ri[1] + aj * rj[1]) / aij; double zij = (ai * ri[2] + aj * rj[2]) / aij; + for (int ksh = ksh0; ksh < ksh1; ++ksh) { __syncthreads(); int k_loc0 = aux_loc[ksh - envs.nbas]; @@ -639,6 +756,9 @@ void unroll_contract_auxvec(RysIntEnvVars& envs, JKMatrix& jk, double ak = env[expk]; double theta = aij * ak / (aij + ak); if (Rt_id == 0) { + // auxvec is already scaled by the aux contraction coefficient + // via the host-side Et_dot_auxvec pre-processing step, so only + // the geometric prefactor is needed here. double fac = PI_FAC/(aij*ak*sqrt(aij+ak)); if (pair_ij >= shl_pair1) { fac = 0; @@ -719,11 +839,12 @@ void unroll_contract_auxvec(RysIntEnvVars& envs, JKMatrix& jk, } } } + if (pair_ij < shl_pair1) { + int ij_loc0 = pair_ij_loc[pair_ij]; #pragma unroll for (int n = 0, i = Rt_id; n < IJ_SIZE; ++n, i += Rt_stride) { if (i >= nf3ij) break; - int ij_loc0 = pair_ij_loc[pair_ij]; atomicAdd(jk.vj+ij_loc0+i, vj_xyz[n]); } } @@ -734,8 +855,26 @@ __global__ static void contract_auxvec_kernel(RysIntEnvVars envs, JKMatrix jk, int *shl_pair_offsets, int *ksh_offsets, uint32_t *bas_ij_idx, int *pair_ij_loc, - int *aux_loc, int *nsp_lookup) + int *aux_loc, int *nsp_lookup + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { + #ifdef USE_SYCL + int ksh_block_id = item.get_group_range(0) - item.get_group(0) - 1; + int ksh = ksh_offsets[ksh_block_id]; + int lk = envs.bas[ANG_OF + ksh*BAS_SLOTS]; + switch (lk) { + case 0: unroll_contract_auxvec<0,35,35>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup, item, shm_mem); break; + case 1: unroll_contract_auxvec<1,21,35>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup, item, shm_mem); break; + case 2: unroll_contract_auxvec<2,15,35>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup, item, shm_mem); break; + case 3: unroll_contract_auxvec<3,11,35>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup, item, shm_mem); break; + case 4: unroll_contract_auxvec<4, 8,35>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup, item, shm_mem); break; + case 5: unroll_contract_auxvec<5, 8,21>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup, item, shm_mem); break; + case 6: unroll_contract_auxvec<6, 8,21>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup, item, shm_mem); break; + } + #else int ksh_block_id = gridDim.y - blockIdx.y - 1; int ksh = ksh_offsets[ksh_block_id]; int lk = envs.bas[ANG_OF + ksh*BAS_SLOTS]; @@ -748,6 +887,7 @@ void contract_auxvec_kernel(RysIntEnvVars envs, JKMatrix jk, case 5: unroll_contract_auxvec<5, 8,21>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup); break; case 6: unroll_contract_auxvec<6, 8,21>(envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, aux_loc, nsp_lookup); break; } + #endif } extern "C" { @@ -759,8 +899,20 @@ int contract_int3c2e_dm(double *vj, double *dm, int n_dm, int naux, int *pair_ij_loc, int *nsp_lookup, double omega) { assert(n_dm == 1); - cudaFuncSetAttribute(contract_int3c2e_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); JKMatrix jk = {vj, NULL, dm, n_dm, 0, omega}; + #ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_shl_pair, nksh); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(shm_size, cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + contract_int3c2e_kernel(dev_envs, jk, shl_pair_offsets, bas_ij_idx, pair_ij_loc, nsp_lookup, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + cudaFuncSetAttribute(contract_int3c2e_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 threads(THREADS); dim3 blocks(nksh, nbatches_shl_pair); contract_int3c2e_kernel<<>>( @@ -770,6 +922,7 @@ int contract_int3c2e_dm(double *vj, double *dm, int n_dm, int naux, fprintf(stderr, "CUDA Error in contract_int3c2e_dm, error message = %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -782,8 +935,20 @@ int contract_int3c2e_auxvec(double *vj, double *auxvec, int n_dm, int naux, int *nsp_lookup, double omega) { assert(n_dm == 1); - cudaFuncSetAttribute(contract_auxvec_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); JKMatrix jk = {vj, NULL, auxvec, n_dm, 0, omega}; + #ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_ksh, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(shm_size, cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + contract_auxvec_kernel(dev_envs, jk, shl_pair_offsets, ksh_offsets, bas_ij_idx, pair_ij_loc, + aux_loc, nsp_lookup, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + cudaFuncSetAttribute(contract_auxvec_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 threads(THREADS); dim3 blocks(nbatches_shl_pair, nbatches_ksh); contract_auxvec_kernel<<>>( @@ -794,6 +959,7 @@ int contract_int3c2e_auxvec(double *vj, double *auxvec, int n_dm, int naux, fprintf(stderr, "CUDA Error in contract_int3c2e_auxvec, error message = %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-md/md_contract_j.cu b/gpu4pyscf/lib/gvhf-md/md_contract_j.cu index e79353629..283b03516 100644 --- a/gpu4pyscf/lib/gvhf-md/md_contract_j.cu +++ b/gpu4pyscf/lib/gvhf-md/md_contract_j.cu @@ -40,11 +40,14 @@ __device__ inline void iter_Rt_n(double *Rt, double rx, double ry, double rz, int l, int nsq_per_block, int gout_id, int gout_stride) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + #endif int nf2 = (l + 1) * (l + 2) / 2; int nf3 = nf2 * (l + 3) / 3; int offsets = nf3 * l / 4 - l; //l*(l+1)*(l+2)*(l+3)/24 - l; - uint16_t *p1 = c_Rt_idx + offsets; - int8_t *tuv_fac = c_Rt_tuv_fac + offsets; + const uint16_t *p1 = c_Rt_idx + offsets; + const int8_t *tuv_fac = c_Rt_tuv_fac + offsets; double Rt_tmp[RT_TMP_SIZE]; nf2 -= 1; // Drop the first element in Rt. It is assigned outside nf3 -= 1; @@ -72,14 +75,35 @@ __global__ void md_j_1dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, int threadsx, int threadsy, int tilex, int tiley, - uint16_t *pRt2_kl_ij, int8_t *efg_phase) + const uint16_t *pRt2_kl_ij, const int8_t *efg_phase + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { +#ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + double *vj_kl_cache = reinterpret_cast(shm_mem); +#else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + extern __shared__ double vj_kl_cache[]; +#endif int *pair_ij_mapping = bounds.pair_ij_mapping; int *pair_kl_mapping = bounds.pair_kl_mapping; int bsizex = threadsx * tilex; int bsizey = threadsy * tiley; - int task_ij0 = blockIdx.x * bsizex; - int task_kl0 = blockIdx.y * bsizey; + int task_ij0 = blockIdx_x * bsizex; + int task_kl0 = blockIdx_y * bsizey; if (q_cond_ij[task_ij0] + q_cond_kl[task_kl0] < bounds.cutoff) { return; } @@ -92,10 +116,10 @@ void md_j_1dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, return; } - int sq_id = threadIdx.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - int nsq_per_block = blockDim.x; + int sq_id = threadIdx_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + int nsq_per_block = blockDim_x; //assert(nsq_per_block == threadsx * threadsy); int t_id = gout_id * nsq_per_block + sq_id; int lane_id = t_id % warpSize; @@ -119,13 +143,12 @@ void md_j_1dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, int npairs_ij = bounds.npairs_ij; int npairs_kl = bounds.npairs_kl; - extern __shared__ double vj_kl_cache[]; double *Rq_cache = vj_kl_cache + nf3kl*bsizey; double *Rp_cache = vj_kl_cache + bsizey*(4+nf3kl); double *dm_ij_cache = vj_kl_cache + bsizey*(4+nf3kl) + threadsx*4 + tx; double *gamma_inc = vj_kl_cache + bsizey*(4+nf3kl) + threadsx*(4+nf3ij) + sq_id; double *Rt = gamma_inc + (order+1) * nsq_per_block; - uint16_t *Rt2_address = pRt2_kl_ij; + uint16_t *Rt2_address = const_cast(pRt2_kl_ij); if (nf3ij * nf3kl <= RT2_IDX_CACHE_SIZE) { int l4 = bounds.lij + bounds.lkl; int nf3 = (l4 + 1) * (l4 + 2) * (l4 + 3) / 6; @@ -143,7 +166,7 @@ void md_j_1dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, } __syncthreads(); for (int n = t_id; n < bsizey; n += threads) { - int task_kl = blockIdx.y * bsizey + n; + int task_kl = blockIdx_y * bsizey + n; if (task_kl < npairs_kl) { int pair_kl = pair_kl_mapping[task_kl]; int ksh = pair_kl / nbas; @@ -169,7 +192,7 @@ void md_j_1dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, } for (int batch_ij = 0; batch_ij < tilex; ++batch_ij) { - int task_ij0 = (blockIdx.x * tilex + batch_ij) * threadsx; + int task_ij0 = (blockIdx_x * tilex + batch_ij) * threadsx; if (task_ij0 >= npairs_ij) { break; } @@ -213,15 +236,15 @@ void md_j_1dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, vj_ij[n] = 0.; } for (int batch_kl = 0; batch_kl < tiley; ++batch_kl) { - int task_kl0 = (blockIdx.y * tiley + batch_kl) * threadsy; + int task_kl0 = (blockIdx_y * tiley + batch_kl) * threadsy; if (task_kl0 >= npairs_kl) { break; } if (pair_ij_mapping == pair_kl_mapping && task_ij0+threadsx <= task_kl0) { break; } - if (qd_ij_max[blockIdx.x*tilex+batch_ij] + q_cond_kl[task_kl0] < bounds.cutoff && - qd_kl_max[blockIdx.y*tiley+batch_kl] + q_cond_ij[task_ij0] < bounds.cutoff) { + if (qd_ij_max[blockIdx_x*tilex+batch_ij] + q_cond_kl[task_kl0] < bounds.cutoff && + qd_kl_max[blockIdx_y*tiley+batch_kl] + q_cond_ij[task_ij0] < bounds.cutoff) { continue; } @@ -389,7 +412,7 @@ void md_j_1dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, int kl = n / tiley; int batch_kl = n - kl * tiley; int sq_kl = ty + batch_kl * threadsy; - int task_kl = blockIdx.y * bsizey + sq_kl; + int task_kl = blockIdx_y * bsizey + sq_kl; if (task_kl < npairs_kl) { int kl_loc0 = pair_kl_loc[task_kl]; atomicAdd(vj+kl_loc0+kl, vj_kl_cache[sq_kl+kl*bsizey]); @@ -402,14 +425,35 @@ __global__ void md_j_4dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, int threadsx, int threadsy, int tilex, int tiley, int dm_size, - uint16_t *pRt2_kl_ij, int8_t *efg_phase) + const uint16_t *pRt2_kl_ij, const int8_t *efg_phase + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { +#ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + double *vj_kl_cache = reinterpret_cast(shm_mem); +#else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + extern __shared__ double vj_kl_cache[]; +#endif int *pair_ij_mapping = bounds.pair_ij_mapping; int *pair_kl_mapping = bounds.pair_kl_mapping; int bsizex = threadsx * tilex; int bsizey = threadsy * tiley; - int task_ij0 = blockIdx.x * bsizex; - int task_kl0 = blockIdx.y * bsizey; + int task_ij0 = blockIdx_x * bsizex; + int task_kl0 = blockIdx_y * bsizey; if (q_cond_ij[task_ij0] + q_cond_kl[task_kl0] < bounds.cutoff) { return; } @@ -422,10 +466,10 @@ void md_j_4dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, return; } - int sq_id = threadIdx.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - int nsq_per_block = blockDim.x; + int sq_id = threadIdx_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + int nsq_per_block = blockDim_x; //assert(nsq_per_block == threadsx * threadsy); int t_id = gout_id * nsq_per_block + sq_id; int lane_id = t_id % warpSize; @@ -453,13 +497,12 @@ void md_j_4dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, int npairs_ij = bounds.npairs_ij; int npairs_kl = bounds.npairs_kl; - extern __shared__ double vj_kl_cache[]; double *Rq_cache = vj_kl_cache + nf3kl*bsizey * DM_BLOCK; double *Rp_cache = Rq_cache + bsizey*4; double *dm_ij_cache = Rp_cache + threadsx*4 + tx; double *gamma_inc = Rp_cache + threadsx*4 + nf3ij * threadsx * DM_BLOCK + sq_id; double *Rt = gamma_inc + (order+1) * nsq_per_block; - uint16_t *Rt2_address = pRt2_kl_ij; + uint16_t *Rt2_address = const_cast(pRt2_kl_ij); // vj_cache requires a size of nthreads*n_dm. order=0 (corresponding to // (ss|ss)) is skipped because the addresses of vj_cache and Rt2_address // overlap. @@ -476,7 +519,7 @@ void md_j_4dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, __syncthreads(); for (int n = t_id; n < bsizey; n += threads) { - int task_kl = blockIdx.y * bsizey + n; + int task_kl = blockIdx_y * bsizey + n; if (task_kl < npairs_kl) { int pair_kl = pair_kl_mapping[task_kl]; int ksh = pair_kl / nbas; @@ -506,7 +549,7 @@ void md_j_4dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, } for (int batch_ij = 0; batch_ij < tilex; ++batch_ij) { - int task_ij0 = (blockIdx.x * tilex + batch_ij) * threadsx; + int task_ij0 = (blockIdx_x * tilex + batch_ij) * threadsx; if (task_ij0 >= npairs_ij) { break; } @@ -554,16 +597,16 @@ void md_j_4dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, vj_ij[n] = 0.; } for (int batch_kl = 0; batch_kl < tiley; ++batch_kl) { - int task_kl0 = (blockIdx.y * tiley + batch_kl) * threadsy; + int task_kl0 = (blockIdx_y * tiley + batch_kl) * threadsy; if (task_kl0 >= npairs_kl) { break; } - int task_ij0 = (blockIdx.x * tilex + batch_ij) * threadsx; + int task_ij0 = (blockIdx_x * tilex + batch_ij) * threadsx; if (pair_ij_mapping == pair_kl_mapping && task_ij0+threadsx <= task_kl0) { break; } - if (qd_ij_max[blockIdx.x*tilex+batch_ij] + q_cond_kl[task_kl0] < bounds.cutoff && - qd_kl_max[blockIdx.y*tiley+batch_kl] + q_cond_ij[task_ij0] < bounds.cutoff) { + if (qd_ij_max[blockIdx_x*tilex+batch_ij] + q_cond_kl[task_kl0] < bounds.cutoff && + qd_kl_max[blockIdx_y*tiley+batch_kl] + q_cond_ij[task_ij0] < bounds.cutoff) { continue; } @@ -894,7 +937,7 @@ void md_j_4dm_kernel(RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, int kl = n / tiley; int batch_kl = n - kl * tiley; int sq_kl = ty + batch_kl * threadsy; - int task_kl = blockIdx.y * bsizey + sq_kl; + int task_kl = blockIdx_y * bsizey + sq_kl; if (task_kl < npairs_kl) { int kl_loc0 = pair_kl_loc[task_kl]; switch (jk.n_dm) { @@ -972,16 +1015,28 @@ int MD_build_j(double *vj, double *dm, int n_dm, int dm_size, int bsizex = threads_ij * tilex; int bsizey = threads_kl * tiley; int nsq_per_block = threads_ij * threads_kl; - dim3 threads(nsq_per_block, gout_stride); int blocks_ij = (npairs_ij + bsizex - 1) / bsizex; int blocks_kl = (npairs_kl + bsizey - 1) / bsizey; + #ifdef USE_SYCL + sycl::range<2> threads(gout_stride, nsq_per_block); + sycl::range<2> blocks(blocks_kl, blocks_ij); + // IMP: SYCL doesnt treat the Rt2_kl_ij, c_Rt2_efg_phase + // pointer arithmetic on host and the obtained pointers are + // not valid on the device. Hence just compute the offset on host + // but obtain the pointer `pRt2_kl_ij` & `efg_phase` in the kernel launch + const int Rt2_kl_ij_syclonly_offset = offset_for_Rt2_idx(lij, lkl); + const int efg_phase_syclonly_offset = offset_for_Rt2_idx(0, lkl); + auto dev_envs = *envs; + #else + dim3 threads(nsq_per_block, gout_stride); dim3 blocks(blocks_ij, blocks_kl); - uint16_t *pRt2_kl_ij; - int8_t *efg_phase; + uint16_t *pRt2_kl_ij = nullptr; + int8_t *efg_phase = nullptr; cudaGetSymbolAddress((void**)&pRt2_kl_ij, Rt2_kl_ij); cudaGetSymbolAddress((void**)&efg_phase, c_Rt2_efg_phase); pRt2_kl_ij += offset_for_Rt2_idx(lij, lkl); efg_phase += offset_for_Rt2_idx(0, lkl); + #endif if (n_dm == 1) { if (!md_j_unrolled(envs, &jk, &bounds, q_cond_ij, q_cond_kl, omega)) { cudaFuncSetAttribute(md_j_1dm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); @@ -993,9 +1048,23 @@ int MD_build_j(double *vj, double *dm, int n_dm, int dm_size, } bounds.qd_ij_max = qd_ij_max + qd_offset_for_threads(npairs_ij, threads_ij); bounds.qd_kl_max = qd_kl_max + qd_offset_for_threads(npairs_kl, threads_kl); + #ifdef USE_SYCL + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(buflen, cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + const uint16_t *pRt2_kl_ij = Rt2_kl_ij + Rt2_kl_ij_syclonly_offset; + const int8_t *efg_phase = c_Rt2_efg_phase + efg_phase_syclonly_offset; + md_j_1dm_kernel(dev_envs, jk, bounds, q_cond_ij, q_cond_kl, + threads_ij, threads_kl, tilex, tiley, + pRt2_kl_ij, efg_phase, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else md_j_1dm_kernel<<>>( *envs, jk, bounds, q_cond_ij, q_cond_kl, threads_ij, threads_kl, tilex, tiley, pRt2_kl_ij, efg_phase); + #endif } } else { if (!md_j_4dm_unrolled(envs, &jk, &bounds, q_cond_ij, q_cond_kl, omega, dm_size)) { @@ -1012,9 +1081,23 @@ int MD_build_j(double *vj, double *dm, int n_dm, int dm_size, jk.vj = vj + dm_offset * dm_size; jk.dm = dm + dm_offset * dm_size; jk.n_dm = n_dm - dm_offset; + #ifdef USE_SYCL + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(buflen, cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + const uint16_t *pRt2_kl_ij = Rt2_kl_ij + Rt2_kl_ij_syclonly_offset; + const int8_t *efg_phase = c_Rt2_efg_phase + efg_phase_syclonly_offset; + md_j_4dm_kernel(dev_envs, jk, bounds, q_cond_ij, q_cond_kl, + threads_ij, threads_kl, tilex, tiley, dm_size, + pRt2_kl_ij, efg_phase, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else md_j_4dm_kernel<<>>( *envs, jk, bounds, q_cond_ij, q_cond_kl, threads_ij, threads_kl, tilex, tiley, dm_size, pRt2_kl_ij, efg_phase); + #endif } } } diff --git a/gpu4pyscf/lib/gvhf-md/md_indices.cu b/gpu4pyscf/lib/gvhf-md/md_indices.cu index 94f6cc05f..3d69235bf 100644 --- a/gpu4pyscf/lib/gvhf-md/md_indices.cu +++ b/gpu4pyscf/lib/gvhf-md/md_indices.cu @@ -18,7 +18,12 @@ #include #include -__device__ int Rt2_idx_offsets[] = { +#ifdef USE_SYCL +inline constexpr int Rt2_idx_offsets[] = +#else +__device__ int Rt2_idx_offsets[] = +#endif + { 0,1,5,15,35,70,126,210,330, 495,499,515,555,635,775,999,1335,1815, 2475,2485,2525,2625,2825,3175,3735,4575,5775, @@ -30,7 +35,12 @@ __device__ int Rt2_idx_offsets[] = { 163350,163515,164175,165825,169125,174900,184140,198000,217800, 245025, }; -__device__ uint16_t Rt2_ij_kl[] = { +#ifdef USE_SYCL +inline constexpr uint16_t Rt2_ij_kl[] = +#else +__device__ uint16_t Rt2_ij_kl[] = +#endif + { 0,0,1,2,3,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4, 5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,0,1,2,3,4, 5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24, @@ -12284,7 +12294,12 @@ __device__ uint16_t Rt2_ij_kl[] = { 944,945,946,947,948,949,950,951,952,953,954,955,956,957,958,959,960,961,962,963, 964,965,966,967,968, }; -__device__ uint16_t Rt2_kl_ij[] = { +#ifdef USE_SYCL +inline constexpr uint16_t Rt2_kl_ij[] = +#else +__device__ uint16_t Rt2_kl_ij[] = +#endif + { 0,0,1,2,3,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4, 5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,0,1,2,3,4, 5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24, @@ -24539,7 +24554,12 @@ __device__ uint16_t Rt2_kl_ij[] = { 964,965,966,967,968, }; -__constant__ int8_t c_Rt2_efg_phase[] = { +#ifdef USE_SYCL +inline constexpr int8_t c_Rt2_efg_phase[] = +#else +__constant__ int8_t c_Rt2_efg_phase[] = +#endif + { // l = 0 1, // l = 1 @@ -24581,7 +24601,12 @@ __constant__ int8_t c_Rt2_efg_phase[] = { }; // offsets = l*(l+1)*(l+2)*(l+3)//24 - l -__constant__ int8_t c_Rt_tuv_fac[] = { +#ifdef USE_SYCL +inline constexpr int8_t c_Rt_tuv_fac[] = +#else +__constant__ int8_t c_Rt_tuv_fac[] = +#endif + { // l = 1 0,0,0, // l = 2 @@ -24847,7 +24872,12 @@ __constant__ int8_t c_Rt_tuv_fac[] = { 11,11,11,11,11,11,11,11,12,12,12,12,12,12,12,12,12,12,13,13, 13,13,13,13,14,14,14,15, }; -__constant__ uint16_t c_Rt_idx[] = { +#ifdef USE_SYCL +inline constexpr uint16_t c_Rt_idx[] = +#else +__constant__ uint16_t c_Rt_idx[] = +#endif + { // l = 1 0,0,0, // l = 2 diff --git a/gpu4pyscf/lib/gvhf-md/md_j.cuh b/gpu4pyscf/lib/gvhf-md/md_j.cuh index 17d8058a1..8f61b5800 100644 --- a/gpu4pyscf/lib/gvhf-md/md_j.cuh +++ b/gpu4pyscf/lib/gvhf-md/md_j.cuh @@ -69,9 +69,31 @@ typedef struct { int offset_for_Rt2_idx(int lij, int lkl); int qd_offset_for_threads(int npairs, int threads); +#ifdef USE_SYCL +#include "md_indices.cu" + +// --------------------------------------------------------------------- +// blockIdx / threadIdx shim for the generated unrolled_md_j*.cu kernels. +// +// Those two files are ~15k lines of auto-generated kernel body that index +// the launch geometry directly as blockIdx.x/.y and threadIdx.x/.y. CUDA +// supplies those as built-ins; SYCL has no equivalent. Rather than rewrite +// every reference, KERNEL_SETUP() materialises two locals of this type +// from the nd_item, so the generated bodies stay byte-identical to +// upstream/master and only the macro preamble at the top of each file +// differs between the two backends. +// +// Axis mapping is fixed by the launch: sycl::nd_range<2> dimension 1 is +// the fast-varying axis and carries CUDA's .x, dimension 0 carries .y. +// --------------------------------------------------------------------- +struct md_j_index2 { + int x, y; +}; +#else extern __device__ int Rt2_idx_offsets[]; extern __device__ uint16_t Rt2_ij_kl[]; extern __device__ uint16_t Rt2_kl_ij[]; extern __constant__ int8_t c_Rt2_efg_phase[]; extern __constant__ int8_t c_Rt_tuv_fac[]; extern __constant__ uint16_t c_Rt_idx[]; +#endif diff --git a/gpu4pyscf/lib/gvhf-md/md_pairdata.c b/gpu4pyscf/lib/gvhf-md/md_pairdata.c index f8dbddd5d..a08cceecd 100644 --- a/gpu4pyscf/lib/gvhf-md/md_pairdata.c +++ b/gpu4pyscf/lib/gvhf-md/md_pairdata.c @@ -8,6 +8,9 @@ #define Ey_at(i,j,t) Ey[(i)*stride1+(j)*stride2+t] #define Ez_at(i,j,t) Ez[(i)*stride1+(j)*stride2+t] +#ifdef __cplusplus +extern "C" { +#endif void get_E_cart_components(double *Ecart, int li, int lj, double ai, double aj, double *Ra, double *Rb) { @@ -128,7 +131,7 @@ void Et_dot_dm(double *Et_dm, double *dm, int n_dm, int Et_dm_size, int l2 = 2*LMAX; int Et_size = (l2+1)*(l2+2)*(l2+3)/6*NCART_MAX*NCART_MAX; int Ex_size = (2*LMAX+1)*(LMAX+1)*(LMAX+1); - double *Et = malloc(sizeof(double) * (Et_size+3*Ex_size)); + double *Et = (double*)malloc(sizeof(double) * (Et_size+3*Ex_size)); double *buf = Et + Et_size; size_t nao = ao_loc[nbas]; size_t nao2 = nao * nao; @@ -181,7 +184,7 @@ void Et_dot_auxvec(double *Et_auxvec, double *auxvec, int n_dm, int *aux_xyz_loc { int Et_size = (L_AUX_MAX+1)*(L_AUX_MAX+2)*(L_AUX_MAX+3)/6*NCART_MAX*NCART_MAX; int Ex_size = (2*L_AUX_MAX+1)*(L_AUX_MAX+1)*(L_AUX_MAX+1); - double *Et = malloc(sizeof(double) * (Et_size+3*Ex_size)); + double *Et = (double*)malloc(sizeof(double) * (Et_size+3*Ex_size)); double *buf = Et + Et_size; int naux = aux_loc[nbas]; int Et_auxvec_size = aux_xyz_loc[nbas]; @@ -217,7 +220,7 @@ void jengine_dot_Et(double *vj, double *jvec, int n_dm, int Et_dm_size, int l2 = 2*LMAX; int Et_size = (l2+1)*(l2+2)*(l2+3)/6*NCART_MAX*NCART_MAX; int Ex_size = (2*LMAX+1)*(LMAX+1)*(LMAX+1); - double *Et = malloc(sizeof(double) * (Et_size+3*Ex_size)); + double *Et = (double*)malloc(sizeof(double) * (Et_size+3*Ex_size)); double *buf = Et + Et_size; size_t nao = ao_loc[nbas]; size_t nao2 = nao * nao; @@ -284,7 +287,7 @@ void PBC_Et_dot_dm(double *Et_dm, double *dm, int n_dm, int Et_dm_size, int l2 = 2*LMAX; int Et_size = (l2+1)*(l2+2)*(l2+3)/6*NCART_MAX*NCART_MAX; int Ex_size = (2*LMAX+1)*(LMAX+1)*(LMAX+1); - double *Et = malloc(sizeof(double) * (Et_size+3*Ex_size)); + double *Et = (double*)malloc(sizeof(double) * (Et_size+3*Ex_size)); double *buf = Et + Et_size; double rjL[3]; size_t nao = ao_loc[c_nbas]; // for the unit cell @@ -361,7 +364,7 @@ void PBC_jengine_dot_Et(double *vj, double *jvec, int n_dm, int Et_dm_size, int l2 = 2*LMAX; int Et_size = (l2+1)*(l2+2)*(l2+3)/6*NCART_MAX*NCART_MAX; int Ex_size = (2*LMAX+1)*(LMAX+1)*(LMAX+1); - double *Et = malloc(sizeof(double) * (Et_size+3*Ex_size)); + double *Et = (double*)malloc(sizeof(double) * (Et_size+3*Ex_size)); double *buf = Et + Et_size; double rjL[3]; size_t nao = ao_loc[c_nbas]; @@ -422,3 +425,7 @@ void PBC_jengine_dot_Et(double *vj, double *jvec, int n_dm, int Et_dm_size, free(Et); } } + +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/gpu4pyscf/lib/gvhf-md/pbc_md_contract_j.cu b/gpu4pyscf/lib/gvhf-md/pbc_md_contract_j.cu index 44b265701..a5fad88d0 100644 --- a/gpu4pyscf/lib/gvhf-md/pbc_md_contract_j.cu +++ b/gpu4pyscf/lib/gvhf-md/pbc_md_contract_j.cu @@ -40,11 +40,14 @@ __device__ inline void iter_Rt_n(double *Rt, double rx, double ry, double rz, int l, int nsq_per_block, int gout_id, int gout_stride) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + #endif int nf2 = (l + 1) * (l + 2) / 2; int nf3 = nf2 * (l + 3) / 3; int offsets = nf3 * l / 4 - l; //l*(l+1)*(l+2)*(l+3)/24 - l; - uint16_t *p1 = c_Rt_idx + offsets; - int8_t *tuv_fac = c_Rt_tuv_fac + offsets; + const uint16_t *p1 = c_Rt_idx + offsets; + const int8_t *tuv_fac = c_Rt_tuv_fac + offsets; double Rt_tmp[RT_TMP_SIZE]; nf2 -= 1; // Drop the first element in Rt. It is assigned outside nf3 -= 1; @@ -73,22 +76,43 @@ __global__ void pbc_md_j_kernel(RysIntEnvVars envs, JKMatrix jmat, MDBoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, int threadsx, int threadsy, int tilex, int tiley, - uint16_t *pRt2_kl_ij, int8_t *efg_phase) + const uint16_t *pRt2_kl_ij, const int8_t *efg_phase + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { +#ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + double *dm_kl_cache = reinterpret_cast(shm_mem); +#else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + extern __shared__ double dm_kl_cache[]; +#endif int64_t *pair_ij_mapping = (int64_t*)bounds.pair_ij_mapping; int64_t *pair_kl_mapping = (int64_t*)bounds.pair_kl_mapping; int bsizex = threadsx * tilex; int bsizey = threadsy * tiley; - int pair_ij0 = blockIdx.x * bsizex; - int pair_kl0 = blockIdx.y * bsizey; + int pair_ij0 = blockIdx_x * bsizex; + int pair_kl0 = blockIdx_y * bsizey; if (q_cond_ij[pair_ij0] + q_cond_kl[pair_kl0] < bounds.cutoff) { return; } - int sq_id = threadIdx.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - int nsq_per_block = blockDim.x; + int sq_id = threadIdx_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + int nsq_per_block = blockDim_x; //assert(nsq_per_block == threadsx * threadsy); int t_id = gout_id * nsq_per_block + sq_id; int tx = sq_id % threadsx; @@ -106,12 +130,11 @@ void pbc_md_j_kernel(RysIntEnvVars envs, JKMatrix jmat, MDBoundsInfo bounds, int npairs_ij = bounds.npairs_ij; int npairs_kl = bounds.npairs_kl; - extern __shared__ double dm_kl_cache[]; double *Rq_cache = dm_kl_cache + nf3kl*bsizey; double *Rp_cache = dm_kl_cache + bsizey*(4+nf3kl); double *gamma_inc = dm_kl_cache + bsizey*(4+nf3kl) + threadsx*4 + sq_id; double *Rt = gamma_inc + (order+1) * nsq_per_block; - uint16_t *Rt2_address = pRt2_kl_ij; + uint16_t *Rt2_address = const_cast(pRt2_kl_ij); if (nf3ij * nf3kl <= RT2_IDX_CACHE_SIZE) { int l4 = bounds.lij + bounds.lkl; int nf3 = (l4 + 1) * (l4 + 2) * (l4 + 3) / 6; @@ -125,7 +148,7 @@ void pbc_md_j_kernel(RysIntEnvVars envs, JKMatrix jmat, MDBoundsInfo bounds, __syncthreads(); for (int n = t_id; n < bsizey; n += threads) { - int pair_kl = blockIdx.y * bsizey + n; + int pair_kl = blockIdx_y * bsizey + n; if (pair_kl < npairs_kl) { int64_t bas_kl = pair_kl_mapping[pair_kl]; int ksh = bas_kl / NBAS_MAX; @@ -157,7 +180,7 @@ void pbc_md_j_kernel(RysIntEnvVars envs, JKMatrix jmat, MDBoundsInfo bounds, int kl = n / tiley; int batch_kl = n - kl * tiley; int sq_kl = ty + batch_kl * threadsy; - int pair_kl = blockIdx.y * bsizey + sq_kl; + int pair_kl = blockIdx_y * bsizey + sq_kl; if (pair_kl < npairs_kl) { int kl_loc0 = pair_kl_loc[pair_kl]; dm_kl_cache[sq_kl+kl*bsizey] = dm[kl_loc0+kl]; @@ -166,7 +189,7 @@ void pbc_md_j_kernel(RysIntEnvVars envs, JKMatrix jmat, MDBoundsInfo bounds, } for (int batch_ij = 0; batch_ij < tilex; ++batch_ij) { - int pair_ij0 = (blockIdx.x * tilex + batch_ij) * threadsx; + int pair_ij0 = (blockIdx_x * tilex + batch_ij) * threadsx; if (pair_ij0 >= npairs_ij) { break; } @@ -202,12 +225,12 @@ void pbc_md_j_kernel(RysIntEnvVars envs, JKMatrix jmat, MDBoundsInfo bounds, vj_ij[n] = 0.; } for (int batch_kl = 0; batch_kl < tiley; ++batch_kl) { - int pair_kl0 = (blockIdx.y * tiley + batch_kl) * threadsy; + int pair_kl0 = (blockIdx_y * tiley + batch_kl) * threadsy; if (pair_kl0 >= npairs_kl) { break; } - if (qd_ij_max[blockIdx.x*tilex+batch_ij] + q_cond_kl[pair_kl0] < bounds.cutoff && - qd_kl_max[blockIdx.y*tiley+batch_kl] + q_cond_ij[pair_ij0] < bounds.cutoff) { + if (qd_ij_max[blockIdx_x*tilex+batch_ij] + q_cond_kl[pair_kl0] < bounds.cutoff && + qd_kl_max[blockIdx_y*tiley+batch_kl] + q_cond_ij[pair_ij0] < bounds.cutoff) { continue; } @@ -366,9 +389,40 @@ int PBC_build_j(double *vj, double *dm, int n_dm, } } int nsq_per_block = threads_ij * threads_kl; - dim3 threads(nsq_per_block, gout_stride); int blocks_ij = (npairs_ij + bsizex - 1) / bsizex; int blocks_kl = (npairs_kl + bsizey - 1) / bsizey; + int dm_size = dm_xyz_size * nimgs_uniq_pair; + + #ifdef USE_SYCL + sycl::range<2> blocks(blocks_kl, blocks_ij); + sycl::range<2> threads(gout_stride, nsq_per_block); + // IMP: SYCL doesnt treat the Rt2_kl_ij, c_Rt2_efg_phase + // pointer arithmetic on host and the obtained pointers are + // not valid on the device. Hence just compute the offset on host + // but obtain the pointer `pRt2_kl_ij` & `efg_phase` in the kernel launch + const int Rt2_kl_ij_syclonly_offset = offset_for_Rt2_idx(lij, lkl); + const int efg_phase_syclonly_offset = offset_for_Rt2_idx(0, lkl); + auto dev_envs = *envs; + for (int i_dm = 0; i_dm < n_dm; ++i_dm) { + JKMatrix jmat = {vj+i_dm*dm_size, NULL, dm+i_dm*dm_size, n_dm, 0, omega}; + if (1){//!pbc_md_j_unrolled(envs, &jmat, &bounds, omega)) { + bounds.qd_ij_max = qd_ij_max + qd_offset_for_threads(npairs_ij, threads_ij); + bounds.qd_kl_max = qd_kl_max + qd_offset_for_threads(npairs_kl, threads_kl); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(buflen, cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + const uint16_t *pRt2_kl_ij = Rt2_kl_ij + Rt2_kl_ij_syclonly_offset; + const int8_t *efg_phase = c_Rt2_efg_phase + efg_phase_syclonly_offset; + pbc_md_j_kernel(dev_envs, jmat, bounds, q_cond_ij, q_cond_kl, + threads_ij, threads_kl, tilex, tiley, + pRt2_kl_ij, efg_phase, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + } + } + #else + dim3 threads(nsq_per_block, gout_stride); dim3 blocks(blocks_ij, blocks_kl); uint16_t *pRt2_kl_ij; int8_t *efg_phase; @@ -376,7 +430,6 @@ int PBC_build_j(double *vj, double *dm, int n_dm, cudaGetSymbolAddress((void**)&efg_phase, c_Rt2_efg_phase); pRt2_kl_ij += offset_for_Rt2_idx(lij, lkl); efg_phase += offset_for_Rt2_idx(0, lkl); - int dm_size = dm_xyz_size * nimgs_uniq_pair; for (int i_dm = 0; i_dm < n_dm; ++i_dm) { JKMatrix jmat = {vj+i_dm*dm_size, NULL, dm+i_dm*dm_size, n_dm, 0, omega}; if (1){//!pbc_md_j_unrolled(envs, &jmat, &bounds, omega)) { @@ -392,6 +445,7 @@ int PBC_build_j(double *vj, double *dm, int n_dm, fprintf(stderr, "CUDA Error in MD_build_j: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-md/unrolled_md_j.cu b/gpu4pyscf/lib/gvhf-md/unrolled_md_j.cu index b73e763a4..33fa4f7c3 100644 --- a/gpu4pyscf/lib/gvhf-md/unrolled_md_j.cu +++ b/gpu4pyscf/lib/gvhf-md/unrolled_md_j.cu @@ -6,6 +6,37 @@ #include "gvhf-md/boys.cu" #include "gvhf-md/md_j.cuh" +#ifdef USE_SYCL + +#define KERNEL_ARGS \ + RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, \ + float *q_cond_ij, float *q_cond_kl, \ + sycl::nd_item<2> &item, double *vj_kl_cache + +#define KERNEL_SETUP() \ + const md_j_index2 blockIdx {(int)item.get_group(1), (int)item.get_group(0)}; \ + const md_j_index2 threadIdx {(int)item.get_local_id(1), (int)item.get_local_id(0)}; \ + int tx = threadIdx.x; \ + int ty = threadIdx.y; \ + int block_x = blockIdx.x; \ + int block_y = blockIdx.y; + +#define LAUNCH_KERNEL(KERNEL, SHM, BLOCKS_IJ, BLOCKS_KL) { \ + auto dev_envs = *envs; auto dev_jk = *jk; auto dev_bounds = *bounds; \ + sycl::range<2> threads(16, 16); \ + sycl::range<2> blocks((npairs_kl + (BLOCKS_KL) - 1) / (BLOCKS_KL), \ + (npairs_ij + (BLOCKS_IJ) - 1) / (BLOCKS_IJ)); \ + sycl_get_queue()->submit([&](sycl::handler &cgh) { \ + sycl::local_accessor local_acc(sycl::range<1>((SHM)+addition_buf), cgh); \ + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), \ + [=](sycl::nd_item<2> item) { \ + KERNEL(dev_envs, dev_jk, dev_bounds, q_cond_ij, q_cond_kl, \ + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); \ + }); \ + }); \ +} + +#else // USE_SYCL #define KERNEL_ARGS \ RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, \ @@ -18,8 +49,17 @@ int block_y = blockIdx.y; \ extern __shared__ double vj_kl_cache[]; -#define LAUNCH_KERNEL(KERNEL, SHMSIZE) \ - KERNEL<<>>(*envs, *jk, *bounds, q_cond_ij, q_cond_kl) +#define LAUNCH_KERNEL(KERNEL, SHM, BLOCKS_IJ, BLOCKS_KL) { \ + dim3 threads(16, 16); \ + dim3 blocks((npairs_ij + (BLOCKS_IJ) - 1) / (BLOCKS_IJ), \ + (npairs_kl + (BLOCKS_KL) - 1) / (BLOCKS_KL), 1); \ + cudaFuncSetAttribute(KERNEL, cudaFuncAttributeMaxDynamicSharedMemorySize, \ + ((SHM)+addition_buf)*sizeof(double)); \ + KERNEL<<>>( \ + *envs, *jk, *bounds, q_cond_ij, q_cond_kl); \ +} + +#endif // USE_SYCL // TILEX=30, TILEY=30 @@ -4668,75 +4708,35 @@ int md_j_unrolled(RysIntEnvVars *envs, JKMatrix *jk, MDBoundsInfo *bounds, addition_buf = 256; } switch (ijkl) { - case 0: { // lij=0, lkl=0, tilex=30, tiley=30 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 479) / 480, (npairs_kl + 479) / 480, 1); - LAUNCH_KERNEL(md_j_0_0, (2992+addition_buf)*sizeof(double)); break; - } break; - case 11: { // lij=1, lkl=0, tilex=48, tiley=23 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 367) / 368, 1); - LAUNCH_KERNEL(md_j_1_0, (2992+addition_buf)*sizeof(double)); break; - } break; - case 12: { // lij=1, lkl=1, tilex=10, tiley=10 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 159) / 160, (npairs_kl + 159) / 160, 1); - LAUNCH_KERNEL(md_j_1_1, (2944+addition_buf)*sizeof(double)); break; - } break; - case 22: { // lij=2, lkl=0, tilex=48, tiley=16 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 255) / 256, 1); - LAUNCH_KERNEL(md_j_2_0, (3040+addition_buf)*sizeof(double)); break; - } break; - case 23: { // lij=2, lkl=1, tilex=48, tiley=30 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 479) / 480, 1); - cudaFuncSetAttribute(md_j_2_1, cudaFuncAttributeMaxDynamicSharedMemorySize, (6112+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_2_1, (6112+addition_buf)*sizeof(double)); break; - } break; - case 24: { // lij=2, lkl=2, tilex=14, tiley=14 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 223) / 224, (npairs_kl + 223) / 224, 1); - cudaFuncSetAttribute(md_j_2_2, cudaFuncAttributeMaxDynamicSharedMemorySize, (5920+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_2_2, (5920+addition_buf)*sizeof(double)); break; - } break; - case 33: { // lij=3, lkl=0, tilex=48, tiley=46 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 735) / 736, 1); - cudaFuncSetAttribute(md_j_3_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6112+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_3_0, (6112+addition_buf)*sizeof(double)); break; - } break; - case 34: { // lij=3, lkl=1, tilex=48, tiley=24 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 383) / 384, 1); - cudaFuncSetAttribute(md_j_3_1, cudaFuncAttributeMaxDynamicSharedMemorySize, (6016+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_3_1, (6016+addition_buf)*sizeof(double)); break; - } break; - case 35: { // lij=3, lkl=2, tilex=48, tiley=11 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 175) / 176, 1); - cudaFuncSetAttribute(md_j_3_2, cudaFuncAttributeMaxDynamicSharedMemorySize, (5920+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_3_2, (5920+addition_buf)*sizeof(double)); break; - } break; - case 44: { // lij=4, lkl=0, tilex=48, tiley=36 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 575) / 576, 1); - cudaFuncSetAttribute(md_j_4_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6064+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4_0, (6064+addition_buf)*sizeof(double)); break; - } break; - case 45: { // lij=4, lkl=1, tilex=48, tiley=18 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 287) / 288, 1); - cudaFuncSetAttribute(md_j_4_1, cudaFuncAttributeMaxDynamicSharedMemorySize, (6000+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4_1, (6000+addition_buf)*sizeof(double)); break; - } break; - case 55: { // lij=5, lkl=0, tilex=48, tiley=26 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 415) / 416, 1); - cudaFuncSetAttribute(md_j_5_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6112+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_5_0, (6112+addition_buf)*sizeof(double)); break; - } break; + case 0: // lij=0, lkl=0, tilex=30, tiley=30 + LAUNCH_KERNEL(md_j_0_0, 2992, 480, 480) break; + case 11: // lij=1, lkl=0, tilex=48, tiley=23 + LAUNCH_KERNEL(md_j_1_0, 2992, 768, 368) break; + case 12: // lij=1, lkl=1, tilex=10, tiley=10 + LAUNCH_KERNEL(md_j_1_1, 2944, 160, 160) break; + case 22: // lij=2, lkl=0, tilex=48, tiley=16 + LAUNCH_KERNEL(md_j_2_0, 3040, 768, 256) break; + case 23: // lij=2, lkl=1, tilex=48, tiley=30 + LAUNCH_KERNEL(md_j_2_1, 6112, 768, 480) break; + case 24: // lij=2, lkl=2, tilex=14, tiley=14 + LAUNCH_KERNEL(md_j_2_2, 5920, 224, 224) break; + case 33: // lij=3, lkl=0, tilex=48, tiley=46 + LAUNCH_KERNEL(md_j_3_0, 6112, 768, 736) break; + case 34: // lij=3, lkl=1, tilex=48, tiley=24 + LAUNCH_KERNEL(md_j_3_1, 6016, 768, 384) break; + case 35: // lij=3, lkl=2, tilex=48, tiley=11 + LAUNCH_KERNEL(md_j_3_2, 5920, 768, 176) break; + case 44: // lij=4, lkl=0, tilex=48, tiley=36 + LAUNCH_KERNEL(md_j_4_0, 6064, 768, 576) break; + case 45: // lij=4, lkl=1, tilex=48, tiley=18 + LAUNCH_KERNEL(md_j_4_1, 6000, 768, 288) break; + case 55: // lij=5, lkl=0, tilex=48, tiley=26 + LAUNCH_KERNEL(md_j_5_0, 6112, 768, 416) break; default: return 0; } return 1; } + +#undef LAUNCH_KERNEL +#undef KERNEL_SETUP +#undef KERNEL_ARGS diff --git a/gpu4pyscf/lib/gvhf-md/unrolled_md_j_4dm.cu b/gpu4pyscf/lib/gvhf-md/unrolled_md_j_4dm.cu index 18ee25dec..313dee884 100644 --- a/gpu4pyscf/lib/gvhf-md/unrolled_md_j_4dm.cu +++ b/gpu4pyscf/lib/gvhf-md/unrolled_md_j_4dm.cu @@ -6,6 +6,38 @@ #include "gvhf-md/boys.cu" #include "gvhf-md/md_j.cuh" +#ifdef USE_SYCL + +#define KERNEL_ARGS \ + RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, \ + float *q_cond_ij, float *q_cond_kl, int dm_size, \ + sycl::nd_item<2> &item, double *vj_kl_cache + +#define KERNEL_SETUP() \ + const md_j_index2 blockIdx {(int)item.get_group(1), (int)item.get_group(0)}; \ + const md_j_index2 threadIdx {(int)item.get_local_id(1), (int)item.get_local_id(0)}; \ + int tx = threadIdx.x; \ + int ty = threadIdx.y; \ + int block_x = blockIdx.x; \ + int block_y = blockIdx.y; + +#define LAUNCH_KERNEL(KERNEL, SHM, BLOCKS_IJ, BLOCKS_KL) { \ + auto dev_envs = *envs; auto dev_jk = *jk; auto dev_bounds = *bounds; \ + sycl::range<2> threads(16, 16); \ + sycl::range<2> blocks((npairs_kl + (BLOCKS_KL) - 1) / (BLOCKS_KL), \ + (npairs_ij + (BLOCKS_IJ) - 1) / (BLOCKS_IJ)); \ + sycl_get_queue()->submit([&](sycl::handler &cgh) { \ + sycl::local_accessor local_acc(sycl::range<1>((SHM)+addition_buf), cgh); \ + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), \ + [=](sycl::nd_item<2> item) { \ + KERNEL(dev_envs, dev_jk, dev_bounds, q_cond_ij, q_cond_kl, dm_size, \ + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); \ + }); \ + }); \ +} + +#else // USE_SYCL + #define KERNEL_ARGS \ RysIntEnvVars envs, JKMatrix jk, MDBoundsInfo bounds, \ float *q_cond_ij, float *q_cond_kl, int dm_size @@ -17,8 +49,17 @@ int block_y = blockIdx.y; \ extern __shared__ double vj_kl_cache[]; -#define LAUNCH_KERNEL(KERNEL, SHMSIZE) \ - KERNEL<<>>(*envs, *jk, *bounds, q_cond_ij, q_cond_kl, dm_size) +#define LAUNCH_KERNEL(KERNEL, SHM, BLOCKS_IJ, BLOCKS_KL) { \ + dim3 threads(16, 16); \ + dim3 blocks((npairs_ij + (BLOCKS_IJ) - 1) / (BLOCKS_IJ), \ + (npairs_kl + (BLOCKS_KL) - 1) / (BLOCKS_KL), 1); \ + cudaFuncSetAttribute(KERNEL, cudaFuncAttributeMaxDynamicSharedMemorySize, \ + ((SHM)+addition_buf)*sizeof(double)); \ + KERNEL<<>>( \ + *envs, *jk, *bounds, q_cond_ij, q_cond_kl, dm_size); \ +} + +#endif // USE_SYCL // TILEX=21, TILEY=21 @@ -11007,71 +11048,33 @@ int md_j_4dm_unrolled(RysIntEnvVars *envs, JKMatrix *jk, MDBoundsInfo *bounds, addition_buf = 256; } switch (ijkl) { - case 0: { // lij=0, lkl=0, tilex=21, tiley=21 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 335) / 336, (npairs_kl + 335) / 336, 1); - cudaFuncSetAttribute(md_j_4dm_0_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6080+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_0_0, (6080+addition_buf)*sizeof(double)); break; - } break; - case 9: { // lij=1, lkl=0, tilex=48, tiley=21 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 335) / 336, 1); - cudaFuncSetAttribute(md_j_4dm_1_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6080+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_1_0, (6080+addition_buf)*sizeof(double)); break; - } break; - case 10: { // lij=1, lkl=1, tilex=6, tiley=6 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 95) / 96, (npairs_kl + 95) / 96, 1); - LAUNCH_KERNEL(md_j_4dm_1_1, (5568+addition_buf)*sizeof(double)); break; - } break; - case 18: { // lij=2, lkl=0, tilex=48, tiley=16 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 255) / 256, 1); - cudaFuncSetAttribute(md_j_4dm_2_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (5952+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_2_0, (5952+addition_buf)*sizeof(double)); break; - } break; - case 19: { // lij=2, lkl=1, tilex=48, tiley=10 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 159) / 160, 1); - cudaFuncSetAttribute(md_j_4dm_2_1, cudaFuncAttributeMaxDynamicSharedMemorySize, (5952+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_2_1, (5952+addition_buf)*sizeof(double)); break; - } break; - case 20: { // lij=2, lkl=2, tilex=4, tiley=4 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 63) / 64, (npairs_kl + 63) / 64, 1); - cudaFuncSetAttribute(md_j_4dm_2_2, cudaFuncAttributeMaxDynamicSharedMemorySize, (6080+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_2_2, (6080+addition_buf)*sizeof(double)); break; - } break; - case 27: { // lij=3, lkl=0, tilex=48, tiley=21 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 335) / 336, 1); - cudaFuncSetAttribute(md_j_4dm_3_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6080+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_3_0, (6080+addition_buf)*sizeof(double)); break; - } break; - case 28: { // lij=3, lkl=1, tilex=48, tiley=6 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 95) / 96, 1); - LAUNCH_KERNEL(md_j_4dm_3_1, (5824+addition_buf)*sizeof(double)); break; - } break; - case 36: { // lij=4, lkl=0, tilex=48, tiley=24 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 383) / 384, 1); - cudaFuncSetAttribute(md_j_4dm_4_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6048+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_4_0, (6048+addition_buf)*sizeof(double)); break; - } break; - case 37: { // lij=4, lkl=1, tilex=48, tiley=9 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 143) / 144, 1); - cudaFuncSetAttribute(md_j_4dm_4_1, cudaFuncAttributeMaxDynamicSharedMemorySize, (5984+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_4_1, (5984+addition_buf)*sizeof(double)); break; - } break; - case 45: { // lij=5, lkl=0, tilex=48, tiley=12 - dim3 threads(16, 16); - dim3 blocks((npairs_ij + 767) / 768, (npairs_kl + 191) / 192, 1); - cudaFuncSetAttribute(md_j_4dm_5_0, cudaFuncAttributeMaxDynamicSharedMemorySize, (6080+addition_buf)*sizeof(double)); - LAUNCH_KERNEL(md_j_4dm_5_0, (6080+addition_buf)*sizeof(double)); break; - } break; + case 0: // lij=0, lkl=0, tilex=21, tiley=21 + LAUNCH_KERNEL(md_j_4dm_0_0, 6080, 336, 336) break; + case 9: // lij=1, lkl=0, tilex=48, tiley=21 + LAUNCH_KERNEL(md_j_4dm_1_0, 6080, 768, 336) break; + case 10: // lij=1, lkl=1, tilex=6, tiley=6 + LAUNCH_KERNEL(md_j_4dm_1_1, 5568, 96, 96) break; + case 18: // lij=2, lkl=0, tilex=48, tiley=16 + LAUNCH_KERNEL(md_j_4dm_2_0, 5952, 768, 256) break; + case 19: // lij=2, lkl=1, tilex=48, tiley=10 + LAUNCH_KERNEL(md_j_4dm_2_1, 5952, 768, 160) break; + case 20: // lij=2, lkl=2, tilex=4, tiley=4 + LAUNCH_KERNEL(md_j_4dm_2_2, 6080, 64, 64) break; + case 27: // lij=3, lkl=0, tilex=48, tiley=21 + LAUNCH_KERNEL(md_j_4dm_3_0, 6080, 768, 336) break; + case 28: // lij=3, lkl=1, tilex=48, tiley=6 + LAUNCH_KERNEL(md_j_4dm_3_1, 5824, 768, 96) break; + case 36: // lij=4, lkl=0, tilex=48, tiley=24 + LAUNCH_KERNEL(md_j_4dm_4_0, 6048, 768, 384) break; + case 37: // lij=4, lkl=1, tilex=48, tiley=9 + LAUNCH_KERNEL(md_j_4dm_4_1, 5984, 768, 144) break; + case 45: // lij=5, lkl=0, tilex=48, tiley=12 + LAUNCH_KERNEL(md_j_4dm_5_0, 6080, 768, 192) break; default: return 0; } return 1; } + +#undef LAUNCH_KERNEL +#undef KERNEL_SETUP +#undef KERNEL_ARGS diff --git a/gpu4pyscf/lib/gvhf-rys/CMakeLists.txt b/gpu4pyscf/lib/gvhf-rys/CMakeLists.txt index 400edaa90..ae592978b 100644 --- a/gpu4pyscf/lib/gvhf-rys/CMakeLists.txt +++ b/gpu4pyscf/lib/gvhf-rys/CMakeLists.txt @@ -1,10 +1,7 @@ -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128") - -add_library(gvhf_rys SHARED +set(GPU_SRCS rys_jk_driver.cu rys_roots_dat.cu nr_sr_estimator.c nr_sr_estimator.cu rys_contract_j.cu cart2xyz.c #unrolled_rys_j.cu - rys_constant.cu rys_contract_k.cu unrolled_rys_k.cu rys_contract_jk.cu unrolled_rys_jk.cu rys_contract_jk_ip1.cu unrolled_rys_jk_ip1.cu unrolled_ejk_ip1.cu @@ -18,14 +15,69 @@ add_library(gvhf_rys SHARED e_int2c2e_ip2.cu ) +if (USE_SYCL) + add_library(gvhf_rys SHARED ${GPU_SRCS}) + + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + + # Per-TU SYCL kernel-name disambiguation. + # + # unrolled_rys_jk.cu and unrolled_rys_k.cu both define 19 kernels named + # rys_k_0000 .. rys_k_3200 with different bodies. SYCL derives the kernel + # identity from the name-class type, whose host-side registry symbols are + # vague-linkage and collapse at link time -- both launch sites would then + # dispatch to the same body, silently, with no diagnostic. + # + # unrolled_kernels.cuh appends RYS_UNROLLED_KERNEL_TAG to every kernel + # name class and #errors if it is missing. The tag MUST be injected here + # rather than in the .cu files: those are auto-generated upstream and must + # stay byte-identical to upstream/master. + # + # Every file that includes unrolled_kernels.cuh needs a UNIQUE tag. + set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/unrolled_rys_jk.cu + PROPERTIES COMPILE_DEFINITIONS "RYS_UNROLLED_KERNEL_TAG=jk") + set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/unrolled_rys_k.cu + PROPERTIES COMPILE_DEFINITIONS "RYS_UNROLLED_KERNEL_TAG=k") + set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/unrolled_rys_jk_ip1.cu + PROPERTIES COMPILE_DEFINITIONS "RYS_UNROLLED_KERNEL_TAG=vjk_ip1") + set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/unrolled_ejk_ip1.cu + PROPERTIES COMPILE_DEFINITIONS "RYS_UNROLLED_KERNEL_TAG=ejk_ip1") + set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/unrolled_ejk_ip2_type12.cu + PROPERTIES COMPILE_DEFINITIONS "RYS_UNROLLED_KERNEL_TAG=ejk_ip2_t12") + set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/unrolled_ejk_ip2_type3.cu + PROPERTIES COMPILE_DEFINITIONS "RYS_UNROLLED_KERNEL_TAG=ejk_ip2_t3") + + # Note: this is needed in SYCL to prevent compiling it as a + # a seperate TU and cause multi-definition error from all the + # `#inlcudes <>` of this file + set_source_files_properties( + ${CMAKE_CURRENT_SOURCE_DIR}/rys_roots_for_k.cu + ${CMAKE_CURRENT_SOURCE_DIR}/rys_roots.cu + ${CMAKE_CURRENT_SOURCE_DIR}/rys_contract_k.cuh + ${CMAKE_CURRENT_SOURCE_DIR}/vhf.cuh + ${CMAKE_CURRENT_SOURCE_DIR}/rys_roots.cuh + PROPERTIES HEADER_FILE_ONLY ON + ) + + set_target_properties(gvhf_rys PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + target_compile_options(gvhf_rys PRIVATE -x c++ -nocudainc -nocudalib) + target_link_libraries(gvhf_rys PRIVATE sycl_compat gsycl) +else() + list(APPEND GPU_SRCS rys_roots_dat.cu rys_constant.cu) + add_library(gvhf_rys SHARED ${GPU_SRCS}) + + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v") + + set_target_properties(gvhf_rys PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} + CUDA_SEPARABLE_COMPILATION ON) +endif() + #option(BUILD_SHARED_LIBS "build shared libraries" 1) #option(ENABLE_STATIC "Enforce static library build" 0) #if(ENABLE_STATIC) # set(BUILD_SHARED_LIBS 0) #endif() -set_target_properties(gvhf_rys PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_SEPARABLE_COMPILATION ON) - -target_link_libraries(gvhf_rys OpenMP::OpenMP_C) +target_link_libraries(gvhf_rys PRIVATE OpenMP::OpenMP_C) diff --git a/gpu4pyscf/lib/gvhf-rys/cart2xyz.c b/gpu4pyscf/lib/gvhf-rys/cart2xyz.c index acc611f26..30fc5c1b7 100644 --- a/gpu4pyscf/lib/gvhf-rys/cart2xyz.c +++ b/gpu4pyscf/lib/gvhf-rys/cart2xyz.c @@ -152,6 +152,9 @@ static void _dm_xyz_to_dm(double* dm_xyz, double* dm, int nao, int li, int lj, d } } +#ifdef __cplusplus +extern "C" { +#endif void transform_cart_to_xyz(double *dm_xyz, double *dm, int *ao_loc, int *pair_loc, int *bas, int nbas, double *env) { @@ -184,8 +187,14 @@ void transform_cart_to_xyz(double *dm_xyz, double *dm, int *ao_loc, int *pair_lo } } } +#ifdef __cplusplus +} // extern "C" +#endif +#ifdef __cplusplus +extern "C" { +#endif void transform_xyz_to_cart(double *vj, double *vj_xyz, int *ao_loc, int *pair_loc, int *bas, int nbas, double *env) { @@ -218,3 +227,6 @@ void transform_xyz_to_cart(double *vj, double *vj_xyz, int *ao_loc, int *pair_lo } } } +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/gpu4pyscf/lib/gvhf-rys/contract_int3c2e.cu b/gpu4pyscf/lib/gvhf-rys/contract_int3c2e.cu index d5f65d9ad..a2ce54fa8 100644 --- a/gpu4pyscf/lib/gvhf-rys/contract_int3c2e.cu +++ b/gpu4pyscf/lib/gvhf-rys/contract_int3c2e.cu @@ -117,15 +117,44 @@ void _store_vj(double *out, double *vj_cache, int li, int lj, size_t nao, __global__ static void contract_int3c2e_dm_kernel(double *out, double *dm, int n_dm, int naux, RysIntEnvVars envs, int *shl_pair_offsets, - uint32_t *bas_ij_idx, int *gout_stride_lookup) + uint32_t *bas_ij_idx, int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - int thread_id = threadIdx.x; - int nbas = envs.nbas; - int ksh = blockIdx.x + nbas; - int *bas = envs.bas; - int *ao_loc = envs.ao_loc; - double *env = envs.env; - double omega = env[PTR_RANGE_OMEGA]; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int gridDim_y = item.get_group_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &dm_id0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &xk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &yk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &zk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ck = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int gridDim_y = gridDim.y; + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1; __shared__ int li, lj, lk, nroots; @@ -135,8 +164,17 @@ void contract_int3c2e_dm_kernel(double *out, double *dm, int n_dm, int naux, __shared__ int dm_id0; __shared__ double xk, yk, zk; __shared__ int expk, ck; + #endif + + int thread_id = threadIdx_x; + int nbas = envs.nbas; + int ksh = blockIdx_x + nbas; + int *bas = envs.bas; + int *ao_loc = envs.ao_loc; + double *env = envs.env; + double omega = env[PTR_RANGE_OMEGA]; if (thread_id == 0) { - int sp_block_id = gridDim.y - blockIdx.y - 1; + int sp_block_id = gridDim_y - blockIdx_y - 1; shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; int bas_ij0 = bas_ij_idx[shl_pair0]; @@ -356,23 +394,58 @@ __global__ static void contract_int3c2e_auxvec_kernel(double *out, double *auxvec, int n_dm, int naux, RysIntEnvVars envs, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *ksh_offsets, - int *gout_stride_lookup) + int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { // For better load balance, consume blocks in the reversed order - int thread_id = threadIdx.x; - int nbas = envs.nbas; - int *bas = envs.bas; - double *env = envs.env; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1; __shared__ int ksh0, ksh1; __shared__ int li, lj, lk, nroots; __shared__ int iprim, jprim, kprim; __shared__ int gout_stride, nst_per_block; + #endif + + int thread_id = threadIdx_x; + int nbas = envs.nbas; + int *bas = envs.bas; + double *env = envs.env; double omega = env[PTR_RANGE_OMEGA]; if (thread_id == 0) { - int sp_block_id = gridDim.x - blockIdx.x - 1; - int ksh_block_id = gridDim.y - blockIdx.y - 1; + int sp_block_id = gridDim_x - blockIdx_x - 1; + int ksh_block_id = gridDim_y - blockIdx_y - 1; ksh0 = ksh_offsets[ksh_block_id]; ksh1 = ksh_offsets[ksh_block_id+1]; shl_pair0 = shl_pair_offsets[sp_block_id]; @@ -561,6 +634,19 @@ int contract_int3c2e_dm(double *out, double *dm, int n_dm, int naux, int nbas_aux, int nbatches_shl_pair, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *gout_stride_lookup) { +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_shl_pair, nbas_aux); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + contract_int3c2e_dm_kernel( + out, dm, n_dm, naux, dev_envs, shl_pair_offsets, bas_ij_idx, gout_stride_lookup, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else cudaFuncSetAttribute(contract_int3c2e_dm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 blocks(nbas_aux, nbatches_shl_pair); contract_int3c2e_dm_kernel<<>>( @@ -570,6 +656,7 @@ int contract_int3c2e_dm(double *out, double *dm, int n_dm, int naux, fprintf(stderr, "CUDA Error in contract_int3c2e_dm: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -580,6 +667,20 @@ int contract_int3c2e_auxvec(double *vj, double *auxvec, int n_dm, int naux, int *shl_pair_offsets, int *ksh_offsets, uint32_t *bas_ij_idx, int *gout_stride_lookup) { +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_ksh, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + contract_int3c2e_auxvec_kernel( + vj, auxvec, n_dm, naux, dev_envs, shl_pair_offsets, bas_ij_idx, ksh_offsets, + gout_stride_lookup, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else cudaFuncSetAttribute(contract_int3c2e_auxvec_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 blocks(nbatches_shl_pair, nbatches_ksh); contract_int3c2e_auxvec_kernel<<>>( @@ -590,6 +691,8 @@ int contract_int3c2e_auxvec(double *vj, double *auxvec, int n_dm, int naux, fprintf(stderr, "CUDA Error in contract_int3c2e_auxvec, error message = %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } + } diff --git a/gpu4pyscf/lib/gvhf-rys/create_tasks.cu b/gpu4pyscf/lib/gvhf-rys/create_tasks.cu index 94821df2d..d4e4d18ed 100644 --- a/gpu4pyscf/lib/gvhf-rys/create_tasks.cu +++ b/gpu4pyscf/lib/gvhf-rys/create_tasks.cu @@ -24,10 +24,32 @@ #define Q_COND_MARGIN 4.f +#ifdef USE_SYCL + +#define KERNEL_SETUP() \ +auto item = syclex::this_work_item::get_nd_item<2>(); \ +int threadIdx_x = item.get_local_id(1); \ +int threadIdx_y = item.get_local_id(0); \ +int blockDim_x = item.get_local_range(1); \ +int blockDim_y = item.get_local_range(0); + +#else // USE_SYCL + +#define KERNEL_SETUP() \ + int threadIdx_x = threadIdx.x; \ + int threadIdx_y = threadIdx.y; \ + int blockDim_x = blockDim.x; \ + int blockDim_y = blockDim.y; + +#endif // USE_SYCL + // np.where(threads_mask)[0] __device__ inline int mask_to_index(int keep, int *tmp_storage, int threads, int t_id) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + #endif tmp_storage[t_id] = keep; __syncthreads(); for (int offset = 1; offset < threads; offset <<= 1) { @@ -50,8 +72,9 @@ void _fill_vk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, int *swap, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -61,6 +84,7 @@ void _fill_vk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -104,7 +128,7 @@ void _fill_vk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, __syncthreads(); } // pad data to avoid overflow - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -117,8 +141,9 @@ void _fill_vjk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, int *swap, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -128,6 +153,7 @@ void _fill_vjk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -173,7 +199,8 @@ void _fill_vjk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -186,8 +213,9 @@ void _fill_vj_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, int *swap, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -197,6 +225,7 @@ void _fill_vj_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -236,7 +265,7 @@ void _fill_vj_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -249,8 +278,9 @@ void _fill_sr_vk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, int *swap, double omega, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -260,6 +290,7 @@ void _fill_sr_vk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -360,7 +391,7 @@ void _fill_sr_vk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -373,8 +404,9 @@ void _fill_sr_vjk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, int *swap, double omega, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -384,6 +416,7 @@ void _fill_sr_vjk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -487,7 +520,8 @@ void _fill_sr_vjk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -500,8 +534,9 @@ void _fill_sr_vj_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, int *swap, double omega, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -511,6 +546,7 @@ void _fill_sr_vj_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -609,7 +645,8 @@ void _fill_sr_vj_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -622,8 +659,9 @@ void _fill_vjk_tasks_nosym(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, int *swap, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -633,6 +671,7 @@ void _fill_vjk_tasks_nosym(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -677,7 +716,8 @@ void _fill_vjk_tasks_nosym(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -690,8 +730,9 @@ void _fill_sr_vjk_tasks_nosym(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, int *swap, double omega, RysIntEnvVars &envs, BoundsInfo &bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -701,6 +742,7 @@ void _fill_sr_vjk_tasks_nosym(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + dm_penalty + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -811,7 +853,7 @@ void _fill_sr_vjk_tasks_nosym(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -822,10 +864,11 @@ static void _fill_ejk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, int pair_ij, int ish, int jsh, float *q_cond_ij, float *q_cond_kl, int *swap, - JKEnergy &jk, RysIntEnvVars envs, BoundsInfo bounds) + JKEnergy jk, RysIntEnvVars envs, BoundsInfo bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -835,6 +878,7 @@ static void _fill_ejk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -879,7 +923,8 @@ static void _fill_ejk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -891,10 +936,11 @@ static void _fill_sr_ejk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float *q_cond_ij, float *q_cond_kl, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, int *swap, - JKEnergy &jk, RysIntEnvVars envs, BoundsInfo bounds) + JKEnergy jk, RysIntEnvVars envs, BoundsInfo bounds) { - int t_id = threadIdx.y * blockDim.x + threadIdx.x; - int threads = blockDim.x * blockDim.y; + KERNEL_SETUP(); + int t_id = threadIdx_y * blockDim_x + threadIdx_x; + int threads = blockDim_x * blockDim_y; __syncthreads(); if (t_id == 0) { ntasks = 0; @@ -904,6 +950,7 @@ static void _fill_sr_ejk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, float q_ij = q_cond_ij[pair_ij]; float kl_cutoff = cutoff - q_ij; if (q_cond_kl[pair_kl0] + Q_COND_MARGIN < kl_cutoff) { + __syncthreads(); return; } @@ -1010,8 +1057,14 @@ static void _fill_sr_ejk_tasks(int& ntasks, int& pair_kl0, uint32_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { + + if (threadIdx_y == 0 && ntasks + t_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+t_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); } + +// KERNEL_SETUP is local to the task-filling helpers above. Undefine it so that +// translation units which #include this file can define their own KERNEL_SETUP +// (e.g. the unrolled_*.cu kernels) without triggering -Wmacro-redefined. +#undef KERNEL_SETUP diff --git a/gpu4pyscf/lib/gvhf-rys/decompress.cu b/gpu4pyscf/lib/gvhf-rys/decompress.cu index 4741270fc..1b091abd1 100644 --- a/gpu4pyscf/lib/gvhf-rys/decompress.cu +++ b/gpu4pyscf/lib/gvhf-rys/decompress.cu @@ -28,9 +28,16 @@ __global__ static void write_kernel(double *out, double *inp, size_t ncol, int col0, int col1) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int threads = item.get_local_range(0); + size_t row = item.get_group(0); +#else int thread_id = threadIdx.x; int threads = blockDim.x; size_t row = blockIdx.x; +#endif int dcol = col1 - col0; out += row * ncol + col0; inp += row * dcol; @@ -43,9 +50,16 @@ __global__ static void transpose_write_kernel(double *out, double *inp, size_t nrow, size_t ncol, int col0, int col1) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int threads = item.get_local_range(0); + int row_id = item.get_group(0); +#else int thread_id = threadIdx.x; int threads = blockDim.x; int row_id = blockIdx.x; +#endif int dcol = col1 - col0; out = out + row_id * ncol + col0; for (int k = thread_id; k < dcol; k += threads) { @@ -56,6 +70,17 @@ void transpose_write_kernel(double *out, double *inp, size_t nrow, size_t ncol, extern "C" { int store_col_segment(double *out_cpu, double *inp, int nrow, int ncol, int col0, int col1) { +#ifdef USE_SYCL + // Host USM allocations are directly device-accessible; no address mapping. + double *out_gpu = out_cpu; + size_t Ncol = ncol; + sycl::range<1> threads(512); + sycl::range<1> blocks(nrow); + sycl_get_queue()->parallel_for( + sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + write_kernel(out_gpu, inp, Ncol, col0, col1); + }); +#else double *out_gpu; cudaError_t err = cudaHostGetDevicePointer(&out_gpu, out_cpu, 0); if(err != cudaSuccess){ @@ -68,11 +93,24 @@ int store_col_segment(double *out_cpu, double *inp, int nrow, int ncol, int col0 fprintf(stderr, "store_col_segment error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } int transpose_write(double *out_cpu, double *inp, int nrow, int ncol, int col0, int col1) { +#ifdef USE_SYCL + // Host USM allocations are directly device-accessible; no address mapping. + double *out_gpu = out_cpu; + size_t Nrow = nrow; + size_t Ncol = ncol; + sycl::range<1> threads(512); + sycl::range<1> blocks(nrow); + sycl_get_queue()->parallel_for( + sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + transpose_write_kernel(out_gpu, inp, Nrow, Ncol, col0, col1); + }); +#else double *out_gpu; cudaError_t err = cudaHostGetDevicePointer(&out_gpu, out_cpu, 0); if(err != cudaSuccess){ @@ -85,6 +123,7 @@ int transpose_write(double *out_cpu, double *inp, int nrow, int ncol, int col0, fprintf(stderr, "transpose_write error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/e_int2c2e_ip2.cu b/gpu4pyscf/lib/gvhf-rys/e_int2c2e_ip2.cu index cf56605b5..dc6dbd8b7 100644 --- a/gpu4pyscf/lib/gvhf-rys/e_int2c2e_ip2.cu +++ b/gpu4pyscf/lib/gvhf-rys/e_int2c2e_ip2.cu @@ -33,19 +33,47 @@ __global__ static void e_int2c2e_ip2_kernel(double *out, double *dm, PBCIntEnvVars envs, double omega, double lr_factor, double sr_factor, int *shl_pair_offsets, uint32_t *bas_ij_idx, - int *gout_stride_lookup) + int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<1> &item, char *shm_mem + #endif + ) { + #ifdef USE_SYCL + int sp_block_id = item.get_group(0); + int thread_id = item.get_local_id(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else int sp_block_id = blockIdx.x; int thread_id = threadIdx.x; - int *bas = envs.bas; - double *env = envs.env; - int nbas = envs.nbas; + __shared__ int shl_pair0, shl_pair1; __shared__ int li, lj, nroots; __shared__ int iprim, jprim; __shared__ int g_size; __shared__ int nao; __shared__ int gout_stride, nsp_per_block; + + extern __shared__ double shared_memory[]; + #endif + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -72,7 +100,6 @@ void e_int2c2e_ip2_kernel(double *out, double *dm, PBCIntEnvVars envs, int gout_id = thread_id / nsp_per_block; int gx_len = g_size * nsp_per_block; - extern __shared__ double shared_memory[]; double *rw = shared_memory + sp_id; double *gx = shared_memory + nsp_per_block * nroots*2 + sp_id; double *Rpq = shared_memory + nsp_per_block * (g_size*3+nroots*2) + sp_id; @@ -361,6 +388,19 @@ int e_int2c2e_ip2(double *out, double *dm, PBCIntEnvVars *envs, int nbatches_shl_pair, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *gout_stride_lookup) { + #ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + e_int2c2e_ip2_kernel(out, dm, dev_envs, omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, gout_stride_lookup, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(e_int2c2e_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); e_int2c2e_ip2_kernel<<>>( out, dm, *envs, omega, lr_factor, sr_factor, @@ -370,6 +410,7 @@ int e_int2c2e_ip2(double *out, double *dm, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in int2c2e_ip2 kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip1.cu b/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip1.cu index 9d807c8fe..d18d183e7 100644 --- a/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip1.cu +++ b/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip1.cu @@ -29,7 +29,20 @@ #define DM_BLOCK 7 #define GOUT_WIDTH 54 +// unrolled_ejk_int3c2e_ip1.cu is auto-generated upstream and must stay +// byte-identical. Its kernels use __syncthreads() but never declare an +// nd_item, so swap in an item-free barrier for the duration of the include. +// The kernels are only ever instantiated from 2-D nd_range launches, so +// get_nd_item<2>() is well-formed. +#ifdef USE_SYCL +#pragma push_macro("__syncthreads") +#undef __syncthreads +#define __syncthreads() (sycl::group_barrier(syclex::this_work_item::get_nd_item<2>().get_group())) +#endif #include "unrolled_ejk_int3c2e_ip1.cu" +#ifdef USE_SYCL +#pragma pop_macro("__syncthreads") +#endif __global__ static void sum_ejk_int3c2e_ip1_kernel(double *ejk, double *ejk_aux, @@ -38,12 +51,48 @@ void sum_ejk_int3c2e_ip1_kernel(double *ejk, double *ejk_aux, double omega, double lr_factor, double sr_factor, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *ksh_offsets, int *gout_stride_lookup, - int *ao_pair_loc, int aux_offset, int naux) + int *ao_pair_loc, int aux_offset, int naux + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - // For better load balance, consume blocks in the reversed order - int thread_id = threadIdx.x; - int sp_block_id = gridDim.x - blockIdx.x - 1; - int ksh_block_id = gridDim.y - blockIdx.y - 1; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nksh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &aux_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1; __shared__ int ksh0, ksh1, nksh; @@ -52,7 +101,12 @@ void sum_ejk_int3c2e_ip1_kernel(double *ejk, double *ejk_aux, __shared__ int g_size; __shared__ int nao; __shared__ int gout_stride, nst_per_block, aux_per_block, nsp_per_block; + #endif + // For better load balance, consume blocks in the reversed order + int thread_id = threadIdx_x; + int sp_block_id = gridDim_x - blockIdx_x - 1; + int ksh_block_id = gridDim_y - blockIdx_y - 1; int nbas = envs.nbas; int *bas = envs.bas; double *env = envs.env; @@ -389,12 +443,48 @@ void ejk_int3c2e_ip1_kernel(double *ejk, double *ejk_aux, RysIntEnvVars envs, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *ksh_offsets, int *gout_stride_lookup, - int *ao_pair_loc, int aux_offset, int npairs, int naux) + int *ao_pair_loc, int aux_offset, int npairs, int naux + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - // For better load balance, consume blocks in the reversed order - int thread_id = threadIdx.x; - int sp_block_id = gridDim.x - blockIdx.x - 1; - int ksh_block_id = gridDim.y - blockIdx.y - 1; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nksh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &aux_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1; __shared__ int ksh0, ksh1, nksh; @@ -403,7 +493,12 @@ void ejk_int3c2e_ip1_kernel(double *ejk, double *ejk_aux, __shared__ int g_size; __shared__ int nao; __shared__ int gout_stride, nst_per_block, aux_per_block, nsp_per_block; + #endif + // For better load balance, consume blocks in the reversed order + int thread_id = threadIdx_x; + int sp_block_id = gridDim_x - blockIdx_x - 1; + int ksh_block_id = gridDim_y - blockIdx_y - 1; int nbas = envs.nbas; int *bas = envs.bas; double *env = envs.env; @@ -741,6 +836,22 @@ int sum_ejk_int3c2e_ip1(double *ejk, double *ejk_aux, int *ao_pair_loc, int aux_offset, int nao, int npairs, int naux, int natm) { +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_ksh, nbatches_shl_pair); + auto dev_envs = *envs; + sycl::queue* q = sycl_get_queue(); + q->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + sum_ejk_int3c2e_ip1_kernel(ejk, ejk_aux, dm, density_auxvec, n_dm, dev_envs, + omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, ksh_offsets, gout_stride_lookup, + ao_pair_loc, aux_offset, naux, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else cudaFuncSetAttribute(sum_ejk_int3c2e_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 blocks(nbatches_shl_pair, nbatches_ksh); sum_ejk_int3c2e_ip1_kernel<<>>( @@ -753,6 +864,7 @@ int sum_ejk_int3c2e_ip1(double *ejk, double *ejk_aux, fprintf(stderr, "CUDA Error in ejk_int3c2e_ip1: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -765,9 +877,33 @@ int ejk_int3c2e_ip1(double *ejk, double *ejk_aux, int *ao_pair_loc, int aux_offset, int nao, int npairs, int naux, int natm) { + size_t nao2 = nao * nao; + +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_ksh, nbatches_shl_pair); + auto dev_envs = *envs; + for (int n = 0; n < n_dm; n += DM_BLOCK) { + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ejk_int3c2e_ip1_kernel(ejk+n*natm*3, ejk_aux+n*natm*3, dm, density_auxvec, n_dm-n, + omega, lr_factor, sr_factor, dev_envs, + shl_pair_offsets, bas_ij_idx, ksh_offsets, gout_stride_lookup, + ao_pair_loc, aux_offset, npairs, naux, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + if (density_auxvec == NULL) { // for exchange + dm += DM_BLOCK * (size_t)npairs * naux; + } else { + dm += DM_BLOCK * nao2; + density_auxvec += DM_BLOCK * naux; + } + } +#else cudaFuncSetAttribute(ejk_int3c2e_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 blocks(nbatches_shl_pair, nbatches_ksh); - size_t nao2 = nao * nao; for (int n = 0; n < n_dm; n += DM_BLOCK) { ejk_int3c2e_ip1_kernel<<>>( ejk+n*natm*3, ejk_aux+n*natm*3, dm, density_auxvec, n_dm-n, @@ -786,6 +922,7 @@ int ejk_int3c2e_ip1(double *ejk, double *ejk_aux, fprintf(stderr, "CUDA Error in ejk_int3c2e_ip1: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip2.cu b/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip2.cu index 2be4a7d6f..74647b185 100644 --- a/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip2.cu +++ b/gpu4pyscf/lib/gvhf-rys/ejk_int3c2e_ip2.cu @@ -17,8 +17,10 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "vhf.cuh" #include "gvhf-rys/rys_roots.cu" #include "gvhf-rys/rys_contract_k.cuh" @@ -33,15 +35,49 @@ void ejk_int3c2e_ip2_kernel(double *ejk, double *dm, double *density_auxvec, double omega, double lr_factor, double sr_factor, RysIntEnvVars envs, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *ksh_offsets, int *gout_stride_lookup, - int *ao_pair_loc, int aux_offset, int naux) + int *ao_pair_loc, int aux_offset, int naux + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - // For better load balance, consume blocks in the reversed order - int thread_id = threadIdx.x; - int sp_block_id = gridDim.x - blockIdx.x - 1; - int ksh_block_id = gridDim.y - blockIdx.y - 1; - int nbas = envs.nbas; - int *bas = envs.bas; - double *env = envs.env; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nksh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &aux_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1; __shared__ int ksh0, ksh1, nksh; __shared__ int li, lj, lk, nroots, nf; @@ -49,6 +85,15 @@ void ejk_int3c2e_ip2_kernel(double *ejk, double *dm, double *density_auxvec, __shared__ int g_size; __shared__ int nao; __shared__ int gout_stride, nst_per_block, aux_per_block, nsp_per_block; + #endif + + // For better load balance, consume blocks in the reversed order + int thread_id = threadIdx_x; + int sp_block_id = gridDim_x - blockIdx_x - 1; + int ksh_block_id = gridDim_y - blockIdx_y - 1; + int nbas = envs.nbas; + int *bas = envs.bas; + double *env = envs.env; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -90,7 +135,6 @@ void ejk_int3c2e_ip2_kernel(double *ejk, double *dm, double *density_auxvec, register int aux_id = st_id - sp_id * aux_per_block; int gx_len = g_size * nst_per_block; - extern __shared__ double shared_memory[]; double *rjri = shared_memory + sp_id; double *Rpq = shared_memory + nsp_per_block * 3 + st_id; double *gx = shared_memory + nst_per_block * 6 + st_id; @@ -535,6 +579,21 @@ int ejk_int3c2e_ip2(double *ejk, double *dm, double *density_auxvec, int *ksh_offsets, int *gout_stride_lookup, int *ao_pair_loc, int aux_offset, int naux) { +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_ksh, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ejk_int3c2e_ip2_kernel( + ejk, dm, density_auxvec, omega, lr_factor, sr_factor, dev_envs, + shl_pair_offsets, bas_ij_idx, ksh_offsets, + gout_stride_lookup, ao_pair_loc, aux_offset, naux, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else cudaFuncSetAttribute(ejk_int3c2e_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 blocks(nbatches_shl_pair, nbatches_ksh); ejk_int3c2e_ip2_kernel<<>>( @@ -546,6 +605,7 @@ int ejk_int3c2e_ip2(double *ejk, double *dm, double *density_auxvec, fprintf(stderr, "CUDA Error in ejk_int3c2e_ip2: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/fill_int3c2e.cu b/gpu4pyscf/lib/gvhf-rys/fill_int3c2e.cu index 6c3713c18..26618628b 100644 --- a/gpu4pyscf/lib/gvhf-rys/fill_int3c2e.cu +++ b/gpu4pyscf/lib/gvhf-rys/fill_int3c2e.cu @@ -22,7 +22,19 @@ #include "gvhf-rys/vhf.cuh" #include "gvhf-rys/rys_roots_for_k.cu" #include "gvhf-rys/rys_contract_k.cuh" +// unrolled_int3c2e.cu is auto-generated upstream and must stay byte-identical. +// Its kernels use __syncthreads() but never declare an nd_item, so swap in an +// item-free barrier for the duration of the include. The kernels are only ever +// instantiated from 2-D nd_range launches, so get_nd_item<2>() is well-formed. +#ifdef USE_SYCL +#pragma push_macro("__syncthreads") +#undef __syncthreads +#define __syncthreads() (sycl::group_barrier(syclex::this_work_item::get_nd_item<2>().get_group())) +#endif #include "unrolled_int3c2e.cu" +#ifdef USE_SYCL +#pragma pop_macro("__syncthreads") +#endif #include "build_rys_gxyz.cuh" #define THREADS 256 @@ -36,10 +48,42 @@ void int3c2e_kernel(double *out, RysIntEnvVars envs, double *pool, int *ksh_offsets, int *gout_stride_lookup, int *ao_pair_loc, int ao_pair_offset, int aux_offset, int naux, int reorder_aux, int to_sph, - int *head, int nbatches_shl_pair, int nbatches_ksh) + int *head, int nbatches_shl_pair, int nbatches_ksh + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - int thread_id = threadIdx.x; - int worker_id = blockIdx.x; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nksp = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &aux_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &sp_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1, nksp; __shared__ int ksh0, ksh1; @@ -49,6 +93,10 @@ void int3c2e_kernel(double *out, RysIntEnvVars envs, double *pool, __shared__ int g_size; __shared__ int gout_stride, nst_per_block; __shared__ int sp_block_id, ksh_block_id; + #endif + + int thread_id = threadIdx_x; + int worker_id = blockIdx_x; while (1) { __syncthreads(); if (thread_id == 0) { @@ -912,9 +960,17 @@ void cart2sph_kernel(double *out, double *input, PBCIntEnvVars envs, int naux, int nbas, int nao_sph, int pair_compressed) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int pair_ij = item.get_group(1); + int thread_id = item.get_local_id(1); + int aux_id = item.get_group(0) * item.get_local_range(1) + thread_id; + #else int pair_ij = blockIdx.x; int thread_id = threadIdx.x; int aux_id = blockIdx.y * blockDim.x + thread_id; + #endif + if (aux_id >= naux) { return; } @@ -1581,6 +1637,27 @@ int fill_int3c2e(double *out, RysIntEnvVars *envs, double *pool, int ao_pair_offset, int aux_offset, int naux, int reorder_aux, int to_sph) { + #ifdef USE_SYCL + cudaDeviceProp prop; + cudaGetDeviceProperties(&prop, 0); + int workers = prop.multiProcessorCount; + int *head = (int *)(pool + workers * POOL_SIZE); + cudaMemset(head, 0, sizeof(int)); + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(1, workers); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + int3c2e_kernel( + out, dev_envs, pool, omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, ksh_offsets, + gout_stride_lookup, ao_pair_loc, ao_pair_offset, aux_offset, naux, + reorder_aux, to_sph, head, nbatches_shl_pair, nbatches_ksh, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(int3c2e_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); @@ -1597,6 +1674,7 @@ int fill_int3c2e(double *out, RysIntEnvVars *envs, double *pool, fprintf(stderr, "CUDA Error in fill_int3c2e: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -1607,6 +1685,14 @@ int int3c2e_cart2sph(double *out, double *input, PBCIntEnvVars *envs, { constexpr int threads = 256; int aux_batches = (naux + threads - 1) / threads; + #ifdef USE_SYCL + sycl::range<2> thread(1, threads); + sycl::range<2> blocks(aux_batches, nshl_pair); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { + cart2sph_kernel(out, input, dev_envs, bas_ij_idx, out_offsets, input_offsets, naux, nbas, nao_sph, pair_compressed); + }); + #else dim3 blocks(nshl_pair, aux_batches); cart2sph_kernel<<>>( out, input, *envs, bas_ij_idx, out_offsets, input_offsets, @@ -1616,6 +1702,7 @@ int int3c2e_cart2sph(double *out, double *input, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in int3c2e_cart2sph kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/fill_int3c2e_ip1.cu b/gpu4pyscf/lib/gvhf-rys/fill_int3c2e_ip1.cu index eadbab617..2749cfd3d 100644 --- a/gpu4pyscf/lib/gvhf-rys/fill_int3c2e_ip1.cu +++ b/gpu4pyscf/lib/gvhf-rys/fill_int3c2e_ip1.cu @@ -17,8 +17,10 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "gvhf-rys/rys_roots.cu" #include "gvhf-rys/rys_contract_k.cuh" @@ -34,20 +36,63 @@ void int3c2e_ip1_kernel(double *out, RysIntEnvVars envs, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *ksh_offsets, int *gout_stride_lookup, int *ao_pair_loc, int ao_pair_offset, int aux_offset, - int nao_pairs, int naux) + int nao_pairs, int naux + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - int thread_id = threadIdx.x; - int sp_block_id = gridDim.x - blockIdx.x - 1; - int ksh_block_id = gridDim.y - blockIdx.y - 1; - int nbas = envs.nbas; - int *bas = envs.bas; - double *env = envs.env; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nshl_pair = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nksh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1, nshl_pair; __shared__ int ksh0, ksh1, nksh; __shared__ int li, lj, lij, lk, nroots, nf, nao; __shared__ int iprim, jprim, kprim; __shared__ int g_size; __shared__ int gout_stride, nst_per_block; + #endif + + int thread_id = threadIdx_x; + int sp_block_id = gridDim_x - blockIdx_x - 1; + int ksh_block_id = gridDim_y - blockIdx_y - 1; + int nbas = envs.nbas; + int *bas = envs.bas; + double *env = envs.env; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -85,7 +130,6 @@ void int3c2e_ip1_kernel(double *out, RysIntEnvVars envs, int gout_id = thread_id / nst_per_block; int st_id = thread_id - gout_id * nst_per_block; int gx_len = g_size * nst_per_block; - extern __shared__ double shared_memory[]; double *rjri = shared_memory + st_id; double *Rpq = shared_memory + nst_per_block * 3 + st_id; double *gx = shared_memory + nst_per_block * 6 + st_id; @@ -249,20 +293,63 @@ void int3c2e_ipaux_kernel(double *out, RysIntEnvVars envs, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *ksh_offsets, int *gout_stride_lookup, int *ao_pair_loc, int ao_pair_offset, int aux_offset, - int nao_pairs, int naux) + int nao_pairs, int naux + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - int thread_id = threadIdx.x; - int sp_block_id = gridDim.x - blockIdx.x - 1; - int ksh_block_id = gridDim.y - blockIdx.y - 1; - int nbas = envs.nbas; - int *bas = envs.bas; - double *env = envs.env; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int gridDim_x = item.get_group_range(1); + int gridDim_y = item.get_group_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nshl_pair = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nksh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int gridDim_x = gridDim.x; + int gridDim_y = gridDim.y; + + extern __shared__ double shared_memory[]; __shared__ int shl_pair0, shl_pair1, nshl_pair; __shared__ int ksh0, ksh1, nksh; __shared__ int li, lj, lij, lk, nroots, nf, nao; __shared__ int iprim, jprim, kprim; __shared__ int g_size; __shared__ int gout_stride, nst_per_block; + #endif + + int thread_id = threadIdx_x; + int sp_block_id = gridDim_x - blockIdx_x - 1; + int ksh_block_id = gridDim_y - blockIdx_y - 1; + int nbas = envs.nbas; + int *bas = envs.bas; + double *env = envs.env; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -300,7 +387,6 @@ void int3c2e_ipaux_kernel(double *out, RysIntEnvVars envs, int gout_id = thread_id / nst_per_block; int st_id = thread_id - gout_id * nst_per_block; int gx_len = g_size * nst_per_block; - extern __shared__ double shared_memory[]; double *rjri = shared_memory + st_id; double *Rpq = shared_memory + nst_per_block * 3 + st_id; double *gx = shared_memory + nst_per_block * 6 + st_id; @@ -466,6 +552,21 @@ int fill_int3c2e_ip1(double *out, RysIntEnvVars *envs, int *ksh_offsets, int *gout_stride_lookup, int *ao_pair_loc, int ao_pair_offset, int aux_offset, int nao_pairs, int naux) { +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_ksh, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + int3c2e_ip1_kernel( + out, dev_envs, omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, ksh_offsets, + gout_stride_lookup, ao_pair_loc, ao_pair_offset, aux_offset, + nao_pairs, naux, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else cudaFuncSetAttribute(int3c2e_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 blocks(nbatches_shl_pair, nbatches_ksh); int3c2e_ip1_kernel<<>>( @@ -478,6 +579,7 @@ int fill_int3c2e_ip1(double *out, RysIntEnvVars *envs, fprintf(stderr, "CUDA Error in fill_int3c2e: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -488,6 +590,21 @@ int fill_int3c2e_ipaux(double *out, RysIntEnvVars *envs, int *ksh_offsets, int *gout_stride_lookup, int *ao_pair_loc, int ao_pair_offset, int aux_offset, int nao_pairs, int naux) { +#ifdef USE_SYCL + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(nbatches_ksh, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + int3c2e_ipaux_kernel( + out, dev_envs, omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, ksh_offsets, + gout_stride_lookup, ao_pair_loc, ao_pair_offset, aux_offset, + nao_pairs, naux, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else cudaFuncSetAttribute(int3c2e_ipaux_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 blocks(nbatches_shl_pair, nbatches_ksh); int3c2e_ipaux_kernel<<>>( @@ -500,6 +617,7 @@ int fill_int3c2e_ipaux(double *out, RysIntEnvVars *envs, fprintf(stderr, "CUDA Error in fill_int3c2e: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/int3c2e_recontract.cu b/gpu4pyscf/lib/gvhf-rys/int3c2e_recontract.cu index bc7951fc6..035739bf5 100644 --- a/gpu4pyscf/lib/gvhf-rys/int3c2e_recontract.cu +++ b/gpu4pyscf/lib/gvhf-rys/int3c2e_recontract.cu @@ -27,9 +27,16 @@ static __global__ void recontract_kernel(double *out, double *input, int *out_idx, int *inp_idx, double *coef, int naux) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int threads = item.get_local_range(0); + int row_id = item.get_group(0); +#else int thread_id = threadIdx.x; int threads = blockDim.x; int row_id = blockIdx.x; +#endif size_t Naux = naux; out = out + out_idx[row_id] * Naux; input = input + inp_idx[row_id] * Naux; @@ -43,12 +50,21 @@ extern "C" { int recontract_ao_pair(double *out, double *input, int *out_idx, int *inp_idx, double *coef, int naux, int count) { +#ifdef USE_SYCL + sycl::range<1> threads(256); + sycl::range<1> blocks(count); + sycl_get_queue()->parallel_for( + sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + recontract_kernel(out, input, out_idx, inp_idx, coef, naux); + }); +#else recontract_kernel<<>>(out, input, out_idx, inp_idx, coef, naux); cudaError_t err = cudaGetLastError(); if(err != cudaSuccess){ fprintf(stderr, "recontract_ao_pair error %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } diff --git a/gpu4pyscf/lib/gvhf-rys/mole_helper.cu b/gpu4pyscf/lib/gvhf-rys/mole_helper.cu index f92721ad1..40e299488 100644 --- a/gpu4pyscf/lib/gvhf-rys/mole_helper.cu +++ b/gpu4pyscf/lib/gvhf-rys/mole_helper.cu @@ -29,22 +29,88 @@ #define NPRIM_MAX 32 #define PTR_PBAS_IDX 4 +// Macros to abstract CUDA/SYCL thread-indexing and kernel launch differences. +// Each pattern appears 4 times in this file, so macros are warranted. + +#ifdef USE_SYCL +#define SETUP_BRA_KERNEL() \ + auto item = syclex::this_work_item::get_nd_item<3>(); \ + int thread_id = item.get_local_id(2); \ + int col0 = item.get_group(2) * COL_BLKSIZE; \ + int c_bas_id = item.get_group(1); \ + int count = item.get_group(0); \ + int (&p_ao_offsets)[NPRIM_MAX] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else +#define SETUP_BRA_KERNEL() \ + int thread_id = threadIdx.x; \ + int col0 = blockIdx.x * COL_BLKSIZE; \ + int c_bas_id = blockIdx.y; \ + int count = blockIdx.z; \ + __shared__ int p_ao_offsets[NPRIM_MAX]; +#endif + +#ifdef USE_SYCL +#define SETUP_KET_KERNEL() \ + auto item = syclex::this_work_item::get_nd_item<2>(); \ + int tx = item.get_local_id(1); \ + int ty = item.get_local_id(0); \ + int row0 = item.get_group(1) * ROW_BLKSIZE; \ + int c_bas_id = item.get_group(0) * TILE_X + tx; \ + int (&p_ao_offsets)[NPRIM_MAX*TILE_X] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else +#define SETUP_KET_KERNEL() \ + int tx = threadIdx.x; \ + int ty = threadIdx.y; \ + int row0 = blockIdx.x * ROW_BLKSIZE; \ + int c_bas_id = blockIdx.y * TILE_X + tx; \ + __shared__ int p_ao_offsets[NPRIM_MAX*TILE_X]; +#endif + +#ifdef USE_SYCL +#define LAUNCH_BRA_KERNEL(KERNEL, counts_, nbas_, nbatch_col_, ...) { \ + sycl::range<3> _threads(1, 1, THREADS); \ + sycl::range<3> _blocks(counts_, nbas_, nbatch_col_); \ + sycl_get_queue()->parallel_for( \ + sycl::nd_range<3>(_blocks * _threads, _threads), [=](auto item) { \ + KERNEL(__VA_ARGS__); \ + }); \ +} +#else +#define LAUNCH_BRA_KERNEL(KERNEL, counts_, nbas_, nbatch_col_, ...) { \ + dim3 _blocks(nbatch_col_, nbas_, counts_); \ + KERNEL<<<_blocks, THREADS>>>(__VA_ARGS__); \ +} +#endif + +#ifdef USE_SYCL +#define LAUNCH_KET_KERNEL(KERNEL, nbas_, nrow_, ...) { \ + sycl::range<2> _threads(TILE_Y, TILE_X); \ + sycl::range<2> _blocks((nbas_+TILE_X-1)/TILE_X, (nrow_+ROW_BLKSIZE-1)/ROW_BLKSIZE); \ + sycl_get_queue()->parallel_for( \ + sycl::nd_range<2>(_blocks * _threads, _threads), [=](auto item) { \ + KERNEL(__VA_ARGS__); \ + }); \ +} +#else +#define LAUNCH_KET_KERNEL(KERNEL, nbas_, nrow_, ...) { \ + dim3 _threads(TILE_X, TILE_Y); \ + dim3 _blocks((nrow_+ROW_BLKSIZE-1)/ROW_BLKSIZE, (nbas_+TILE_X-1)/TILE_X); \ + KERNEL<<<_blocks, _threads>>>(__VA_ARGS__); \ +} +#endif + static __global__ -void bra_from_sorted_kernel(double *out, double *input, double *recontract_coef, +void bra_sorted2cart_kernel(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, - int *c_ao_loc, int *p_ao_loc, - int nbas, int npbas, int ncol, int cart) + int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int ncol) { - int thread_id = threadIdx.x; - int col0 = blockIdx.x * COL_BLKSIZE; + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + + SETUP_BRA_KERNEL(); int col1 = min(col0 + COL_BLKSIZE, ncol); - int c_bas_id = blockIdx.y; - int count = blockIdx.z; int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; - int nfi = 2 * li + 1; - if (cart) { - nfi = c_nf[li]; - } + int nfi = (li + 1) * (li + 2) / 2; int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; int *pbas_idx = pbas_idx_recontraction + recontract_bas[c_bas_id*BAS_SLOTS+PTR_PBAS_IDX]; @@ -53,9 +119,6 @@ void bra_from_sorted_kernel(double *out, double *input, double *recontract_coef, size_t p_nao = p_ao_loc[npbas]; size_t stride = nfi * ncol; double *pgto = input + count * p_nao * ncol; - constexpr int BLKSIZE = 8; - double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX]; if (thread_id < nprim) { int p_bas_id = pbas_idx[thread_id]; p_ao_offsets[thread_id] = p_ao_loc[p_bas_id]; @@ -90,21 +153,17 @@ void bra_from_sorted_kernel(double *out, double *input, double *recontract_coef, } static __global__ -void bra_to_sorted_kernel(double *out, double *input, double *recontract_coef, - int *recontract_bas, int *pbas_idx_recontraction, - int *c_ao_loc, int *p_ao_loc, - int nbas, int npbas, int ncol, int cart) +void bra_cart2sorted_kernel(double *out, double *input, double *recontract_coef, + int *recontract_bas, int *pbas_idx_recontraction, + int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int ncol) { - int thread_id = threadIdx.x; - int col0 = blockIdx.x * COL_BLKSIZE; + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + + SETUP_BRA_KERNEL(); int col1 = min(col0 + COL_BLKSIZE, ncol); - int c_bas_id = blockIdx.y; - int count = blockIdx.z; int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; - int nfi = 2 * li + 1; - if (cart) { - nfi = c_nf[li]; - } + int nfi = (li + 1) * (li + 2) / 2; int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; int *pbas_idx = pbas_idx_recontraction + recontract_bas[c_bas_id*BAS_SLOTS+PTR_PBAS_IDX]; @@ -113,9 +172,6 @@ void bra_to_sorted_kernel(double *out, double *input, double *recontract_coef, size_t p_nao = p_ao_loc[npbas]; size_t stride = nfi * ncol; double *pgto = out + count * p_nao * ncol; - constexpr int BLKSIZE = 8; - double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX]; if (thread_id < nprim) { int p_bas_id = pbas_idx[thread_id]; p_ao_offsets[thread_id] = p_ao_loc[p_bas_id]; @@ -153,13 +209,13 @@ void bra_sorted2sph_kernel(double *out, double *input, double *recontract_coef, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int ncol) { - int thread_id = threadIdx.x; - int col0 = blockIdx.x * COL_BLKSIZE; + constexpr int BLKSIZE = 4; + double cval[BLKSIZE]; + + SETUP_BRA_KERNEL(); int col1 = min(col0 + COL_BLKSIZE, ncol); - int c_bas_id = blockIdx.y; - int count = blockIdx.z; int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; - int nfi = c_nf[li]; + int nfi = (li + 1) * (li + 2) / 2; int di = li * 2 + 1; int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; @@ -168,9 +224,6 @@ void bra_sorted2sph_kernel(double *out, double *input, double *recontract_coef, size_t c_nao = c_ao_loc[nbas]; size_t p_nao = p_ao_loc[npbas]; double *pgto = input + count * p_nao * ncol; - constexpr int BLKSIZE = 4; - double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX]; if (thread_id < nprim) { int p_bas_id = pbas_idx[thread_id]; p_ao_offsets[thread_id] = p_ao_loc[p_bas_id]; @@ -550,11 +603,11 @@ void bra_sph2sorted_kernel(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int ncol) { - int thread_id = threadIdx.x; - int col0 = blockIdx.x * COL_BLKSIZE; + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + + SETUP_BRA_KERNEL(); int col1 = min(col0 + COL_BLKSIZE, ncol); - int c_bas_id = blockIdx.y; - int count = blockIdx.z; int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; int di = li * 2 + 1; int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; @@ -565,9 +618,6 @@ void bra_sph2sorted_kernel(double *out, double *input, double *recontract_coef, size_t p_nao = p_ao_loc[npbas]; size_t stride = di * ncol; double *pgto = out + count * p_nao * ncol; - constexpr int BLKSIZE = 8; - double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX]; if (thread_id < nprim) { int p_bas_id = pbas_idx[thread_id]; p_ao_offsets[thread_id] = p_ao_loc[p_bas_id]; @@ -870,17 +920,129 @@ void bra_sph2sorted_kernel(double *out, double *input, double *recontract_coef, } } +static __global__ +void ket_sorted2cart_kernel(double *out, double *input, double *recontract_coef, + int *recontract_bas, int *pbas_idx_recontraction, + int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow) +{ + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + + SETUP_KET_KERNEL(); + int thread_id = ty * TILE_X + tx; + int row1 = min(row0 + ROW_BLKSIZE, nrow); + int valid = c_bas_id < nbas; + if (!valid) { + c_bas_id = 0; + } + int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; + int nfi = (li + 1) * (li + 2) / 2; + int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; + int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; + int *pbas_idx = pbas_idx_recontraction + recontract_bas[c_bas_id*BAS_SLOTS+PTR_PBAS_IDX]; + double *coef = recontract_coef + recontract_bas[c_bas_id*BAS_SLOTS+PTR_COEFF]; + size_t c_nao = c_ao_loc[nbas]; + size_t p_nao = p_ao_loc[npbas]; + for (int ip = ty; ip < nprim; ip += TILE_Y) { + int p_bas_id = pbas_idx[ip]; + p_ao_offsets[ip*TILE_X+tx] = p_ao_loc[p_bas_id]; + } + __syncthreads(); + if (!valid) { + return; + } + + for (int ctr0 = 0; ctr0 < n_ctr; ctr0 += BLKSIZE) { + int sub_nctr = min(n_ctr - ctr0, BLKSIZE); + for (int row_id = row0+ty; row_id < row1; row_id += TILE_Y) { + double *cgto = out + row_id*c_nao + c_ao_loc[c_bas_id] + ctr0*nfi; + double *pgto = input + row_id*p_nao; + for (int i = 0; i < nfi; ++i) { + for (int n = 0; n < sub_nctr; ++n) { + if (n == sub_nctr) break; + cval[n] = 0; + } + for (int ip = 0; ip < nprim; ++ip) { + double s = pgto[p_ao_offsets[ip*TILE_X+tx]+i]; + double *c = coef + ctr0*nprim + ip; + for (int n = 0; n < sub_nctr; ++n) { + if (n == sub_nctr) break; + cval[n] += s * c[n*nprim]; + } + } + for (int n = 0; n < sub_nctr; ++n) { + if (n == sub_nctr) break; + cgto[n*nfi+i] = cval[n]; + } + } + } + } +} + +static __global__ +void bra_from_sorted_kernel(double *out, double *input, double *recontract_coef, + int *recontract_bas, int *pbas_idx_recontraction, + int *c_ao_loc, int *p_ao_loc, + int nbas, int npbas, int ncol, int cart) +{ + SETUP_BRA_KERNEL(); + int col1 = min(col0 + COL_BLKSIZE, ncol); + int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; + int nfi = 2 * li + 1; + if (cart) { + nfi = c_nf[li]; + } + int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; + int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; + int *pbas_idx = pbas_idx_recontraction + recontract_bas[c_bas_id*BAS_SLOTS+PTR_PBAS_IDX]; + double *coef = recontract_coef + recontract_bas[c_bas_id*BAS_SLOTS+PTR_COEFF]; + size_t c_nao = c_ao_loc[nbas]; + size_t p_nao = p_ao_loc[npbas]; + size_t stride = nfi * ncol; + double *pgto = input + count * p_nao * ncol; + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + if (thread_id < nprim) { + int p_bas_id = pbas_idx[thread_id]; + p_ao_offsets[thread_id] = p_ao_loc[p_bas_id]; + } + __syncthreads(); + + for (int ctr0 = 0; ctr0 < n_ctr; ctr0 += BLKSIZE) { + int sub_nctr = min(n_ctr - ctr0, BLKSIZE); + size_t c_off = (count * c_nao + c_ao_loc[c_bas_id] + ctr0*nfi) * ncol; + for (int col_id = col0+thread_id; col_id < col1; col_id += THREADS) { + for (int i = 0; i < nfi; ++i) { + for (int n = 0; n < BLKSIZE; ++n) { + if (n == sub_nctr) break; + cval[n] = 0; + } + for (int ip = 0; ip < nprim; ++ip) { + double s = pgto[(size_t)(p_ao_offsets[ip]+i)*ncol+col_id]; + double *c = coef + ctr0*nprim + ip; + for (int n = 0; n < BLKSIZE; ++n) { + if (n == sub_nctr) break; + cval[n] += s * c[n*nprim]; + } + } + double *cgto = out + c_off + (size_t)i * ncol + col_id; + for (int n = 0; n < BLKSIZE; ++n) { + if (n == sub_nctr) break; + cgto[n*stride] = cval[n]; + } + } + } + } +} + static __global__ void ket_from_sorted_kernel(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow, int cart) { - int tx = threadIdx.x; - int ty = threadIdx.y; - int row0 = blockIdx.x * ROW_BLKSIZE; + SETUP_KET_KERNEL(); int row1 = min(row0 + ROW_BLKSIZE, nrow); - int c_bas_id = blockIdx.y * TILE_X + tx; int valid = c_bas_id < nbas; if (!valid) { c_bas_id = 0; @@ -898,7 +1060,6 @@ void ket_from_sorted_kernel(double *out, double *input, double *recontract_coef, size_t p_nao = p_ao_loc[npbas]; constexpr int BLKSIZE = 8; double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX*TILE_X]; for (int ip = ty; ip < nprim; ip += TILE_Y) { int p_bas_id = pbas_idx[ip]; p_ao_offsets[ip*TILE_X+tx] = p_ao_loc[p_bas_id]; @@ -935,17 +1096,68 @@ void ket_from_sorted_kernel(double *out, double *input, double *recontract_coef, } } +static __global__ +void bra_to_sorted_kernel(double *out, double *input, double *recontract_coef, + int *recontract_bas, int *pbas_idx_recontraction, + int *c_ao_loc, int *p_ao_loc, + int nbas, int npbas, int ncol, int cart) +{ + SETUP_BRA_KERNEL(); + int col1 = min(col0 + COL_BLKSIZE, ncol); + int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; + int nfi = 2 * li + 1; + if (cart) { + nfi = c_nf[li]; + } + int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; + int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; + int *pbas_idx = pbas_idx_recontraction + recontract_bas[c_bas_id*BAS_SLOTS+PTR_PBAS_IDX]; + double *coef = recontract_coef + recontract_bas[c_bas_id*BAS_SLOTS+PTR_COEFF]; + size_t c_nao = c_ao_loc[nbas]; + size_t p_nao = p_ao_loc[npbas]; + size_t stride = nfi * ncol; + double *pgto = out + count * p_nao * ncol; + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + if (thread_id < nprim) { + int p_bas_id = pbas_idx[thread_id]; + p_ao_offsets[thread_id] = p_ao_loc[p_bas_id]; + } + __syncthreads(); + + for (int ctr0 = 0; ctr0 < n_ctr; ctr0 += BLKSIZE) { + int sub_nctr = min(n_ctr - ctr0, BLKSIZE); + size_t c_off = (count * c_nao + c_ao_loc[c_bas_id] + ctr0*nfi) * ncol; + for (int i = 0; i < nfi; ++i) { + for (int col_id = col0+thread_id; col_id < col1; col_id += THREADS) { + double *cgto = input + c_off + (size_t)i * ncol + col_id; + for (int n = 0; n < BLKSIZE; ++n) { + if (n == sub_nctr) break; + cval[n] = cgto[n*stride]; + } + for (int ip = 0; ip < nprim; ++ip) { + double *c = coef + ctr0*nprim + ip; + double s = cval[0] * c[0]; + for (int n = 1; n < BLKSIZE; ++n) { + if (n == sub_nctr) break; + s += cval[n] * c[n*nprim]; + } + pgto[(size_t)(p_ao_offsets[ip]+i)*ncol+col_id] += s; + //atomicAdd(pgto+(p_ao_offsets[ip]+i)*ncol+col_id, s); + } + } + } + } +} + static __global__ void ket_to_sorted_kernel(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow, int cart) { - int tx = threadIdx.x; - int ty = threadIdx.y; - int row0 = blockIdx.x * ROW_BLKSIZE; + SETUP_KET_KERNEL(); int row1 = min(row0 + ROW_BLKSIZE, nrow); - int c_bas_id = blockIdx.y * TILE_X + tx; int valid = c_bas_id < nbas; if (!valid) { c_bas_id = 0; @@ -963,7 +1175,62 @@ void ket_to_sorted_kernel(double *out, double *input, double *recontract_coef, size_t p_nao = p_ao_loc[npbas]; constexpr int BLKSIZE = 8; double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX*TILE_X]; + for (int ip = ty; ip < nprim; ip += TILE_Y) { + int p_bas_id = pbas_idx[ip]; + p_ao_offsets[ip*TILE_X+tx] = p_ao_loc[p_bas_id]; + } + __syncthreads(); + if (!valid) { + return; + } + + for (int ctr0 = 0; ctr0 < n_ctr; ctr0 += BLKSIZE) { + int sub_nctr = min(n_ctr - ctr0, BLKSIZE); + for (int row_id = row0+ty; row_id < row1; row_id += TILE_Y) { + double *cgto = input + row_id*c_nao + c_ao_loc[c_bas_id] + ctr0*nfi; + double *pgto = out + row_id*p_nao; + for (int i = 0; i < nfi; ++i) { + for (int n = 0; n < sub_nctr; ++n) { + if (n == sub_nctr) break; + cval[n] = cgto[n*nfi+i]; + } + for (int ip = 0; ip < nprim; ++ip) { + double *c = coef + ctr0*nprim + ip; + double s = cval[0] * c[0]; + for (int n = 1; n < sub_nctr; ++n) { + if (n == sub_nctr) break; + s += cval[n] * c[n*nprim]; + } + pgto[p_ao_offsets[ip*TILE_X+tx]+i] += s; + } + } + } + } +} + +static __global__ +void ket_cart2sorted_kernel(double *out, double *input, double *recontract_coef, + int *recontract_bas, int *pbas_idx_recontraction, + int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow) +{ + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + + SETUP_KET_KERNEL(); + int thread_id = ty * TILE_X + tx; + int row1 = min(row0 + ROW_BLKSIZE, nrow); + int valid = c_bas_id < nbas; + if (!valid) { + c_bas_id = 0; + } + int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; + int nfi = (li + 1) * (li + 2) / 2; + int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; + int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; + int *pbas_idx = pbas_idx_recontraction + recontract_bas[c_bas_id*BAS_SLOTS+PTR_PBAS_IDX]; + double *coef = recontract_coef + recontract_bas[c_bas_id*BAS_SLOTS+PTR_COEFF]; + size_t c_nao = c_ao_loc[nbas]; + size_t p_nao = p_ao_loc[npbas]; for (int ip = ty; ip < nprim; ip += TILE_Y) { int p_bas_id = pbas_idx[ip]; p_ao_offsets[ip*TILE_X+tx] = p_ao_loc[p_bas_id]; @@ -1002,17 +1269,18 @@ void ket_sorted2sph_kernel(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow) { - int tx = threadIdx.x; - int ty = threadIdx.y; - int row0 = blockIdx.x * ROW_BLKSIZE; + constexpr int BLKSIZE = 4; + double cval[BLKSIZE]; + + SETUP_KET_KERNEL(); + int thread_id = ty * TILE_X + tx; int row1 = min(row0 + ROW_BLKSIZE, nrow); - int c_bas_id = blockIdx.y * TILE_X + tx; int valid = c_bas_id < nbas; if (!valid) { c_bas_id = 0; } int li = recontract_bas[c_bas_id*BAS_SLOTS+ANG_OF]; - int nfi = c_nf[li]; + int nfi = (li + 1) * (li + 2) / 2; int di = li * 2 + 1; int nprim = recontract_bas[c_bas_id*BAS_SLOTS+NPRIM_OF]; int n_ctr = recontract_bas[c_bas_id*BAS_SLOTS+NCTR_OF ]; @@ -1020,9 +1288,6 @@ void ket_sorted2sph_kernel(double *out, double *input, double *recontract_coef, double *coef = recontract_coef + recontract_bas[c_bas_id*BAS_SLOTS+PTR_COEFF]; size_t c_nao = c_ao_loc[nbas]; size_t p_nao = p_ao_loc[npbas]; - constexpr int BLKSIZE = 4; - double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX*TILE_X]; for (int ip = ty; ip < nprim; ip += TILE_Y) { int p_bas_id = pbas_idx[ip]; p_ao_offsets[ip*TILE_X+tx] = p_ao_loc[p_bas_id]; @@ -1406,11 +1671,12 @@ void ket_sph2sorted_kernel(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow) { - int tx = threadIdx.x; - int ty = threadIdx.y; - int row0 = blockIdx.x * ROW_BLKSIZE; + constexpr int BLKSIZE = 8; + double cval[BLKSIZE]; + + SETUP_KET_KERNEL(); + int thread_id = ty * TILE_X + tx; int row1 = min(row0 + ROW_BLKSIZE, nrow); - int c_bas_id = blockIdx.y * TILE_X + tx; int valid = c_bas_id < nbas; if (!valid) { c_bas_id = 0; @@ -1423,9 +1689,6 @@ void ket_sph2sorted_kernel(double *out, double *input, double *recontract_coef, double *coef = recontract_coef + recontract_bas[c_bas_id*BAS_SLOTS+PTR_COEFF]; size_t c_nao = c_ao_loc[nbas]; size_t p_nao = p_ao_loc[npbas]; - constexpr int BLKSIZE = 8; - double cval[BLKSIZE]; - __shared__ int p_ao_offsets[NPRIM_MAX*TILE_X]; for (int ip = ty; ip < nprim; ip += TILE_Y) { int p_bas_id = pbas_idx[ip]; p_ao_offsets[ip*TILE_X+tx] = p_ao_loc[p_bas_id]; @@ -1738,8 +2001,7 @@ int bra_from_sorted(double *out, double *input, double *recontract_coef, int nbas, int npbas, int ncol, int counts, int cart) { int nbatch_col = (ncol + COL_BLKSIZE-1) / COL_BLKSIZE; - dim3 blocks(nbatch_col, nbas, counts); - bra_from_sorted_kernel<<>>( + LAUNCH_BRA_KERNEL(bra_from_sorted_kernel, counts, nbas, nbatch_col, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, ncol, cart); cudaError_t err = cudaGetLastError(); @@ -1765,13 +2027,12 @@ int bra_to_sorted(double *out, double *input, double *recontract_coef, int nbas, int npbas, int ncol, int counts, int cart) { int nbatch_col = (ncol + COL_BLKSIZE-1) / COL_BLKSIZE; - dim3 blocks(nbatch_col, nbas, counts); - bra_to_sorted_kernel<<>>( + LAUNCH_BRA_KERNEL(bra_to_sorted_kernel, counts, nbas, nbatch_col, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, ncol, cart); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in bra_cart2sorted kernel: %s\n", cudaGetErrorString(err)); + fprintf(stderr, "CUDA Error in bra_to_sorted kernel: %s\n", cudaGetErrorString(err)); return 1; } return 0; @@ -1791,8 +2052,7 @@ int bra_sorted2sph(double *out, double *input, double *recontract_coef, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int ncol, int counts) { int nbatch_col = (ncol + COL_BLKSIZE-1) / COL_BLKSIZE; - dim3 blocks(nbatch_col, nbas, counts); - bra_sorted2sph_kernel<<>>( + LAUNCH_BRA_KERNEL(bra_sorted2sph_kernel, counts, nbas, nbatch_col, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, ncol); cudaError_t err = cudaGetLastError(); @@ -1808,8 +2068,7 @@ int bra_sph2sorted(double *out, double *input, double *recontract_coef, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int ncol, int counts) { int nbatch_col = (ncol + COL_BLKSIZE-1) / COL_BLKSIZE; - dim3 blocks(nbatch_col, nbas, counts); - bra_sph2sorted_kernel<<>>( + LAUNCH_BRA_KERNEL(bra_sph2sorted_kernel, counts, nbas, nbatch_col, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, ncol); cudaError_t err = cudaGetLastError(); @@ -1824,14 +2083,12 @@ int ket_from_sorted(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow, int cart) { - dim3 threads(TILE_X, TILE_Y); - dim3 blocks((nrow+ROW_BLKSIZE-1)/ROW_BLKSIZE, (nbas+TILE_X-1)/TILE_X); - ket_from_sorted_kernel<<>>( + LAUNCH_KET_KERNEL(ket_from_sorted_kernel, nbas, nrow, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, nrow, cart); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in ket_sorted2cart kernel: %s\n", cudaGetErrorString(err)); + fprintf(stderr, "CUDA Error in ket_from_sorted kernel: %s\n", cudaGetErrorString(err)); return 1; } return 0; @@ -1850,14 +2107,12 @@ int ket_to_sorted(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow, int cart) { - dim3 threads(TILE_X, TILE_Y); - dim3 blocks((nrow+ROW_BLKSIZE-1)/ROW_BLKSIZE, (nbas+TILE_X-1)/TILE_X); - ket_to_sorted_kernel<<>>( + LAUNCH_KET_KERNEL(ket_to_sorted_kernel, nbas, nrow, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, nrow, cart); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in ket_cart2sorted kernel: %s\n", cudaGetErrorString(err)); + fprintf(stderr, "CUDA Error in ket_to_sorted kernel: %s\n", cudaGetErrorString(err)); return 1; } return 0; @@ -1876,9 +2131,7 @@ int ket_sorted2sph(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow) { - dim3 threads(TILE_X, TILE_Y); - dim3 blocks((nrow+ROW_BLKSIZE-1)/ROW_BLKSIZE, (nbas+TILE_X-1)/TILE_X); - ket_sorted2sph_kernel<<>>( + LAUNCH_KET_KERNEL(ket_sorted2sph_kernel, nbas, nrow, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, nrow); cudaError_t err = cudaGetLastError(); @@ -1893,9 +2146,7 @@ int ket_sph2sorted(double *out, double *input, double *recontract_coef, int *recontract_bas, int *pbas_idx_recontraction, int *c_ao_loc, int *p_ao_loc, int nbas, int npbas, int nrow) { - dim3 threads(TILE_X, TILE_Y); - dim3 blocks((nrow+ROW_BLKSIZE-1)/ROW_BLKSIZE, (nbas+TILE_X-1)/TILE_X); - ket_sph2sorted_kernel<<>>( + LAUNCH_KET_KERNEL(ket_sph2sorted_kernel, nbas, nrow, out, input, recontract_coef, recontract_bas, pbas_idx_recontraction, c_ao_loc, p_ao_loc, nbas, npbas, nrow); cudaError_t err = cudaGetLastError(); diff --git a/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.c b/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.c index ee7c78b6d..c68701150 100644 --- a/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.c +++ b/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.c @@ -7,11 +7,14 @@ // sqrt(-log(1e-9)) #define R_GUESS_FAC 4.5f +#ifdef __cplusplus +extern "C" { +#endif void sr_eri_s_estimator(float *s_estimator, float omega, float *diffuse_exps, float *diffuse_ctr_coef, int *atm, int natm, int *bas, int nbas, double *env) { - float *rx = malloc(sizeof(float) * nbas * 3); + float *rx = (float *)malloc(sizeof(float) * nbas * 3); float *ry = rx + nbas; float *rz = ry + nbas; @@ -73,3 +76,7 @@ void sr_eri_s_estimator(float *s_estimator, float omega, } free(rx); } + +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.cu b/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.cu index 73fc811bc..587049d31 100644 --- a/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.cu +++ b/gpu4pyscf/lib/gvhf-rys/nr_sr_estimator.cu @@ -42,13 +42,19 @@ void fill_s_estimator_kernel(float *s_estimator, RysIntEnvVars envs, uint32_t *bas_ij_idx, float *diffuse_exps, float *diffuse_ctr_coef, int npairs, double omega) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + uint32_t sp_block_id = item.get_group(0); + int t_id = item.get_local_id(0); +#else uint32_t sp_block_id = blockIdx.x; int t_id = threadIdx.x; +#endif int *bas = envs.bas; double *env = envs.env; uint32_t nbas = envs.nbas; uint32_t shl_pair0 = sp_block_id * SP_BLOCK_SIZE; - uint32_t shl_pair1 = min((sp_block_id+1) * SP_BLOCK_SIZE, npairs); + uint32_t shl_pair1 = min((sp_block_id+1) * SP_BLOCK_SIZE, (uint32_t)npairs); float omega2 = omega * omega; for (uint32_t pair_ij = shl_pair0+t_id; pair_ij < shl_pair1; pair_ij += THREADS) { @@ -99,11 +105,43 @@ void fill_s_estimator_kernel(float *s_estimator, RysIntEnvVars envs, static __global__ void int2e_qcond_kernel(float *q_out, RysIntEnvVars envs, uint32_t *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - double omega, double lr_factor, double sr_factor) + double omega, double lr_factor, double sr_factor + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte* shm_mem + #endif + ) { + #ifdef USE_SYCL + int sp_block_id = item.get_group(1); + int thread_id = item.get_local_id(1); + int threads = item.get_local_range(1); + + float* shared_memory = reinterpret_cast(shm_mem); + + auto thread_block = item.get_group(); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &stride_k = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else int sp_block_id = blockIdx.x; int thread_id = threadIdx.x; int threads = blockDim.x; + + extern __shared__ float shared_memory[]; + + __shared__ int li, lj; + __shared__ int iprim, jprim; + __shared__ int nroots; + __shared__ int stride_k, g_size; + __shared__ int gout_stride, nsp_per_block; + #endif + int shl_pair0 = shl_pair_offsets[sp_block_id]; int shl_pair1 = shl_pair_offsets[sp_block_id+1]; int bas_ij0 = bas_ij_idx[shl_pair0]; @@ -113,7 +151,7 @@ void int2e_qcond_kernel(float *q_out, RysIntEnvVars envs, uint32_t *bas_ij_idx, int *bas = envs.bas; double *env = envs.env; - __shared__ int li, lj; + if (thread_id == 0) { li = bas[ish0*BAS_SLOTS+ANG_OF]; lj = bas[jsh0*BAS_SLOTS+ANG_OF]; @@ -123,10 +161,6 @@ void int2e_qcond_kernel(float *q_out, RysIntEnvVars envs, uint32_t *bas_ij_idx, return; } - __shared__ int iprim, jprim; - __shared__ int nroots; - __shared__ int stride_k, g_size; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; @@ -152,7 +186,7 @@ void int2e_qcond_kernel(float *q_out, RysIntEnvVars envs, uint32_t *bas_ij_idx, int lprim = jprim; int stride_j = li + 1; int nfij = nfi * nfj; - extern __shared__ float shared_memory[]; + float *rjri = shared_memory + sp_id; float *Rpq = shared_memory + nsp_per_block * 3 + sp_id; float *rw = shared_memory + nsp_per_block * 6 + sp_id; @@ -161,8 +195,8 @@ void int2e_qcond_kernel(float *q_out, RysIntEnvVars envs, uint32_t *bas_ij_idx, float *gx = shared_memory + nsp_per_block * (nroots * 2 + 6) + sp_id; // gz can be reused for gbuf float *gbuf = gx + g_size * nsp_per_block * 2; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); for (int task_id = shl_pair0+sp_id; task_id < shl_pair1+sp_id; task_id += nsp_per_block) { float gout[GOUT_WIDTH]; @@ -402,6 +436,12 @@ int fill_s_estimator(float *s_estimator, RysIntEnvVars *envs, float *diffuse_ctr_coef, int npairs, double omega) { int sp_blocks = (npairs + SP_BLOCK_SIZE - 1) / SP_BLOCK_SIZE; + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<1>(sp_blocks * THREADS, THREADS), [=](auto item) { + fill_s_estimator_kernel(s_estimator, dev_envs, bas_ij_idx, diffuse_exps, diffuse_ctr_coef, npairs, omega); + }); + #else fill_s_estimator_kernel<<>>( s_estimator, *envs, bas_ij_idx, diffuse_exps, diffuse_ctr_coef, npairs, omega); @@ -410,6 +450,7 @@ int fill_s_estimator(float *s_estimator, RysIntEnvVars *envs, fprintf(stderr, "CUDA Error in fill_s_estimator_kernel %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -418,6 +459,21 @@ int int2e_qcond_estimator(float *q_out, RysIntEnvVars *envs, int shm_size, int *shl_pair_offsets, int *gout_stride_lookup, double omega, double lr_factor, double sr_factor) { + #ifdef USE_SYCL + // Note: Though the kernel is 1D launch in CUDA, SYCL had to do 2D because of the free-functions used in + // rys_roots_for_k() method + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(1, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + int2e_qcond_kernel(q_out, dev_envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + omega, lr_factor, sr_factor, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else int2e_qcond_kernel<<>>( q_out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, omega, lr_factor, sr_factor); @@ -426,6 +482,7 @@ int int2e_qcond_estimator(float *q_out, RysIntEnvVars *envs, int shm_size, fprintf(stderr, "CUDA Error in int2e_qcond_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/rys_contract_j.cu b/gpu4pyscf/lib/gvhf-rys/rys_contract_j.cu index 78277d7ff..df85124b5 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_contract_j.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_contract_j.cu @@ -24,28 +24,68 @@ #include "rys_roots.cu" #include "create_tasks.cu" +#ifdef USE_SYCL +SYCL_EXTERNAL sycl_device_global s_rys_i_in_fold2idx; +SYCL_EXTERNAL sycl_device_global s_rys_i_in_fold3idx; +#else __constant__ Fold2Index c_i_in_fold2idx[165]; __constant__ Fold3Index c_i_in_fold3idx[495]; +#endif __global__ static void rys_j_kernel(RysIntEnvVars envs, JKMatrix jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, int *head, int reserved_shm_size) + uint32_t *pool, int *head, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double *shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int blockIdx_x = item.get_group(1); + + auto c_i_in_fold2idx = s_rys_i_in_fold2idx.get(); + auto c_i_in_fold3idx = s_rys_i_in_fold3idx.get(); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; - __shared__ int ish, jsh; + __shared__ int ish; + __shared__ int jsh; __shared__ double ri[3]; __shared__ double rjri[3]; __shared__ double aij_cache[2]; __shared__ int expi; __shared__ int expj; + #endif + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -499,21 +539,55 @@ __global__ static void rys_j_with_gout_kernel(RysIntEnvVars envs, JKMatrix jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, int *head, int reserved_shm_size) + uint32_t *pool, int *head, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double *shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int blockIdx_x = item.get_group(1); + + auto c_i_in_fold3idx = s_rys_i_in_fold3idx.get(); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; - __shared__ int ish, jsh; + __shared__ int ish; + __shared__ int jsh; __shared__ double ri[3]; __shared__ double rjri[3]; __shared__ double aij_cache[2]; __shared__ int expi; __shared__ int expj; + #endif + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -877,12 +951,25 @@ int RYS_build_j(double *vj, double *dm, int n_dm, int nao, int nmax = MAX(lij, lkl); int nf3_ij = (lij+1)*(lij+2)*(lij+3)/6; int nf3_kl = (lkl+1)*(lkl+2)*(lkl+3)/6; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block; if (with_gout) { buflen += nf3_ij*nf3_kl * quartets_per_block; int reserved_shm_size = buflen; buflen += iprim * jprim; + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_j_with_gout_kernel(dev_envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, head, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); buflen *= sizeof(double); if (buflen > 48000) { cudaFuncSetAttribute(rys_j_with_gout_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); @@ -896,11 +983,26 @@ int RYS_build_j(double *vj, double *dm, int n_dm, int nao, rys_j_with_gout_kernel<<>>( *envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head, reserved_shm_size); + #endif } else { buflen += (nf3_ij+nf3_kl*2+(lij+1)*(lkl+1)*(nmax+2)) * quartets_per_block; int reserved_shm_size = buflen; buflen += iprim * jprim; buflen += nf3_ij; // dm_ij_cache + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_j_kernel(dev_envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, head, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); buflen *= sizeof(double); if (buflen > 48000) { cudaFuncSetAttribute(rys_j_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); @@ -914,6 +1016,7 @@ int RYS_build_j(double *vj, double *dm, int n_dm, int nao, rys_j_kernel<<>>( *envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head, reserved_shm_size); + #endif } } cudaError_t err = cudaGetLastError(); @@ -949,8 +1052,13 @@ int RYS_init_rysj_constant() } } } } + #ifdef USE_SYCL + sycl_get_queue()->memcpy(s_rys_i_in_fold2idx, i_in_fold2idx, 165*sizeof(Fold2Index)).wait(); + sycl_get_queue()->memcpy(s_rys_i_in_fold3idx, i_in_fold3idx, 495*sizeof(Fold3Index)).wait(); + #else cudaMemcpyToSymbol(c_i_in_fold2idx, i_in_fold2idx, 165*sizeof(Fold2Index)); cudaMemcpyToSymbol(c_i_in_fold3idx, i_in_fold3idx, 495*sizeof(Fold3Index)); + #endif return 0; } } diff --git a/gpu4pyscf/lib/gvhf-rys/rys_contract_jk.cu b/gpu4pyscf/lib/gvhf-rys/rys_contract_jk.cu index d71d2750c..7fcc80a48 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_contract_jk.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_contract_jk.cu @@ -33,15 +33,46 @@ __global__ static void rys_jk_kernel(RysIntEnvVars envs, JKMatrix jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, int *head, int gout_pattern, int reserved_shm_size) + uint32_t *pool, int *head, int gout_pattern, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { - // sq is short for shl_quartet - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int blockIdx_x = item.get_group(1); + + double *shared_memory = reinterpret_cast(shm_mem); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + auto c_gxyz_offset = s_rys_gxyz_offset.get(); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int ish, jsh; __shared__ int i0, j0, nao; @@ -50,6 +81,13 @@ void rys_jk_kernel(RysIntEnvVars envs, JKMatrix jk, BoundsInfo bounds, __shared__ double aij_cache[2]; __shared__ int expi; __shared__ int expj; + #endif + // sq is short for shl_quartet + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -404,13 +442,33 @@ int RYS_build_jk(double *vj, double *vk, double *dm, int n_dm, int nao, int n_tiles = ntiles_i * ntiles_j * ntiles_k * ntiles_l; auto launch = [&](auto offset, int tile_chunk) { + #ifdef USE_SYCL + sycl_get_queue()->memcpy(s_rys_gxyz_offset, gxyz_offset+offset, + tile_chunk*sizeof(GXYZOffset)).wait(); + #else checkCudaErrors( cudaMemcpyToSymbol(c_gxyz_offset, gxyz_offset+offset, tile_chunk*sizeof(GXYZOffset), 0, cudaMemcpyHostToDevice)); + #endif int scheme[4]; threads_scheme_for_k(scheme, bounds, shm_size, tile_chunk); int buflen = scheme[2]; + int reserved_shm_size = scheme[3]; + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(scheme[1], scheme[0]); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_jk_kernel(dev_envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, + head + offset/256, gout_pattern, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else if (buflen > 48000) { cudaFuncSetAttribute(rys_jk_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); cudaError_t err = cudaGetLastError(); @@ -421,11 +479,11 @@ int RYS_build_jk(double *vj, double *vk, double *dm, int n_dm, int nao, } } dim3 threads(scheme[0], scheme[1]); - int reserved_shm_size = scheme[3]; rys_jk_kernel<<>>( *envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head + offset/256, gout_pattern, reserved_shm_size); + #endif }; launch(0, 256); diff --git a/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip1.cu b/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip1.cu index 7e5c5ece0..ce335edfc 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip1.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip1.cu @@ -33,21 +33,54 @@ __global__ static void rys_vjk_ip1_kernel(RysIntEnvVars envs, JKMatrix jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, int *head, int nf, int reserved_shm_size) + uint32_t *pool, int *head, int nf, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double *shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; - __shared__ int ish, jsh; + __shared__ int ish; + __shared__ int jsh; __shared__ double ri[3]; __shared__ double rjri[3]; __shared__ double aij_cache[3]; __shared__ int expi; __shared__ int expj; + #endif + + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -355,15 +388,41 @@ void rys_ejk_ip1_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, uint32_t *pool, int *head, double *dd_pool, int nf, - int reserved_shm_size) + int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double* shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - double *dd_cache = dd_pool + blockIdx.x * nf * blockDim.x + sq_id; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int ish, jsh; __shared__ int i0, j0; @@ -371,6 +430,14 @@ void rys_ejk_ip1_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, __shared__ double rjri[3]; __shared__ double aij_cache[2]; __shared__ int expi, expj; + #endif + + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; + double *dd_cache = dd_pool + blockIdx_x * nf * blockDim_x + sq_id; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -404,10 +471,10 @@ void rys_ejk_ip1_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, double *gx = shared_memory + nsq_per_block * 6 + sq_id; double *rw = shared_memory + nsq_per_block * (g_size*3+6) + sq_id; double *cicj_cache = shared_memory + reserved_shm_size; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); - int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(lk); - int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(ll); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(lk); + const int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(ll); while (1) { __syncthreads(); if (t_id == 0) { @@ -763,15 +830,41 @@ void rys_ejk_ip1_multidm_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo boun float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, uint32_t *pool, int *head, double *dd_pool, - int dd_cache_size, int nf, int reserved_shm_size) + int dd_cache_size, int nf, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double* shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - double *dd_cache = dd_pool + blockIdx.x * dd_cache_size + sq_id; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int ish, jsh; __shared__ int i0, j0; @@ -779,6 +872,14 @@ void rys_ejk_ip1_multidm_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo boun __shared__ double rjri[3]; __shared__ double aij_cache[2]; __shared__ int expi, expj; + #endif + + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; + double *dd_cache = dd_pool + blockIdx_x * dd_cache_size + sq_id; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -1226,19 +1327,45 @@ while (1) { __global__ static void rys_ejk_ip1_sum_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, - double *j_factor, double *k_factor, double *dm1, double *dm2, - float *q_cond_ij, float *q_cond_kl, float dm_penalty, - float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, int *head, double *dd_pool, - int dd_cache_size, int nf, int reserved_shm_size) + double *j_factor, double *k_factor, double *dm1, double *dm2, + float *q_cond_ij, float *q_cond_kl, float dm_penalty, + float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, + uint32_t *pool, int *head, double *dd_pool, + int dd_cache_size, int nf, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double* shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - double *dd_cache = dd_pool + blockIdx.x * dd_cache_size + sq_id; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int ish, jsh; __shared__ int i0, j0; @@ -1246,6 +1373,14 @@ void rys_ejk_ip1_sum_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, __shared__ double rjri[3]; __shared__ double aij_cache[2]; __shared__ int expi, expj; + #endif + + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; + double *dd_cache = dd_pool + blockIdx_x * dd_cache_size + sq_id; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -1708,23 +1843,38 @@ int RYS_build_jk_ip1(double *vj, double *vk, double *dm, int n_dm, int nao, int int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int reserved_shm_size = (nroots*2 + g_size*3 + 6) * quartets_per_block; int buflen = reserved_shm_size + ij_prims; - buflen *= sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_vjk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_vjk_ip1_kernel(envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, head, nf, + reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + int shm_bytes = buflen * sizeof(double); + if (shm_bytes > 48000) { + cudaFuncSetAttribute(rys_vjk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_bytes); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", shm_bytes, cudaGetErrorString(err)); return 1; } } - rys_vjk_ip1_kernel<<>>( + rys_vjk_ip1_kernel<<>>( envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head, nf, reserved_shm_size); + #endif } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -1803,23 +1953,39 @@ int RYS_per_atom_jk_ip1(double *ejk, double j_factor, double k_factor, int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block; int reserved_shm_size = max(buflen, 6*gout_stride*quartets_per_block); - buflen = (reserved_shm_size + ij_prims)*sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_ejk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + buflen = reserved_shm_size + ij_prims; + + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_ejk_ip1_kernel(envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, head, dd_pool, nf, + reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + int shm_bytes = buflen * sizeof(double); + if (shm_bytes > 48000) { + cudaFuncSetAttribute(rys_ejk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_bytes); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", shm_bytes, cudaGetErrorString(err)); return 1; } } - rys_ejk_ip1_kernel<<>>( + rys_ejk_ip1_kernel<<>>( envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head, dd_pool, nf, reserved_shm_size); + #endif } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -1889,19 +2055,21 @@ int RYS_per_atom_jk_ip1_multidm(double *ejk, double *j_factor, double *j_factor_ int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block; int reserved_shm_size = max(buflen, 6*gout_stride*quartets_per_block); - buflen = (reserved_shm_size + ij_prims)*sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_ejk_ip1_multidm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + buflen = reserved_shm_size + ij_prims; + #ifndef USE_SYCL + if (buflen*(int)sizeof(double) > 48000) { + cudaFuncSetAttribute(rys_ejk_ip1_multidm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + buflen*sizeof(double)); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, - cudaGetErrorString(err)); + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", + (int)(buflen*sizeof(double)), cudaGetErrorString(err)); return 1; } } + #endif size_t nao2 = nao * nao; int nf = nfi * nfj * nfk * nfl; for (int n = 0; n < n_dm; n += DM_BLOCK) { @@ -1917,10 +2085,26 @@ int RYS_per_atom_jk_ip1_multidm(double *ejk, double *j_factor, double *j_factor_ jk.k_factor = 1.; // ensure not dropping by schwarz screening } } - rys_ejk_ip1_multidm_kernel<<>>( + + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_ejk_ip1_multidm_kernel(envs, jk, bounds, j_factor+n, k_factor+n, dm1+n*nao2, dm2+n*nao2, + q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, + pool, head, dd_pool, dd_cache_size, nf, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + rys_ejk_ip1_multidm_kernel<<>>( envs, jk, bounds, j_factor+n, k_factor+n, dm1+n*nao2, dm2+n*nao2, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head, dd_pool, dd_cache_size, nf, reserved_shm_size); + #endif head++; } cudaError_t err = cudaGetLastError(); @@ -1989,19 +2173,21 @@ int RYS_per_atom_jk_ip1_sum(double *ejk, double *j_factor, double *j_factor_cpu, int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block; int reserved_shm_size = max(buflen, 6*gout_stride*quartets_per_block); - buflen = (reserved_shm_size + ij_prims)*sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_ejk_ip1_sum_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + buflen = reserved_shm_size + ij_prims; + #ifndef USE_SYCL + if (buflen*(int)sizeof(double) > 48000) { + cudaFuncSetAttribute(rys_ejk_ip1_sum_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + buflen*sizeof(double)); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, - cudaGetErrorString(err)); + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", + (int)(buflen*sizeof(double)), cudaGetErrorString(err)); return 1; } } + #endif for (int i = 0; i < n_dm; ++i) { if (j_factor_cpu[i] != 0.) { jk.j_factor = 1.; // ensure not dropping by schwarz screening @@ -2010,10 +2196,26 @@ int RYS_per_atom_jk_ip1_sum(double *ejk, double *j_factor, double *j_factor_cpu, jk.k_factor = 1.; // ensure not dropping by schwarz screening } } - rys_ejk_ip1_sum_kernel<<>>( + + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_ejk_ip1_sum_kernel(envs, jk, bounds, j_factor, k_factor, dm1, dm2, + q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, + pool, head, dd_pool, dd_cache_size, nf, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + rys_ejk_ip1_sum_kernel<<>>( envs, jk, bounds, j_factor, k_factor, dm1, dm2, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head, dd_pool, dd_cache_size, nf, reserved_shm_size); + #endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in RYS_per_atom_jk_ip1_sum, li,lj,lk,ll = %d,%d,%d,%d, error message = %s\n", diff --git a/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip2.cu b/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip2.cu index 5b5d31101..e6e510536 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip2.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_contract_jk_ip2.cu @@ -34,20 +34,52 @@ __global__ static void rys_ejk_ip2_type12_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, double *dd_pool, int *head, int nf) + uint32_t *pool, double *dd_pool, int *head, int nf + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double *shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - double *dd_cache = dd_pool + blockIdx.x * nf * blockDim.x + sq_id; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int blockIdx_x = item.get_group(1); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[4] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ri = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int ish, jsh; __shared__ double rjri[3]; __shared__ double aij_cache[4]; __shared__ int ri, expi, expj; + #endif + + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; + double *dd_cache = dd_pool + blockIdx_x * nf * blockDim_x + sq_id; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -79,10 +111,10 @@ void rys_ejk_ip2_type12_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bound double *gx = shared_memory + nsq_per_block * 6 + sq_id; double *rw = shared_memory + nsq_per_block * (g_size*3+6) + sq_id; double *cicj_cache = shared_memory + nsq_per_block * (g_size*3+nroots*2+6); - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); - int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(lk); - int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(ll); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(lk); + const int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(ll); while (1) { __syncthreads(); if (t_id == 0) { @@ -116,6 +148,7 @@ while (1) { expj = bas[jsh*BAS_SLOTS+PTR_EXP]; ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; } + __syncthreads(); if (t_id < 3) { int ri_ptr = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; int rj_ptr = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; @@ -625,20 +658,52 @@ __global__ static void rys_ejk_ip2_type3_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, double *dd_pool, int *head, int nf) + uint32_t *pool, double *dd_pool, int *head, int nf + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double *shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - double *dd_cache = dd_pool + blockIdx.x * nf * blockDim.x + sq_id; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int blockIdx_x = item.get_group(1); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[4] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ri = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int ish, jsh; __shared__ double rjri[3]; __shared__ double aij_cache[4]; __shared__ int ri, expi, expj; + #endif + + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; + double *dd_cache = dd_pool + blockIdx_x * nf * blockDim_x + sq_id; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -670,10 +735,10 @@ void rys_ejk_ip2_type3_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds double *gx = shared_memory + nsq_per_block * 6 + sq_id; double *rw = shared_memory + nsq_per_block * (g_size*3+6) + sq_id; double *cicj_cache = shared_memory + nsq_per_block * (g_size*3+nroots*2+6); - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); - int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(lk); - int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(ll); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(lk); + const int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(ll); while (1) { __syncthreads(); if (t_id == 0) { @@ -707,6 +772,7 @@ while (1) { expj = bas[jsh*BAS_SLOTS+PTR_EXP]; ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; } + __syncthreads(); if (t_id < 3) { int ri_ptr = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; int rj_ptr = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; @@ -1263,22 +1329,36 @@ int RYS_per_atom_jk_ip2_type12(double *ejk, double j_factor, double k_factor, int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block + ij_prims; - buflen *= sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_ejk_ip2_type12_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_ejk_ip2_type12_kernel(envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, dd_pool, head, nf, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + int shm_bytes = buflen * sizeof(double); + if (shm_bytes > 48000) { + cudaFuncSetAttribute(rys_ejk_ip2_type12_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_bytes); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", shm_bytes, cudaGetErrorString(err)); return 1; } } - rys_ejk_ip2_type12_kernel<<>>( + rys_ejk_ip2_type12_kernel<<>>( envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, dd_pool, head, nf); + #endif } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -1354,23 +1434,37 @@ int RYS_per_atom_jk_ip2_type3(double *ejk, double j_factor, double k_factor, int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block + ij_prims; buflen = max(buflen, 9*gout_stride*quartets_per_block); - buflen *= sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_ejk_ip2_type3_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_ejk_ip2_type3_kernel(envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, dd_pool, head, nf, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + int shm_bytes = buflen * sizeof(double); + if (shm_bytes > 48000) { + cudaFuncSetAttribute(rys_ejk_ip2_type3_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_bytes); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", shm_bytes, cudaGetErrorString(err)); return 1; } } - rys_ejk_ip2_type3_kernel<<>>( + rys_ejk_ip2_type3_kernel<<>>( envs, jk, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, dd_pool, head, nf); + #endif } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { diff --git a/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cu b/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cu index 2ef4883e8..be6464eec 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cu @@ -27,6 +27,10 @@ #include "create_tasks.cu" #include "rys_contract_k.cuh" +#ifdef USE_SYCL +SYCL_EXTERNAL sycl_device_global s_rys_gxyz_offset; +#endif + #define GOUT_WIDTH1 81 // gout_pattern = ((li == 0) << 3) | ((lj == 0) << 2) | ((lk == 0) << 1) | (ll == 0); @@ -34,14 +38,44 @@ __global__ static void rys_k_kernel(RysIntEnvVars envs, JKMatrix kmat, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float dm_penalty, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - uint32_t *pool, int *head, int gout_pattern, int reserved_shm_size) + uint32_t *pool, int *head, int gout_pattern, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { - // sq is short for shl_quartet - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - uint32_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; + #ifdef USE_SYCL + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int blockIdx_x = item.get_group(1); + + double *shared_memory = reinterpret_cast(shm_mem); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + auto c_gxyz_offset = s_rys_gxyz_offset.get(); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int ish, jsh; @@ -51,6 +85,13 @@ void rys_k_kernel(RysIntEnvVars envs, JKMatrix kmat, BoundsInfo bounds, __shared__ double aij_cache[2]; __shared__ int expi; __shared__ int expj; + #endif + // sq is short for shl_quartet + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + uint32_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; int t_id = gout_id * nsq_per_block + sq_id; int threads = nsq_per_block * gout_stride; @@ -670,13 +711,33 @@ int RYS_build_k(double *vk, double *dm, int n_dm, int nao, int n_tiles = ntiles_i * ntiles_j * ntiles_k * ntiles_l; auto launch = [&](auto offset, int tile_chunk) { + #ifdef USE_SYCL + sycl_get_queue()->memcpy(s_rys_gxyz_offset, gxyz_offset+offset, + tile_chunk*sizeof(GXYZOffset)).wait(); + #else checkCudaErrors( cudaMemcpyToSymbol(c_gxyz_offset, gxyz_offset+offset, tile_chunk*sizeof(GXYZOffset), 0, cudaMemcpyHostToDevice)); + #endif int scheme[4]; threads_scheme_for_k(scheme, bounds, shm_size, tile_chunk); int buflen = scheme[2]; + int reserved_shm_size = scheme[3]; + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(scheme[1], scheme[0]); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_k_kernel(dev_envs, kmat, bounds, q_cond_ij, q_cond_kl, dm_penalty, + s_cond_ij, s_cond_kl, diffuse_exps, pool, + head + offset/256, gout_pattern, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else if (buflen > 48000) { cudaFuncSetAttribute(rys_k_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); cudaError_t err = cudaGetLastError(); @@ -687,11 +748,11 @@ int RYS_build_k(double *vk, double *dm, int n_dm, int nao, } } dim3 threads(scheme[0], scheme[1]); - int reserved_shm_size = scheme[3]; rys_k_kernel<<>>( *envs, kmat, bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head + offset/256, gout_pattern, reserved_shm_size); + #endif }; launch(0, 256); diff --git a/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cuh b/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cuh index da5864af7..9ca6551a6 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cuh +++ b/gpu4pyscf/lib/gvhf-rys/rys_contract_k.cuh @@ -15,31 +15,144 @@ */ #pragma once - #include #include + #include "gvhf-rys/vhf.cuh" -__device__ __forceinline__ +// Abstracts 2D-kernel gout thread-index setup. Used 4x in this header. +#ifdef USE_SYCL +#define SETUP_GOUT_KERNEL() \ + auto item = syclex::this_work_item::get_nd_item<2>(); \ + int nsq_per_block = item.get_local_range(1); \ + int gout_id = item.get_local_id(0); \ + int gout_stride = item.get_local_range(0); +#else +#define SETUP_GOUT_KERNEL() \ + int nsq_per_block = blockDim.x; \ + int gout_id = threadIdx.y; \ + int gout_stride = blockDim.y; +#endif + +#ifdef USE_SYCL + +// Please mind that this is a copy of the values in rys_constant.cu +// Given the support for C++ seperate constexpr declaration and definition +// is until C++20, we do use this var as `static constexpr` +static constexpr int _c_cartesian_lexical_xyz[252] = { + // s + 0, 0, 0, + // p + 1, 0, 0, + 0, 1, 0, + 0, 0, 1, + // d + 2, 0, 0, + 1, 1, 0, + 1, 0, 1, + 0, 2, 0, + 0, 1, 1, + 0, 0, 2, + // f + 3, 0, 0, + 2, 1, 0, + 2, 0, 1, + 1, 2, 0, + 1, 1, 1, + 1, 0, 2, + 0, 3, 0, + 0, 2, 1, + 0, 1, 2, + 0, 0, 3, + // g + 4, 0, 0, + 3, 1, 0, + 3, 0, 1, + 2, 2, 0, + 2, 1, 1, + 2, 0, 2, + 1, 3, 0, + 1, 2, 1, + 1, 1, 2, + 1, 0, 3, + 0, 4, 0, + 0, 3, 1, + 0, 2, 2, + 0, 1, 3, + 0, 0, 4, + // h + 5, 0, 0, + 4, 1, 0, + 4, 0, 1, + 3, 2, 0, + 3, 1, 1, + 3, 0, 2, + 2, 3, 0, + 2, 2, 1, + 2, 1, 2, + 2, 0, 3, + 1, 4, 0, + 1, 3, 1, + 1, 2, 2, + 1, 1, 3, + 1, 0, 4, + 0, 5, 0, + 0, 4, 1, + 0, 3, 2, + 0, 2, 3, + 0, 1, 4, + 0, 0, 5, + // j + 6, 0, 0, + 5, 1, 0, + 5, 0, 1, + 4, 2, 0, + 4, 1, 1, + 4, 0, 2, + 3, 3, 0, + 3, 2, 1, + 3, 1, 2, + 3, 0, 3, + 2, 4, 0, + 2, 3, 1, + 2, 2, 2, + 2, 1, 3, + 2, 0, 4, + 1, 5, 0, + 1, 4, 1, + 1, 3, 2, + 1, 2, 3, + 1, 1, 4, + 1, 0, 5, + 0, 6, 0, + 0, 5, 1, + 0, 4, 2, + 0, 3, 3, + 0, 2, 4, + 0, 1, 5, + 0, 0, 6, +}; +#endif // USE_SYCL + + +DEVICE_INLINE int lex_xyz_offset(int l) { // the offsets for _c_cartesian_lexical_xyz = l*(l+1)*(l+2)/6 * 3 return l*(l+1)*(l+2) / 2; } -__device__ __forceinline__ +DEVICE_INLINE int lex_xyz_address(int l, int i) { return _c_cartesian_lexical_xyz[lex_xyz_offset(l) + i]; } template -__device__ __forceinline__ +DEVICE_INLINE void vrr(double *g, double *ri, double *rj, double *Rpq, double aj_aij, double rt_aij, double b10, int g_size) { - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; + SETUP_GOUT_KERNEL(); for (int n = gout_id; n < 3; n += gout_stride) { double *_gx = g + n * g_size * nsq_per_block; double Rpa = (rj[n] - ri[n]) * aj_aij; @@ -58,13 +171,11 @@ void vrr(double *g, double *ri, double *rj, double *Rpq, double aj_aij, double r } template -__device__ __forceinline__ +DEVICE_INLINE void trr(double *g, double *rlrk, double *Rpq, double al_akl, double rt_akl, double b00, double b01, int lij3, int stride_k, int g_size) { - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; + SETUP_GOUT_KERNEL(); for (int n = gout_id; n < lij3+gout_id; n += gout_stride) { __syncthreads(); int i = n / 3; //for i in range(lij+1): @@ -99,12 +210,10 @@ void trr(double *g, double *rlrk, double *Rpq, double al_akl, double rt_akl, } template -__device__ __forceinline__ +DEVICE_INLINE void hrr_ij(double *g, double *rjri, int count, int g_size) { - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; + SETUP_GOUT_KERNEL(); constexpr int lij = LI + LJ; constexpr int stride_j = LI + 1; constexpr int stride_k = stride_j * (LJ + 1); @@ -128,12 +237,10 @@ void hrr_ij(double *g, double *rjri, int count, int g_size) } } template -__device__ __forceinline__ +DEVICE_INLINE void hrr_kl(double *g, double *rlrk, int stride_k) { - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; + SETUP_GOUT_KERNEL(); constexpr int lkl = LK + LL; for (int n = gout_id; n < stride_k*3; n += gout_stride) { int i = n / 3; @@ -157,7 +264,7 @@ void hrr_kl(double *g, double *rlrk, int stride_k) } template -__device__ __forceinline__ +DEVICE_INLINE void inner_dot(double *gout, double *g, int *addr_i, int *addr_j, int *addr_k, int *addr_l) @@ -187,7 +294,7 @@ void inner_dot(double *gout, double *g, } } } } } -__device__ __forceinline__ +DEVICE_INLINE void load_dm(double *dm, double *dm_cache, int nao, int i0, int j0, int ioff, int joff, int nfi, int nfj) { @@ -206,7 +313,7 @@ void load_dm(double *dm, double *dm_cache, int nao, int i0, int j0, } template -__device__ __forceinline__ +DEVICE_INLINE void dot_dm(double *vk, double *dm_cache, double *gout, int nao, int i0, int l0, int ioff, int loff, int nfi, int nfl) { @@ -230,3 +337,5 @@ void dot_dm(double *vk, double *dm_cache, double *gout, int nao, int i0, int l0, atomicAdd(vk_local+i*nao+l, v); } } } + +#undef SETUP_GOUT_KERNEL diff --git a/gpu4pyscf/lib/gvhf-rys/rys_jk_driver.cu b/gpu4pyscf/lib/gvhf-rys/rys_jk_driver.cu index ad4e14871..b589a6790 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_jk_driver.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_jk_driver.cu @@ -23,6 +23,7 @@ #include "vhf.cuh" +#ifndef USE_SYCL __constant__ int c_nf[] = { 1, 3, @@ -46,6 +47,7 @@ __constant__ float c_div_nf[] = { 0.027778f, 0.022223f, }; +#endif extern "C" { int cuda_version() diff --git a/gpu4pyscf/lib/gvhf-rys/rys_roots.cu b/gpu4pyscf/lib/gvhf-rys/rys_roots.cu index 41c69ca50..59844f8ff 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_roots.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_roots.cu @@ -58,12 +58,12 @@ static void rys_roots(int nroots, double x, double *rw, return; } - double *datax = ROOT_RW_DATA + DEGREE1*INTERVALS * nroots*(nroots-1); + const double *datax = ROOT_RW_DATA + DEGREE1*INTERVALS * nroots*(nroots-1); int it = (int)(x * .4); double u = (x - it * 2.5) * 0.8 - 1.; double u2 = u * 2.; for (int i = rt_id; i < nroots*2; i += stride) { - double *c = datax + i * DEGREE1 * INTERVALS; + const double *c = datax + i * DEGREE1 * INTERVALS; //for i in range(2, degree + 1): // c0, c1 = c[degree-i] - c1, c0 + c1*u2 double c0 = c[it + DEGREE *INTERVALS]; @@ -87,10 +87,17 @@ static void rys_roots(int nroots, double x, double *rw, } // rys_roots for range-separation Coulomb -__device__ __forceinline__ +// NOTE (SYCL/PVC): see rys_roots_for_k.cu -- the barriers are skipped when +// stride == 1, because such callers evaluate every root on the calling +// work-item and invoke this routine from a work-item dependent loop, where a +// group barrier deadlocks on Level Zero. +__device__ __forceinline__ static void rys_roots_rs(int nroots, double theta, double rr, double omega, double *rw, int block_size, int rt_id, int stride) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + #endif double theta_rr = theta * rr; if (omega == 0) { rys_roots(nroots, theta_rr, rw, block_size, rt_id, stride); @@ -115,5 +122,5 @@ void rys_roots_rs(int nroots, double theta, double rr, double omega, rw1[ irys*2 *block_size] *= theta_fac; rw1[(irys*2+1)*block_size] *= sqrt_theta_fac; } - } + } } diff --git a/gpu4pyscf/lib/gvhf-rys/rys_roots.cuh b/gpu4pyscf/lib/gvhf-rys/rys_roots.cuh index 0ed1e3afa..6f4bf792b 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_roots.cuh +++ b/gpu4pyscf/lib/gvhf-rys/rys_roots.cuh @@ -4,6 +4,11 @@ #define INTERVAL 2.5 #define INTERVALS 40 +#ifdef USE_SYCL +#include +#include "rys_roots_dat.cu" + +#else extern __device__ double ROOT_SMALLX_R0[]; extern __device__ double ROOT_SMALLX_R1[]; extern __device__ double ROOT_SMALLX_W0[]; @@ -11,3 +16,4 @@ extern __device__ double ROOT_SMALLX_W1[]; extern __device__ double ROOT_LARGEX_R_DATA[]; extern __device__ double ROOT_LARGEX_W_DATA[]; extern __device__ double ROOT_RW_DATA[]; +#endif diff --git a/gpu4pyscf/lib/gvhf-rys/rys_roots_dat.cu b/gpu4pyscf/lib/gvhf-rys/rys_roots_dat.cu index d713f3f3e..ef244bc25 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_roots_dat.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_roots_dat.cu @@ -1,5 +1,9 @@ -__device__ -double ROOT_SMALLX_R0[] = { +#ifdef USE_SYCL +inline constexpr double ROOT_SMALLX_R0[] = +#else +__device__ double ROOT_SMALLX_R0[] = +#endif + { // nroots = 1 3.3333333333333333e-01, // nroots = 2 @@ -67,8 +71,12 @@ double ROOT_SMALLX_R0[] = { 9.8630441451934892e-01, }; -__device__ -double ROOT_SMALLX_R1[] = { +#ifdef USE_SYCL +inline constexpr double ROOT_SMALLX_R1[] = +#else +__device__ double ROOT_SMALLX_R1[] = +#endif + { // nroots = 1 -8.8888888888888888e-02, // nroots = 2 @@ -136,8 +144,12 @@ double ROOT_SMALLX_R1[] = { -6.5892763019479263e-04, }; -__device__ -double ROOT_SMALLX_W0[] = { +#ifdef USE_SYCL +inline constexpr double ROOT_SMALLX_W0[] = +#else +__device__ double ROOT_SMALLX_W0[] = +#endif + { // nroots = 1 1.0000000000000000e+00, // nroots = 2 @@ -205,8 +217,12 @@ double ROOT_SMALLX_W0[] = { 1.7614007139152118e-02, }; -__device__ -double ROOT_SMALLX_W1[] = { +#ifdef USE_SYCL +inline constexpr double ROOT_SMALLX_W1[] = +#else +__device__ double ROOT_SMALLX_W1[] = +#endif + { // nroots = 1 -3.3333333333333331e-01, // nroots = 2 @@ -274,7 +290,12 @@ double ROOT_SMALLX_W1[] = { -1.6531204416842745e-02, }; -__device__ double ROOT_LARGEX_R_DATA[] = { +#ifdef USE_SYCL +inline constexpr double ROOT_LARGEX_R_DATA[] = +#else +__device__ double ROOT_LARGEX_R_DATA[] = +#endif + { // nroots = 1 5.0000000000000000e-01, // nroots = 2 @@ -342,7 +363,12 @@ __device__ double ROOT_LARGEX_R_DATA[] = { 6.0192063149587915e-02, }; -__device__ double ROOT_LARGEX_W_DATA[] = { +#ifdef USE_SYCL +inline constexpr double ROOT_LARGEX_W_DATA[] = +#else +__device__ double ROOT_LARGEX_W_DATA[] = +#endif + { // nroots = 1 1.0000000000000000e+00, // nroots = 2 @@ -410,7 +436,12 @@ __device__ double ROOT_LARGEX_W_DATA[] = { 5.2158612689910977e-01, }; -__device__ double ROOT_RW_DATA[] = { +#ifdef USE_SYCL +inline constexpr double ROOT_RW_DATA[] = +#else +__device__ double ROOT_RW_DATA[] = +#endif + { // nroots = 1 // root 0 2.44302310464373174e-01, diff --git a/gpu4pyscf/lib/gvhf-rys/rys_roots_for_k.cu b/gpu4pyscf/lib/gvhf-rys/rys_roots_for_k.cu index 00b080cbf..461e50463 100644 --- a/gpu4pyscf/lib/gvhf-rys/rys_roots_for_k.cu +++ b/gpu4pyscf/lib/gvhf-rys/rys_roots_for_k.cu @@ -18,11 +18,27 @@ #include "gvhf-rys/rys_roots.cu" -__device__ __forceinline__ +// NOTE (SYCL/PVC): the barriers below must NOT be taken when `stride == 1`. +// A number of callers (the "unrolled" int3c2e/ejk kernels) pass +// stride=1, rt_id=0, meaning every work-item evaluates *all* of its own roots +// into its own `rw` slot; no cross-thread data is exchanged, so the barrier is +// semantically a no-op. Those callers also drive a work-item dependent loop +// for (idx = st_id; idx < nst; idx += nst_per_block) +// whose trip count differs between work-items. Executing a group barrier in +// such a loop is UB: on CUDA the hardware barrier ignores threads that already +// exited the kernel, so it happens to work, but on Level Zero the work-items +// that left the loop early never arrive and the work-group hangs forever. +// Guarding on `stride > 1` (uniform within the group in every caller) keeps the +// barrier exactly where it is actually needed - the cooperative gout_stride>1 +// callers, whose loops are uniform (`idx < nst + st_id`). +static __device__ __forceinline__ void rys_roots_for_k(int nroots, double theta, double rr, double *rw, double omega, double lr_factor, double sr_factor, int block_size, int stride, int rt_id) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); +#endif double theta_rr = theta * rr; if (omega == 0) { rys_roots(nroots, theta_rr, rw, block_size, rt_id, stride); @@ -58,13 +74,20 @@ void rys_roots_for_k(int nroots, double theta, double rr, double *rw, } } -__device__ __forceinline__ +static __device__ __forceinline__ void rys_roots_for_k(int nroots, double theta, double rr, double *rw, double omega, double lr_factor, double sr_factor) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int block_size = item.get_local_range(1); + int stride = item.get_local_range(0); + int rt_id = item.get_local_id(0); +#else int block_size = blockDim.x; int stride = blockDim.y; int rt_id = threadIdx.y; +#endif rys_roots_for_k(nroots, theta, rr, rw, omega, lr_factor, sr_factor, block_size, stride, rt_id); } diff --git a/gpu4pyscf/lib/gvhf-rys/unrolled_kernels.cuh b/gpu4pyscf/lib/gvhf-rys/unrolled_kernels.cuh index 0668c2494..b007b24c4 100644 --- a/gpu4pyscf/lib/gvhf-rys/unrolled_kernels.cuh +++ b/gpu4pyscf/lib/gvhf-rys/unrolled_kernels.cuh @@ -16,6 +16,107 @@ #include "gvhf-rys/vhf.cuh" +#ifdef USE_SYCL + +// --------------------------------------------------------------------- +// Per-translation-unit kernel-name disambiguation. +// +// unrolled_rys_jk.cu and unrolled_rys_k.cu BOTH define 19 kernels named +// rys_k_0000 .. rys_k_3200 with DIFFERENT bodies. +// +// RYS_UNROLLED_KERNEL_TAG is injected per source file by +// gvhf-rys/CMakeLists.txt. Do NOT define it inside the .cu files -- +// they are auto-generated upstream and must stay byte-identical. +// --------------------------------------------------------------------- +#ifndef RYS_UNROLLED_KERNEL_TAG +#error "RYS_UNROLLED_KERNEL_TAG is not defined. Every unrolled_*.cu that includes unrolled_kernels.cuh must get a unique tag via set_source_files_properties(... COMPILE_DEFINITIONS RYS_UNROLLED_KERNEL_TAG=) in gvhf-rys/CMakeLists.txt. Without it, identically-named kernels in different translation units silently alias." +#endif + +#define RYS_KERNEL_TAG_CAT_(KERNEL, TAG) KERNEL##_##TAG##_sycl +#define RYS_KERNEL_TAG_CAT(KERNEL, TAG) RYS_KERNEL_TAG_CAT_(KERNEL, TAG) +#define RYS_KERNEL_TAG(KERNEL) RYS_KERNEL_TAG_CAT(KERNEL, RYS_UNROLLED_KERNEL_TAG) + +#define dim3 sycl::range<2> + +#define JKMATRIX_KERNEL_ARGS \ + RysIntEnvVars envs, JKMatrix jk, BoundsInfo bounds, \ + float *q_cond_ij, float *q_cond_kl, float dm_penalty, \ + float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, \ + uint32_t *pool, int *head, \ + sycl::nd_item<2> &item, double *shared_memory + +#define JKMATRIX_KERNEL_SETUP() \ + int sq_id = item.get_local_id(1); \ + int gout_id = item.get_local_id(0); \ + int _nsq_per_block = item.get_local_range(1); \ + uint32_t *bas_kl_idx = pool + item.get_group(1) * QUEUE_DEPTH; \ + auto _rys_grp = item.get_group(); \ + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); + +#define LAUNCH_JKMATRIX_KERNEL(KERNEL) { \ + auto _rys_envs = *envs; auto _rys_jk = *jk; auto _rys_bounds = *bounds; \ + sycl::range<2> _rys_blocks(1, workers); \ + sycl::range<2> _rys_threads(gout_stride, nsq_per_block); \ + sycl_get_queue()->submit([&](sycl::handler &cgh) { \ + sycl::local_accessor _rys_lmem(sycl::range<1>(buflen), cgh); \ + cgh.parallel_for( \ + sycl::nd_range<2>(_rys_blocks * _rys_threads, _rys_threads), \ + [=](sycl::nd_item<2> item) { \ + KERNEL(_rys_envs, _rys_jk, _rys_bounds, q_cond_ij, q_cond_kl, \ + dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, head, \ + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(_rys_lmem)); \ + }); \ + }); \ + } + +#define JKENERGY_KERNEL_ARGS \ + RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, \ + float *q_cond_ij, float *q_cond_kl, float dm_penalty, \ + float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, \ + uint32_t *pool, double *dd_pool, int *head, \ + sycl::nd_item<2> &item, double *shared_memory + +#define JKENERGY_KERNEL_SETUP() \ + int sq_id = item.get_local_id(1); \ + int gout_id = item.get_local_id(0); \ + int worker_id = item.get_group(1); \ + auto _rys_grp = item.get_group(); \ + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); \ + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(_rys_grp); + +#define LAUNCH_JKENERGY_KERNEL(KERNEL) { \ + auto _rys_envs = *envs; auto _rys_jk = *jk; auto _rys_bounds = *bounds; \ + sycl::range<2> _rys_blocks(1, workers); \ + sycl::range<2> _rys_threads(gout_stride, nsq_per_block); \ + sycl_get_queue()->submit([&](sycl::handler &cgh) { \ + sycl::local_accessor _rys_lmem(sycl::range<1>(buflen), cgh); \ + cgh.parallel_for( \ + sycl::nd_range<2>(_rys_blocks * _rys_threads, _rys_threads), \ + [=](sycl::nd_item<2> item) { \ + KERNEL(_rys_envs, _rys_jk, _rys_bounds, q_cond_ij, q_cond_kl, \ + dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, \ + dd_pool, head, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(_rys_lmem)); \ + }); \ + }); \ + } + +#else // !USE_SYCL -- byte-identical to upstream/master + #define JKMATRIX_KERNEL_ARGS \ RysIntEnvVars envs, JKMatrix jk, BoundsInfo bounds, \ float *q_cond_ij, float *q_cond_kl, float dm_penalty, \ @@ -59,3 +160,5 @@ #define LAUNCH_JKENERGY_KERNEL(KERNEL) \ KERNEL<<>>( \ *envs, *jk, *bounds, q_cond_ij, q_cond_kl, dm_penalty, s_cond_ij, s_cond_kl, diffuse_exps, pool, dd_pool, head) + +#endif // USE_SYCL diff --git a/gpu4pyscf/lib/gvhf-rys/vhf.cuh b/gpu4pyscf/lib/gvhf-rys/vhf.cuh index 208a3bc8e..7a7fbdf5a 100644 --- a/gpu4pyscf/lib/gvhf-rys/vhf.cuh +++ b/gpu4pyscf/lib/gvhf-rys/vhf.cuh @@ -2,6 +2,12 @@ #include #include +#ifdef USE_SYCL +#include +#elif defined(__CUDACC__) +#include +#endif + #define PTR_RANGE_OMEGA 8 // slots of atm #define CHARGE_OF 0 @@ -32,6 +38,13 @@ #define MIN(x, y) ((x) < (y) ? (x) : (y)) #define MAX(x, y) ((x) > (y) ? (x) : (y)) +// Abstracts __device__ __forceinline__ (CUDA) vs static inline (SYCL) on device functions. +#ifdef USE_SYCL +#define DEVICE_INLINE static inline +#else +#define DEVICE_INLINE __device__ __forceinline__ +#endif + // 2*pi**2.5 #define PI_FAC 34.98683665524972497 @@ -161,4 +174,66 @@ extern __constant__ GXYZOffset c_gxyz_offset[]; extern __constant__ int c_nf[]; extern __constant__ float c_div_nf[]; -#endif + +#elif defined(USE_SYCL) + +static inline unsigned get_smid() +{ + auto max_cu = 448; + auto item = syclex::this_work_item::get_nd_item<2>(); + auto g = item.get_group_linear_id(); + return (g % max_cu); +} + +// NOTE: On CUDA, adjust_threads doubles the launch's nsq_per_block only when +// cudaFuncGetAttributes confirms the kernel's actual register usage supports +// 2x occupancy per SM. Each unrolled *_ip1 kernel hardcodes its own internal +// `constexpr int nsq_per_block` used to lay out shared/local memory, so the +// host-side "threads" value driving the nd_range and local_accessor buflen +// MUST stay equal to that constant. Unconditionally doubling it here (as a +// stand-in for the missing SYCL equivalent of cudaFuncGetAttributes) makes +// the launched work-group width diverge from the kernel's baked-in shared +// memory layout, corrupting the block-level reduction (silently wrong +// gradients/JK energies -- worst case is the simplest all-s-function case, +// e.g. RHF/H2 in a minimal basis, since that's the first switch-case hit). +// Until a real SYCL analogue of the CUDA register-occupancy query exists, +// this must be a no-op. +#define adjust_threads(kernel, threads) { } + +extern SYCL_EXTERNAL sycl_device_global s_rys_i_in_fold2idx; +extern SYCL_EXTERNAL sycl_device_global s_rys_i_in_fold3idx; + +//NOTE: `_c_cartesian_lexical_xyz` equvialent in SYCL is converted to +// `static constexpr` var defined in rys_contract_k.cuh becuase this +// particular header is being included in gvhf-rys/rys_contract_jk_ip1.cu, +// gvhf-rys/rys_contract_jk_ip2.cu files that uses this var. Hence it is not +// declared or defined here + +// Here 625 is just a random MAX chosen from rys_constant.cu +extern SYCL_EXTERNAL sycl_device_global s_rys_gxyz_offset; + +__constant__ int c_nf[] = { + 1, + 3, + 6, + 10, + 15, + 21, + 28, + 36, + 45, +}; + +__constant__ float c_div_nf[] = { + 1.f, + 0.333334f, + 0.166667f, + 0.100001f, + 0.066667f, + 0.047620f, + 0.035715f, + 0.027778f, + 0.022223f, +}; + +#endif // __CUDACC__ diff --git a/gpu4pyscf/lib/gvhf/CMakeLists.txt b/gpu4pyscf/lib/gvhf/CMakeLists.txt index 97bfeaaa4..ecb7110bc 100644 --- a/gpu4pyscf/lib/gvhf/CMakeLists.txt +++ b/gpu4pyscf/lib/gvhf/CMakeLists.txt @@ -14,7 +14,7 @@ # set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v -maxrregcount=128") -add_library(gvhf SHARED +set(GPU_SRCS constant.cu #nr_jk_driver.cu #nr_jk_driver_int3c2e_ip1.cu @@ -25,15 +25,24 @@ add_library(gvhf SHARED #get_veff_driver_ip1.cu #rys_roots_dat.cu ) - -#option(BUILD_SHARED_LIBS "build shared libraries" 1) -#option(ENABLE_STATIC "Enforce static library build" 0) -#if(ENABLE_STATIC) -# set(BUILD_SHARED_LIBS 0) -#endif() +add_library(gvhf SHARED ${GPU_SRCS}) set_target_properties(gvhf PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_SEPARABLE_COMPILATION ON) + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) -target_link_libraries(gvhf gint) +if (USE_SYCL) + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + target_compile_options(gvhf PRIVATE -x c++ -nocudainc -nocudalib) + # gvhf/constant.cuh declares its bpcache as s_gvhf_bpcache rather than + # gint's s_bpcache. Both are default-visibility device_global objects, and + # libgint.so and libgvhf.so are co-resident with no DT_NEEDED link between + # them, so identical names would alias to whichever definition the dynamic + # linker resolved first and let one library's host-side memcpy land in the + # other's device image. The rename keeps them distinct; this define maps + # any inlined gint source that still spells s_bpcache onto gvhf's copy. + target_compile_definitions(gvhf PRIVATE s_bpcache=s_gvhf_bpcache) + target_link_libraries(gvhf PRIVATE sycl_compat gsycl gint) +else (USE_SYCL) + set_target_properties(gvhf PROPERTIES CUDA_SEPARABLE_COMPILATION ON) + target_link_libraries(gvhf PRIVATE gint) +endif (USE_SYCL) diff --git a/gpu4pyscf/lib/gvhf/constant.cu b/gpu4pyscf/lib/gvhf/constant.cu index 7e0ea1a4b..59e0f4ab1 100644 --- a/gpu4pyscf/lib/gvhf/constant.cu +++ b/gpu4pyscf/lib/gvhf/constant.cu @@ -16,21 +16,27 @@ #include "constant.cuh" +#ifdef USE_SYCL +SYCL_EXTERNAL sycl_device_global s_gvhf_bpcache; +//__constant__ int16_t c_idx4c[NFffff*3]; +#else __constant__ BasisProdCache c_bpcache; //__constant__ int16_t c_idx4c[NFffff*3]; // Generated with GINTinit_index1d_xyz __constant__ int c_idx[TOT_NF*3] = { - 0, 1, 0, 0, 2, 1, 1, 0, 0, 0, 3, 2, 2, 1, 1, 1, 0, 0, 0, 0, 4, 3, 3, - 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 5, 4, 4, 3, 3, 3, 2, 2, 2, 2, 1, - 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 6, 5, 5, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, - 2, 2, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 2, - 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, - 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, - 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, 1, 0, - 6, 5, 4, 3, 2, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 2, 0, 0, 1, 0, 1, 2, - 0, 1, 2, 3, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 0, 1, 0, - 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 0, 1, 0, 1, 2, + 0, 1, 0, 0, 2, 1, 1, 0, 0, 0, 3, 2, 2, 1, 1, 1, 0, 0, 0, 0, 4, 3, 3, + 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 5, 4, 4, 3, 3, 3, 2, 2, 2, 2, 1, + 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 6, 5, 5, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, + 2, 2, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 2, + 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, + 3, 2, 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, + 1, 0, 0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0, 5, 4, 3, 2, 1, 0, + 6, 5, 4, 3, 2, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 2, 0, 0, 1, 0, 1, 2, + 0, 1, 2, 3, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 0, 1, 0, + 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6}; __constant__ int c_l_locs[GPU_LMAX+2] = {0, 1, 4, 10, 20, 35, 56, 84}; + +#endif diff --git a/gpu4pyscf/lib/gvhf/constant.cuh b/gpu4pyscf/lib/gvhf/constant.cuh index cf82935cd..1e173e455 100644 --- a/gpu4pyscf/lib/gvhf/constant.cuh +++ b/gpu4pyscf/lib/gvhf/constant.cuh @@ -3,9 +3,24 @@ #include "gint/gint.h" +#ifdef USE_SYCL +#include + +// Named distinctly from gint's s_bpcache (gint/cint2e.cuh) -- both are +// GLOBAL DEFAULT-visibility device_global objects and libgint.so/libgvhf.so +// are co-resident in the process with no DT_NEEDED link between them, so +// identical names alias to whichever library's definition the dynamic +// linker resolves first. That let gvhf's host-side bpcache memcpy target +// gint's device image (or vice versa), corrupting whichever kernel ran +// concurrently on the other library's queue. See +// hang_analysis_evidence/DEFECT5_free_and_device_global_audit.md, Finding A. +extern SYCL_EXTERNAL sycl_device_global s_gvhf_bpcache; + +#else // USE_SYCL extern __constant__ BasisProdCache c_bpcache; -//extern __constant__ int16_t c_idx4c[NFffff*3]; +//extern __constaont__ int16_t c_idx4c[NFffff*3]; extern __constant__ int c_idx[TOT_NF*3]; extern __constant__ int c_l_locs[GPU_LMAX+2]; +#endif // USE_SYCL #endif //GPU4PYSCF_CONSTANT_CUH diff --git a/gpu4pyscf/lib/multigrid/CMakeLists.txt b/gpu4pyscf/lib/multigrid/CMakeLists.txt index 812c39f14..7c386c117 100644 --- a/gpu4pyscf/lib/multigrid/CMakeLists.txt +++ b/gpu4pyscf/lib/multigrid/CMakeLists.txt @@ -11,22 +11,14 @@ # LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} # CUDA_SEPARABLE_COMPILATION ON) -add_library(mgrid_v2 SHARED - multigrid_v2/drivers.cu - multigrid_v2/screen.cu - multigrid_v2/eval_xc.cu - multigrid_v2/eval_xc_grad.cu +set(GPU_SRCS_v2 + multigrid_v2/drivers.cu + multigrid_v2/screen.cu + multigrid_v2/eval_xc.cu + multigrid_v2/eval_xc_grad.cu ) -set_target_properties(mgrid_v2 PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_SEPARABLE_COMPILATION ON -) -if(ENABLE_FP32_MULTIGRID) - target_compile_definitions(mgrid_v2 PRIVATE GPU4PYSCF_ENABLE_FP32_MULTIGRID) -endif() - -add_library(mgrid_v3 SHARED +set(GPU_SRCS_v3 multigrid_v3/screen.cu multigrid_v3/drivers.cu multigrid_v3/eval_density.cu @@ -43,7 +35,34 @@ add_library(mgrid_v3 SHARED multigrid_v3/aft_eval_mgga_mat.cu ) -set_target_properties(mgrid_v3 PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_SEPARABLE_COMPILATION ON -) +add_library(mgrid_v3 SHARED ${GPU_SRCS_v3}) + +if (USE_SYCL) + # mgrid_v2 is not the default multigrid engine (multigrid_v3 is) and is + # not SYCL-ported (multigrid_v2.py already tolerates a missing + # libmgrid_v2 via try/except OSError, same as the dead v1 "mgrid" + # target above). Only mgrid_v3 is built here. + file(GLOB CUH_HEADERS_v3 + "${CMAKE_CURRENT_SOURCE_DIR}/multigrid_v3/*.cuh" + ) + set_source_files_properties(${CUH_HEADERS_v3} PROPERTIES LANGUAGE CXX) + set_source_files_properties(${GPU_SRCS_v3} PROPERTIES LANGUAGE CXX) + + set_target_properties(mgrid_v3 PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + target_compile_options(mgrid_v3 PRIVATE -x c++ -nocudainc -nocudalib) + target_link_libraries(mgrid_v3 PRIVATE sycl_compat gsycl) +else() + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128") + + add_library(mgrid_v2 SHARED ${GPU_SRCS_v2}) + set_target_properties(mgrid_v2 PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} + CUDA_SEPARABLE_COMPILATION ON) + if(ENABLE_FP32_MULTIGRID) + target_compile_definitions(mgrid_v2 PRIVATE GPU4PYSCF_ENABLE_FP32_MULTIGRID) + endif() + + set_target_properties(mgrid_v3 PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} + CUDA_SEPARABLE_COMPILATION ON) +endif (USE_SYCL) diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_density.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_density.cu index 653767db4..b26504e4e 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_density.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_density.cu @@ -17,9 +17,11 @@ #include #include #include +#ifndef USE_SYCL #include #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "gvhf-rys/rys_contract_k.cuh" #include "constant_objects.cuh" @@ -44,17 +46,36 @@ void orth_aopair_dm_kernel(double *outR, double *outI, double *dm, int *mesh_cum, int *nimgs_cum, int ntiles, double factor) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int sp_block_id = item.get_group(0) / ntiles; + int tile_id = item.get_group(0) % ntiles; +#else int thread_id = threadIdx.x; - int x_id = thread_id / NGV_PER_BLOCK; - int Gv_id = thread_id % NGV_PER_BLOCK; int sp_block_id = blockIdx.x / ntiles; int tile_id = blockIdx.x % ntiles; +#endif + int x_id = thread_id / NGV_PER_BLOCK; + int Gv_id = thread_id % NGV_PER_BLOCK; +#ifdef USE_SYCL + auto &gx = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &swap = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &mesh_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ri = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &rj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ij_offset = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &fac = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ai = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ double gx[NGV_PER_BLOCK*3*2*LMAX1*LMAX1]; __shared__ double swap[NGV_PER_BLOCK*3*2*(LMAX+LMAX+1)]; __shared__ int mesh_start[3]; __shared__ int ri, rj; __shared__ size_t ij_offset; __shared__ double fac, ai, aj; +#endif int *bas = envs.bas; int nbas = envs.nbas; @@ -232,9 +253,21 @@ int orth_contract_aopair_dm(double *outR, double *outI, int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; int ntiles = ntiles_x * ntiles_y * ntiles_z; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(ntiles*nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + orth_aopair_dm_kernel( + outR, outI, dm, dev_envs, shl_pair_offsets, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, ntiles, factor); + }).wait(); +#else orth_aopair_dm_kernel<<>>( outR, outI, dm, *envs, shl_pair_offsets, bas_ij_idx, G_bases, L_bases, mesh_cum, nimgs_cum, ntiles, factor); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in orth_aopair_dm_kernel: %s\n", cudaGetErrorString(err)); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_lda_mat.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_lda_mat.cu index 294e1c9a1..10a2e28d2 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_lda_mat.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_lda_mat.cu @@ -17,15 +17,27 @@ #include #include #include +#ifndef USE_SYCL #include #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "gvhf-rys/rys_contract_k.cuh" #include "constant_objects.cuh" #include "utils.cuh" #include "aft_recursion.cuh" +#ifdef USE_SYCL +// CUDA cuDoubleComplex stand-in. No SYCL analogue exists anywhere in the +// codebase; this POD exactly matches cuDoubleComplex's memory layout (two +// doubles) so device-pointer reinterpretation from Python callers is +// unaffected. Unlike sycl::double2 (whose .x()/.y() are methods), this is a +// real struct with real members, so existing .x/.y field-access and +// brace-init call sites need zero further changes. +struct alignas(16) cuDoubleComplex { double x, y; }; +#endif + #define WARP_SIZE 32 #define WARPS 8 #define THREADS 256 @@ -47,11 +59,33 @@ void orth_lda_mat_kernel(double *out, cuDoubleComplex *vxcG, int *mesh_cum, int *nimgs_cum, int npair, int ntiles_x, int ntiles_y, int ntiles_z) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int pair_id = item.get_group(0) % npair; +#else int thread_id = threadIdx.x; + int pair_id = blockIdx.x % npair; +#endif int x_id = thread_id / NGV_PER_BLOCK; int Gv_id = thread_id % NGV_PER_BLOCK; +#ifdef USE_SYCL + auto &tile_batch = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + if (thread_id == 0) { + tile_batch = item.get_group(0) / npair; + } + auto &gx = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &swap = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &mesh_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &vjR = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ri = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &rj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ai = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ int tile_batch; - int pair_id = blockIdx.x % npair; if (thread_id == 0) { tile_batch = blockIdx.x / npair; } @@ -61,6 +95,7 @@ void orth_lda_mat_kernel(double *out, cuDoubleComplex *vxcG, __shared__ double vjR[NCART_MAX*NCART_MAX * WARPS]; __shared__ int ri, rj, li, lj; __shared__ double ai, aj; +#endif int mesh_x = mesh_cum[1] - mesh_cum[0]; int mesh_y = mesh_cum[2] - mesh_cum[1]; @@ -245,9 +280,21 @@ int orth_aft_lda_mat(double *out, cuDoubleComplex *vxcG, cuDoubleComplex *placeh int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; int ntiles = ntiles_x * ntiles_y * ntiles_z; int ntile_batch = (ntiles + TILES_PER_BATCH-1) / TILES_PER_BATCH; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(ntile_batch*npair); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + orth_lda_mat_kernel( + out, vxcG, dev_envs, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, npair, ntiles_x, ntiles_y, ntiles_z); + }).wait(); +#else orth_lda_mat_kernel<<>>( out, vxcG, *envs, bas_ij_idx, G_bases, L_bases, mesh_cum, nimgs_cum, npair, ntiles_x, ntiles_y, ntiles_z); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in orth_lda_mat_kernel: %s\n", cudaGetErrorString(err)); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_mgga_mat.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_mgga_mat.cu index 9d415a738..672f3fc7c 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_mgga_mat.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_mgga_mat.cu @@ -17,15 +17,27 @@ #include #include #include +#ifndef USE_SYCL #include #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "gvhf-rys/rys_contract_k.cuh" #include "constant_objects.cuh" #include "utils.cuh" #include "aft_recursion.cuh" +#ifdef USE_SYCL +// CUDA cuDoubleComplex stand-in. No SYCL analogue exists anywhere in the +// codebase; this POD exactly matches cuDoubleComplex's memory layout (two +// doubles) so device-pointer reinterpretation from Python callers is +// unaffected. Unlike sycl::double2 (whose .x()/.y() are methods), this is a +// real struct with real members, so existing .x/.y field-access and +// brace-init call sites need zero further changes. +struct alignas(16) cuDoubleComplex { double x, y; }; +#endif + #define WARP_SIZE 32 #define WARPS 8 #define THREADS 256 @@ -44,11 +56,32 @@ void orth_mgga_mat_kernel(double *out, cuDoubleComplex *vrhoG, int *mesh_cum, int *nimgs_cum, int npair, int ntiles_x, int ntiles_y, int ntiles_z) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int pair_id = item.get_group(0) % npair; +#else int thread_id = threadIdx.x; + int pair_id = blockIdx.x % npair; +#endif int x_id = thread_id / NGV_PER_BLOCK; int Gv_id = thread_id % NGV_PER_BLOCK; +#ifdef USE_SYCL + auto &tile_batch = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + if (thread_id == 0) { + tile_batch = item.get_group(0) / npair; + } + auto &gx = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &mesh_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &vjR = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ri = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &rj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ai = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ int tile_batch; - int pair_id = blockIdx.x % npair; if (thread_id == 0) { tile_batch = blockIdx.x / npair; } @@ -57,6 +90,7 @@ void orth_mgga_mat_kernel(double *out, cuDoubleComplex *vrhoG, __shared__ double vjR[NCART_MAX*NCART_MAX * WARPS]; __shared__ int ri, rj, li, lj; __shared__ double ai, aj; +#endif int mesh_x = mesh_cum[1] - mesh_cum[0]; int mesh_y = mesh_cum[2] - mesh_cum[1]; @@ -297,9 +331,21 @@ int orth_aft_mgga_mat(double *out, cuDoubleComplex *vrhoG, cuDoubleComplex *vtau int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; int ntiles = ntiles_x * ntiles_y * ntiles_z; int ntile_batch = (ntiles + TILES_PER_BATCH-1) / TILES_PER_BATCH; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(ntile_batch*npair); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + orth_mgga_mat_kernel( + out, vrhoG, vtauG, dev_envs, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, npair, ntiles_x, ntiles_y, ntiles_z); + }).wait(); +#else orth_mgga_mat_kernel<<>>( out, vrhoG, vtauG, *envs, bas_ij_idx, G_bases, L_bases, mesh_cum, nimgs_cum, npair, ntiles_x, ntiles_y, ntiles_z); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in orth_mgga_mat_kernel: %s\n", cudaGetErrorString(err)); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_tau.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_tau.cu index 6084ee630..93724587e 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_tau.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_tau.cu @@ -17,9 +17,11 @@ #include #include #include +#ifndef USE_SYCL #include #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "gvhf-rys/rys_contract_k.cuh" #include "constant_objects.cuh" @@ -39,17 +41,36 @@ void orth_ft_tau_dm_kernel(double *densityR, double *densityI, double *tauR, dou int64_t *bas_ij_idx, double *G_bases, double *L_bases, int *mesh_cum, int *nimgs_cum, int ntiles, double factor) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int sp_block_id = item.get_group(0) / ntiles; + int tile_id = item.get_group(0) % ntiles; +#else int thread_id = threadIdx.x; - int x_id = thread_id / NGV_PER_BLOCK; - int Gv_id = thread_id % NGV_PER_BLOCK; int sp_block_id = blockIdx.x / ntiles; int tile_id = blockIdx.x % ntiles; +#endif + int x_id = thread_id / NGV_PER_BLOCK; + int Gv_id = thread_id % NGV_PER_BLOCK; +#ifdef USE_SYCL + auto &gx = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &swap = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &mesh_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ri = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &rj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ij_offset = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &fac = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ai = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ double gx[NGV_PER_BLOCK*3*2*(LMAX1+1)*(LMAX1+1)]; __shared__ double swap[NGV_PER_BLOCK*3*2*(LMAX+LMAX+3)]; __shared__ int mesh_start[3]; __shared__ int ri, rj; __shared__ size_t ij_offset; __shared__ double fac, ai, aj; +#endif int *bas = envs.bas; int nbas = envs.nbas; @@ -250,9 +271,21 @@ int orth_contract_ft_tau_dm(double *densityR, double *densityI, int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; int ntiles = ntiles_x * ntiles_y * ntiles_z; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(ntiles*nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + orth_ft_tau_dm_kernel( + densityR, densityI, tauR, tauI, dm, dev_envs, shl_pair_offsets, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, ntiles, factor); + }).wait(); +#else orth_ft_tau_dm_kernel<<>>( densityR, densityI, tauR, tauI, dm, *envs, shl_pair_offsets, bas_ij_idx, G_bases, L_bases, mesh_cum, nimgs_cum, ntiles, factor); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in orth_ft_tau_dm_kernel: %s\n", cudaGetErrorString(err)); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_recursion.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_recursion.cuh index 7d8ff37b7..e4968b28c 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_recursion.cuh +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_recursion.cuh @@ -16,6 +16,9 @@ #define NGV_PER_BLOCK 16 +#ifdef USE_SYCL +inline +#endif __forceinline__ __device__ void vrr_hrr(double *gx, double *swap, int addrR, int li, int lj, int stride_j, double a2, double xjxi, double aj_aij, @@ -153,6 +156,9 @@ void vrr_hrr(double *gx, int addrR, int stride_j, double a2, double xjxi, } } +#ifdef USE_SYCL +inline +#endif __forceinline__ __device__ void dI_gx(double *gx, int addr, int stride_i, int li, double ai2, double &outR, double &outI) @@ -165,6 +171,9 @@ void dI_gx(double *gx, int addr, int stride_i, int li, } } +#ifdef USE_SYCL +inline +#endif __forceinline__ __device__ void dIdJ_gx(double *gx, int addr, int stride_i, int stride_j, int li, int lj, double ai2, double aj2, double &outR, double &outI) @@ -189,6 +198,9 @@ void dIdJ_gx(double *gx, int addr, int stride_i, int stride_j, int li, int lj, } } +#ifdef USE_SYCL +inline +#endif __forceinline__ __device__ void dIdJ_gx(double *gx, int addr, int stride_i, int li, double ai2, double kx, double &outR, double &outI) @@ -217,6 +229,9 @@ void dIdJ_gx(double *gx, int addr, int stride_i, int li, outI = f1R * kx - outI; } +#ifdef USE_SYCL +inline +#endif __forceinline__ __device__ void d2IdJ_gx(double *gx, int addr, int stride_i, int stride_j, int li, int lj, double ai2, double aj2, double &outR, double &outI) diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/cartesian.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/cartesian.cuh index 29ff2009f..e2f083b7a 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/cartesian.cuh +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/cartesian.cuh @@ -16,6 +16,9 @@ #pragma once +#ifdef USE_SYCL +inline +#endif __forceinline__ __device__ void rename_registers(double i_cartesian[], int start, int stop, int count) { diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/constant_objects.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/constant_objects.cuh index b36d71e00..650813521 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/constant_objects.cuh +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/constant_objects.cuh @@ -16,10 +16,33 @@ #pragma once +#ifdef USE_SYCL +#include + +extern SYCL_EXTERNAL sycl_device_global s_c_lattice_vectors; +extern SYCL_EXTERNAL sycl_device_global s_c_reciprocal_lattice_vectors; // norm to 1 +extern SYCL_EXTERNAL sycl_device_global s_c_dxyz_dabc; + +#define c_lattice_vectors (s_c_lattice_vectors.get()) +#define c_reciprocal_lattice_vectors (s_c_reciprocal_lattice_vectors.get()) +#define c_dxyz_dabc (s_c_dxyz_dabc.get()) + +// c_nf / c_div_nf are defined unconditionally in gvhf-rys/vhf.cuh, which +// every multigrid_v3 TU includes. Under USE_SYCL __constant__ expands to +// `inline constexpr`, so defining them here too is an ODR redefinition +// error. Inherit vhf.cuh's tables (identical values) instead. + +// CUDA float2 stand-in. sycl::float2's element accessors are methods +// (v.x()), not members (v.x), so aliasing to it would break every +// `.x`/`.y` site in screen.cu and the eval_*_v2/strain_grad kernels. This +// POD keeps both branches identical and supports brace-init assignment. +struct alignas(8) float2 { float x, y; }; +#else extern __constant__ double c_lattice_vectors[9]; extern __constant__ double c_reciprocal_lattice_vectors[9]; // norm to 1 extern __constant__ double c_dxyz_dabc[9]; extern __constant__ int c_nf[]; extern __constant__ float c_div_nf[]; +#endif #define NBAS_MAX 16777216 diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/drivers.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/drivers.cu index ab8847dbb..7ddbe731a 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/drivers.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/drivers.cu @@ -15,11 +15,20 @@ */ #include +#ifndef USE_SYCL #include #include #include #include +#endif +#include "constant_objects.cuh" +#ifdef USE_SYCL +SYCL_EXTERNAL sycl_device_global s_c_lattice_vectors; +SYCL_EXTERNAL sycl_device_global s_c_reciprocal_lattice_vectors; +SYCL_EXTERNAL sycl_device_global s_c_dxyz_dabc; +// c_nf / c_div_nf are plain constexpr tables defined in constant_objects.cuh +#else __constant__ double c_lattice_vectors[9]; __constant__ double c_reciprocal_lattice_vectors[9]; __constant__ double c_dxyz_dabc[9]; @@ -47,16 +56,25 @@ __constant__ float c_div_nf[] = { 0.027778f, 0.022223f, }; +#endif // input[nc,nx,ny,nz], output[nc,mx,my,mz] __global__ static void fft_take_kernel(double2* __restrict__ out, double2* __restrict__ in, int mx, int my, int mz, int nx, int ny, int nz, int nc) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int x = item.get_group(0); + int y = item.get_group(1) * item.get_local_range(1) + item.get_local_id(1); + int tx = item.get_local_id(0); + int threadsx = item.get_local_range(0); +#else int x = blockIdx.x; int y = blockIdx.y * blockDim.y + threadIdx.y; int tx = threadIdx.x; int threadsx = blockDim.x; +#endif if (x >= mx || y >= my) return; int sx = x; @@ -81,10 +99,18 @@ __global__ static void fft_takebak_kernel(double2* __restrict__ out, double2* __restrict__ in, int mx, int my, int mz, int nx, int ny, int nz, int nc) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int x = item.get_group(0); + int y = item.get_group(1) * item.get_local_range(1) + item.get_local_id(1); + int tx = item.get_local_id(0); + int threadsx = item.get_local_range(0); +#else int x = blockIdx.x; int y = blockIdx.y * blockDim.y + threadIdx.y; int tx = threadIdx.x; int threadsx = blockDim.x; +#endif if (x >= mx || y >= my) return; int sx = x; @@ -99,8 +125,12 @@ void fft_takebak_kernel(double2* __restrict__ out, double2* __restrict__ in, for (int c = 0; c < nc; ++c) { size_t dst = (((size_t)c*nx + sx)*ny + sy)*nz + sz; size_t src = (((size_t)c*mx + x )*my + y )*mz + z; +#ifdef USE_SYCL + out[dst] = double2{out[dst].x() + in[src].x(), out[dst].y() + in[src].y()}; +#else out[dst].x += in[src].x; out[dst].y += in[src].y; +#endif } } } @@ -109,12 +139,21 @@ extern "C" { void update_lattice_vectors(double *lattice_vectors, double *reciprocal_lattice_vectors) { +#ifdef USE_SYCL + sycl_get_queue()->memcpy(s_c_lattice_vectors, lattice_vectors, 9 * sizeof(double)); + sycl_get_queue()->memcpy(s_c_reciprocal_lattice_vectors, reciprocal_lattice_vectors, 9 * sizeof(double)).wait(); +#else cudaMemcpyToSymbol(c_lattice_vectors, lattice_vectors, 9 * sizeof(double)); cudaMemcpyToSymbol(c_reciprocal_lattice_vectors, reciprocal_lattice_vectors, 9 * sizeof(double)); +#endif } void update_dxyz_dabc(double *dxyz_dabc) { +#ifdef USE_SYCL + sycl_get_queue()->memcpy(s_c_dxyz_dabc, dxyz_dabc, 9 * sizeof(double)).wait(); +#else cudaMemcpyToSymbol(c_dxyz_dabc, dxyz_dabc, 9 * sizeof(double)); +#endif } int fft_take(double2 *out, double2 *in, int *out_shape, int *in_shape, int counts) @@ -122,10 +161,20 @@ int fft_take(double2 *out, double2 *in, int *out_shape, int *in_shape, int count int mx = out_shape[0]; int my = out_shape[1]; int mz = out_shape[2]; +#ifdef USE_SYCL + int nx = in_shape[0], ny = in_shape[1], nz = in_shape[2]; + sycl::range<2> threads(32, 16); + sycl::range<2> grids(mx, (my+15)/16); + sycl_get_queue()->parallel_for + (sycl::nd_range<2>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + fft_take_kernel(out, in, mx, my, mz, nx, ny, nz, counts); + }).wait(); +#else dim3 threads(32, 16); dim3 grids(mx, (my+15)/16); fft_take_kernel<<>>( out, in, mx, my, mz, in_shape[0], in_shape[1], in_shape[2], counts); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in fft_take kernel: %s\n", cudaGetErrorString(err)); @@ -139,10 +188,20 @@ int fft_takebak(double2 *out, double2 *in, int *out_shape, int *in_shape, int co int mx = in_shape[0]; int my = in_shape[1]; int mz = in_shape[2]; +#ifdef USE_SYCL + int nx = out_shape[0], ny = out_shape[1], nz = out_shape[2]; + sycl::range<2> threads(32, 16); + sycl::range<2> grids(mx, (my+15)/16); + sycl_get_queue()->parallel_for + (sycl::nd_range<2>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + fft_takebak_kernel(out, in, mx, my, mz, nx, ny, nz, counts); + }).wait(); +#else dim3 threads(32, 16); dim3 grids(mx, (my+15)/16); fft_takebak_kernel<<>>( out, in, mx, my, mz, out_shape[0], out_shape[1], out_shape[2], counts); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in fft_takebak kernel: %s\n", cudaGetErrorString(err)); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_density.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_density.cu index 2f893d7ac..537a9e98c 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_density.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_density.cu @@ -17,13 +17,20 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "constant_objects.cuh" #include "cartesian.cuh" #include "utils.cuh" +#ifdef USE_SYCL +#define CONCAT_(a,b) a##b +#define CONCAT(a,b) CONCAT_(a,b) +#endif + #define TILE 4 #define WARP_SIZE 32 #define THREADS 64 @@ -40,11 +47,25 @@ void eval_density_kernel(double *density, double *dm, PBCIntEnvVars envs, { constexpr int threads = THREADS; constexpr int WARPS = THREADS / WARP_SIZE; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int tile_id0 = item.get_group(0) * tiles_per_block; + auto thread_block = item.get_group(); + int &a_upper = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &b_upper = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &c_upper = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &start_position_x = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &start_position_y = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &start_position_z = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + auto &density_value = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else int thread_id = threadIdx.x; int tile_id0 = blockIdx.x * tiles_per_block; __shared__ int a_upper, b_upper, c_upper; __shared__ double start_position_x, start_position_y, start_position_z; __shared__ double density_value[TILE*TILE*TILE*WARPS]; +#endif constexpr int nfi = (LI + 1) * (LI + 2) / 2; constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; @@ -305,6 +326,24 @@ for (int tile_id = tile_id0; tile_id < min(tile_id0+tiles_per_block, ntiles); ti } extern "C" { +#ifdef USE_SYCL +#define eval_density_kernel_case(li, lj, slice_i, slice_j, non_orth) \ + case (li * LMAX1 + lj): { \ + sycl::range<1> sycl_threads(THREADS); \ + sycl::range<1> sycl_grids(block_grid); \ + auto dev_envs = *envs; \ + sycl_get_queue()->parallel_for \ + (sycl::nd_range<1>(sycl_grids * sycl_threads, sycl_threads), [=](auto item) [[intel::kernel_args_restrict]] { \ + eval_density_kernel( \ + density, dm, dev_envs, supmol_img_coords, factor, \ + shl_pair_offsets, dressed_bas_ij_idx, \ + grid_tile_index, n_contributing_tiles, tiles_per_block, \ + a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ + mesh_a, mesh_b, mesh_c, negligible); \ + }).wait(); \ + } \ + break +#else #define eval_density_kernel_case(li, lj, slice_i, slice_j, non_orth) \ case (li * LMAX1 + lj): \ eval_density_kernel<<>>( \ @@ -314,6 +353,7 @@ extern "C" { a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ mesh_a, mesh_b, mesh_c, negligible); \ break +#endif int evaluate_density(double *density, double *placeholder, double *dm, PBCIntEnvVars *envs, diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat_v2.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat_v2.cu index bd4f917bf..0b6d113f1 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat_v2.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat_v2.cu @@ -17,13 +17,20 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "constant_objects.cuh" #include "cartesian.cuh" #include "utils.cuh" +#ifdef USE_SYCL +#define CONCAT_(a,b) a##b +#define CONCAT(a,b) CONCAT_(a,b) +#endif + template __global__ static void eval_lda_mat_kernel_v2(double *out, double *vxc_weights, PBCIntEnvVars envs, @@ -32,12 +39,45 @@ void eval_lda_mat_kernel_v2(double *out, double *vxc_weights, PBCIntEnvVars envs int mesh_a, int mesh_b, int mesh_c, int npairs, double negligible) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + constexpr int tile = 16; + int tx = item.get_local_id(1); + int ty = item.get_local_id(0); + int thread_id = ty * tile + tx; + int pair_id = item.get_group(1); +#else constexpr int tile = 16; int tx = threadIdx.x; int ty = threadIdx.y; int thread_id = ty * tile + tx; int pair_id = blockIdx.x; +#endif +#ifdef USE_SYCL + auto &a_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_center = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ij_offset = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &cc = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &exp_da_squared = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &theta_rr = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &swap = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ int a_start, a_stop, a_center; __shared__ int b_start, b_stop; __shared__ int c_start, c_stop; @@ -47,6 +87,7 @@ void eval_lda_mat_kernel_v2(double *out, double *vxc_weights, PBCIntEnvVars envs __shared__ double xj, yj, zj; __shared__ double xij, yij, zij, aij, theta_rr; __shared__ double swap[8]; +#endif int *bas = envs.bas; double *env = envs.env; @@ -244,6 +285,18 @@ void eval_lda_mat_kernel_v2(double *out, double *vxc_weights, PBCIntEnvVars envs } extern "C" { +#ifdef USE_SYCL +#define eval_lda_mat_kernel_v2_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + sycl_get_queue()->parallel_for \ + (sycl::nd_range<2>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { \ + eval_lda_mat_kernel_v2( \ + out, vxc, dev_envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + negligible); \ + }).wait(); \ + break +#else #define eval_lda_mat_kernel_v2_case(li, lj, slice_i, slice_j) \ case (li * LMAX1 + lj): \ eval_lda_mat_kernel_v2<<>>( \ @@ -251,6 +304,7 @@ extern "C" { da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ negligible); \ break +#endif int evaluate_lda_mat_v2(double *out, double *vxc, double *placeholder, PBCIntEnvVars *envs, double *dxyz_dabc, int li, int lj, int64_t *bas_ij_idx, @@ -263,7 +317,13 @@ int evaluate_lda_mat_v2(double *out, double *vxc, double *placeholder, PBCIntEnv double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); +#ifdef USE_SYCL + sycl::range<2> threads(16, 16); + sycl::range<2> grids(1, npairs); + auto dev_envs = *envs; +#else dim3 threads(16, 16); +#endif switch (li * LMAX1 + lj) { eval_lda_mat_kernel_v2_case(0,0, 1, 1); eval_lda_mat_kernel_v2_case(1,0, 3, 1); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_strain_grad.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_strain_grad.cu index 1a53bc3f2..f52d996ef 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_strain_grad.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_strain_grad.cu @@ -17,8 +17,10 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "constant_objects.cuh" #include "cartesian.cuh" @@ -28,6 +30,11 @@ #define WARP_SIZE 32 #define THREADS 64 +#ifdef USE_SYCL +#define CONCAT_(a,b) a##b +#define CONCAT(a,b) CONCAT_(a,b) +#endif + template __global__ static void eval_lda_grad_kernel(double *grad, double *strain, double *dm, @@ -37,15 +44,52 @@ void eval_lda_grad_kernel(double *grad, double *strain, double *dm, int mesh_a, int mesh_b, int mesh_c, int npairs, double factor, double negligible) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + constexpr int tile = 16; + int tx = item.get_local_id(1); + int ty = item.get_local_id(0); + int thread_id = ty * tile + tx; + int pair_id = item.get_group(1); +#else constexpr int tile = 16; int tx = threadIdx.x; int ty = threadIdx.y; int thread_id = ty * tile + tx; int pair_id = blockIdx.x; +#endif constexpr int nfi = (LI + 1) * (LI + 2) / 2; constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; +#ifdef USE_SYCL + auto &a_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_center = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &cc = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &exp_da_squared = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ai = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &theta_rr = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xjxi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yjyi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zjzi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &dm_cache = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ int a_start, a_stop, a_center; __shared__ int b_start, b_stop; __shared__ int c_start, c_stop; @@ -55,6 +99,7 @@ void eval_lda_grad_kernel(double *grad, double *strain, double *dm, __shared__ double xij, yij, zij, ai, aj, aij, theta_rr; __shared__ double xjxi, yjyi, zjzi; __shared__ double dm_cache[nfi*nfj]; +#endif int *bas = envs.bas; double *env = envs.env; @@ -351,6 +396,18 @@ void eval_lda_grad_kernel(double *grad, double *strain, double *dm, } extern "C" { +#ifdef USE_SYCL +#define eval_lda_grad_kernel_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + sycl_get_queue()->parallel_for \ + (sycl::nd_range<2>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { \ + eval_lda_grad_kernel( \ + grad, strain, dm, vxc, dev_envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + factor, negligible); \ + }).wait(); \ + break +#else #define eval_lda_grad_kernel_case(li, lj, slice_i, slice_j) \ case (li * LMAX1 + lj): \ eval_lda_grad_kernel<<>>( \ @@ -358,6 +415,7 @@ extern "C" { da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ factor, negligible); \ break +#endif int evaluate_lda_grad(double *grad, double *strain, double *dm, double *vxc, double *placeholder, PBCIntEnvVars *envs, @@ -371,7 +429,13 @@ int evaluate_lda_grad(double *grad, double *strain, double *dm, double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); +#ifdef USE_SYCL + sycl::range<2> threads(16, 16); + sycl::range<2> grids(1, npairs); + auto dev_envs = *envs; +#else dim3 threads(16, 16); +#endif switch (li * LMAX1 + lj) { eval_lda_grad_kernel_case(0,0, 1, 1); eval_lda_grad_kernel_case(1,0, 3, 1); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat_v2.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat_v2.cu index 164f7f67b..fedc57330 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat_v2.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat_v2.cu @@ -17,13 +17,20 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "constant_objects.cuh" #include "cartesian.cuh" #include "utils.cuh" +#ifdef USE_SYCL +#define CONCAT_(a,b) a##b +#define CONCAT(a,b) CONCAT_(a,b) +#endif + template __global__ static void eval_mgga_mat_kernel_v2(double *out, double *vrho_weights, double *vtau_weights, @@ -33,12 +40,47 @@ void eval_mgga_mat_kernel_v2(double *out, double *vrho_weights, double *vtau_wei int mesh_a, int mesh_b, int mesh_c, int npairs, double negligible) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + constexpr int tile = 16; + int tx = item.get_local_id(1); + int ty = item.get_local_id(0); + int thread_id = ty * tile + tx; + int pair_id = item.get_group(1); +#else constexpr int tile = 16; int tx = threadIdx.x; int ty = threadIdx.y; int thread_id = ty * tile + tx; int pair_id = blockIdx.x; +#endif +#ifdef USE_SYCL + auto &a_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_center = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ij_offset = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &cc = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &exp_da_squared = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ai = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &theta_rr = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &swap = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ int a_start, a_stop, a_center; __shared__ int b_start, b_stop; __shared__ int c_start, c_stop; @@ -48,6 +90,7 @@ void eval_mgga_mat_kernel_v2(double *out, double *vrho_weights, double *vtau_wei __shared__ double xj, yj, zj; __shared__ double xij, yij, zij, ai, aj, aij, theta_rr; __shared__ double swap[8]; +#endif int *bas = envs.bas; double *env = envs.env; @@ -291,6 +334,18 @@ void eval_mgga_mat_kernel_v2(double *out, double *vrho_weights, double *vtau_wei } extern "C" { +#ifdef USE_SYCL +#define eval_mgga_mat_kernel_v2_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + sycl_get_queue()->parallel_for \ + (sycl::nd_range<2>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { \ + eval_mgga_mat_kernel_v2( \ + out, vxc, tau, dev_envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + negligible); \ + }).wait(); \ + break +#else #define eval_mgga_mat_kernel_v2_case(li, lj, slice_i, slice_j) \ case (li * LMAX1 + lj): \ eval_mgga_mat_kernel_v2<<>>( \ @@ -298,6 +353,7 @@ extern "C" { da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ negligible); \ break +#endif int evaluate_mgga_mat_v2(double *out, double *vxc, double *tau, PBCIntEnvVars *envs, double *dxyz_dabc, int li, int lj, int64_t *bas_ij_idx, @@ -310,7 +366,13 @@ int evaluate_mgga_mat_v2(double *out, double *vxc, double *tau, PBCIntEnvVars *e double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); +#ifdef USE_SYCL + sycl::range<2> threads(16, 16); + sycl::range<2> grids(1, npairs); + auto dev_envs = *envs; +#else dim3 threads(16, 16); +#endif switch (li * LMAX1 + lj) { eval_mgga_mat_kernel_v2_case(0,0, 1, 1); eval_mgga_mat_kernel_v2_case(1,0, 3, 1); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_strain_grad.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_strain_grad.cu index 2db72f30e..d29dedb3b 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_strain_grad.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_strain_grad.cu @@ -17,8 +17,10 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "constant_objects.cuh" #include "cartesian.cuh" @@ -28,6 +30,11 @@ #define WARP_SIZE 32 #define THREADS 64 +#ifdef USE_SYCL +#define CONCAT_(a,b) a##b +#define CONCAT(a,b) CONCAT_(a,b) +#endif + template __global__ static void eval_mgga_grad_kernel(double *grad, double *strain, double *dm, @@ -37,15 +44,52 @@ void eval_mgga_grad_kernel(double *grad, double *strain, double *dm, int mesh_a, int mesh_b, int mesh_c, int npairs, double factor, double negligible) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + constexpr int tile = 16; + int tx = item.get_local_id(1); + int ty = item.get_local_id(0); + int thread_id = ty * tile + tx; + int pair_id = item.get_group(1); +#else constexpr int tile = 16; int tx = threadIdx.x; int ty = threadIdx.y; int thread_id = ty * tile + tx; int pair_id = blockIdx.x; +#endif constexpr int nfi = (LI + 1) * (LI + 2) / 2; constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; +#ifdef USE_SYCL + auto &a_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &a_center = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &b_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_start = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &c_stop = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &cc = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &exp_da_squared = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &ai = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aj = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &aij = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &theta_rr = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &xjxi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &yjyi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &zjzi = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + auto &dm_cache = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ int a_start, a_stop, a_center; __shared__ int b_start, b_stop; __shared__ int c_start, c_stop; @@ -55,6 +99,7 @@ void eval_mgga_grad_kernel(double *grad, double *strain, double *dm, __shared__ double xij, yij, zij, ai, aj, aij, theta_rr; __shared__ double xjxi, yjyi, zjzi; __shared__ double dm_cache[nfi*nfj]; +#endif int *bas = envs.bas; double *env = envs.env; @@ -483,6 +528,18 @@ void eval_mgga_grad_kernel(double *grad, double *strain, double *dm, } extern "C" { +#ifdef USE_SYCL +#define eval_mgga_grad_kernel_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + sycl_get_queue()->parallel_for \ + (sycl::nd_range<2>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { \ + eval_mgga_grad_kernel( \ + grad, strain, dm, vxc, tau, dev_envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + factor, negligible); \ + }).wait(); \ + break +#else #define eval_mgga_grad_kernel_case(li, lj, slice_i, slice_j) \ case (li * LMAX1 + lj): \ eval_mgga_grad_kernel<<>>( \ @@ -490,6 +547,7 @@ extern "C" { da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ factor, negligible); \ break +#endif int evaluate_mgga_grad(double *grad, double *strain, double *dm, double *vxc, double *tau, PBCIntEnvVars *envs, @@ -503,7 +561,13 @@ int evaluate_mgga_grad(double *grad, double *strain, double *dm, double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); +#ifdef USE_SYCL + sycl::range<2> threads(16, 16); + sycl::range<2> grids(1, npairs); + auto dev_envs = *envs; +#else dim3 threads(16, 16); +#endif switch (li * LMAX1 + lj) { eval_mgga_grad_kernel_case(0,0, 1, 1); eval_mgga_grad_kernel_case(1,0, 3, 1); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_tau.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_tau.cu index 06d9a6f23..8b8d1c53d 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_tau.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_tau.cu @@ -17,13 +17,20 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "constant_objects.cuh" #include "cartesian.cuh" #include "utils.cuh" +#ifdef USE_SYCL +#define CONCAT_(a,b) a##b +#define CONCAT(a,b) CONCAT_(a,b) +#endif + #define TILE 4 #define WARP_SIZE 32 #define THREADS 64 @@ -40,12 +47,27 @@ void eval_tau_kernel(double *density, double *tau, double *dm, PBCIntEnvVars env { constexpr int threads = THREADS; constexpr int WARPS = THREADS / WARP_SIZE; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int tile_id0 = item.get_group(0) * tiles_per_block; + auto thread_block = item.get_group(); + int &a_upper = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &b_upper = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &c_upper = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &start_position_x = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &start_position_y = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &start_position_z = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + auto &rho_value = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + auto &tau_value = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else int thread_id = threadIdx.x; int tile_id0 = blockIdx.x * tiles_per_block; __shared__ int a_upper, b_upper, c_upper; __shared__ double start_position_x, start_position_y, start_position_z; __shared__ double rho_value[TILE*TILE*TILE*WARPS]; __shared__ double tau_value[TILE*TILE*TILE*WARPS]; +#endif constexpr int nfi = (LI + 1) * (LI + 2) / 2; constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; @@ -323,6 +345,24 @@ for (int tile_id = tile_id0; tile_id < min(tile_id0+tiles_per_block, ntiles); ti } extern "C" { +#ifdef USE_SYCL +#define eval_tau_kernel_case(li, lj, slice_i, slice_j, non_orth) \ + case (li * LMAX1 + lj): { \ + sycl::range<1> sycl_threads(THREADS); \ + sycl::range<1> sycl_grids(block_grid); \ + auto dev_envs = *envs; \ + sycl_get_queue()->parallel_for \ + (sycl::nd_range<1>(sycl_grids * sycl_threads, sycl_threads), [=](auto item) [[intel::kernel_args_restrict]] { \ + eval_tau_kernel( \ + density, tau, dm, dev_envs, supmol_img_coords, factor, \ + shl_pair_offsets, dressed_bas_ij_idx, \ + grid_tile_index, n_contributing_tiles, tiles_per_block, \ + a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ + mesh_a, mesh_b, mesh_c, negligible); \ + }).wait(); \ + } \ + break +#else #define eval_tau_kernel_case(li, lj, slice_i, slice_j, non_orth) \ case (li * LMAX1 + lj): \ eval_tau_kernel<<>>( \ @@ -332,6 +372,7 @@ extern "C" { a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ mesh_a, mesh_b, mesh_c, negligible); \ break +#endif int evaluate_tau(double *density, double *tau, double *dm, PBCIntEnvVars *envs, double *dxyz_dabc, double *supmol_img_coords, diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/screen.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/screen.cu index 9fb1662b9..5d5c09c90 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/screen.cu +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/screen.cu @@ -18,8 +18,10 @@ #include #include #include +#ifndef USE_SYCL #include #include +#endif #include "gvhf-rys/vhf.cuh" #include "constant_objects.cuh" #include "cartesian.cuh" @@ -58,7 +60,12 @@ void grid_ranges_kernel(float2 *grid_frac_ranges, float *pair_ke, int npairs, float log_threshold, float undressed_threshold, float ke_max) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int pair_id = item.get_group(0) * item.get_local_range(0) + item.get_local_id(0); +#else int pair_id = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (pair_id >= npairs) return; int *bas = envs.bas; @@ -183,7 +190,12 @@ void grid_range_to_tiles_kernel(int *grid_tile_idx, int64_t *dressed_bas_ij, int mesh_x, int mesh_y, int mesh_z, int npairs, int nbas, int *head) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int pair_id = item.get_group(0) * item.get_local_range(0) + item.get_local_id(0); +#else int pair_id = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (pair_id >= npairs) return; int64_t bas_ij = bas_ij_idx[pair_id]; @@ -274,8 +286,14 @@ __global__ static void ovlp_mask_estimation_kernel(int8_t *ovlp_mask, PBCIntEnvVars envs, double *img_coords, int nimgs, float log_cutoff) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int jsh = item.get_group(0) * item.get_local_range(0) + item.get_local_id(0); + int ish = item.get_group(1) * item.get_local_range(1) + item.get_local_id(1); +#else int jsh = blockIdx.x * blockDim.x + threadIdx.x; int ish = blockIdx.y * blockDim.y + threadIdx.y; +#endif int nbas = envs.nbas; int bvk_nbas = envs.nbas * envs.bvk_ncells; if (ish >= nbas || jsh >= bvk_nbas) { @@ -338,7 +356,12 @@ void estimate_aft_Ecut_kernel(float *Ecut, int64_t *bas_ij_idx, PBCIntEnvVars en double *img_coords, int nimgs, int npairs, float log_cutoff, float Ecut_max, int is_mgga) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int pair_id = item.get_group(0) * item.get_local_range(0) + item.get_local_id(0); +#else int pair_id = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (pair_id >= npairs) { return; } @@ -412,13 +435,20 @@ void supmol_non_trivial_pairs_kernel(int64_t *supmol_bas_ij, int64_t *bas_ij_idx PBCIntEnvVars envs, int npairs, float log_cutoff, int is_mgga, int *head) { + constexpr int batch_size = 64; +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int pair_id = item.get_group(0) * item.get_local_range(0) + thread_id; + auto &img_cache = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else int thread_id = threadIdx.x; int pair_id = blockIdx.x * blockDim.x + thread_id; + __shared__ int8_t img_cache[THREADS*batch_size]; +#endif if (pair_id >= npairs) { return; } - constexpr int batch_size = 64; - __shared__ int8_t img_cache[THREADS*batch_size]; int bvk_nbas = envs.nbas * envs.bvk_ncells; int nimgs = envs.nimgs; int *bas = envs.bas; @@ -498,9 +528,21 @@ int gaussian_prod_grid_ranges(float2 *grid_frac_ranges, float *pair_ke, float undressed_threshold, float ke_max) { int batches = (npairs + THREADS-1) / THREADS; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(batches); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + grid_ranges_kernel( + grid_frac_ranges, pair_ke, Ecut_by_shell, dev_envs, bas_ij_idx, + li_inc, lj_inc, npairs, log_threshold, undressed_threshold, ke_max); + }).wait(); +#else grid_ranges_kernel<<>>( grid_frac_ranges, pair_ke, Ecut_by_shell, *envs, bas_ij_idx, li_inc, lj_inc, npairs, log_threshold, undressed_threshold, ke_max); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in gaussian_prod_grid_ranges: %s\n", cudaGetErrorString(err)); @@ -521,9 +563,20 @@ int grid_range_to_tiles(int *grid_tile_idx, int64_t *dressed_bas_ij, int mesh_y = mesh[1]; int mesh_z = mesh[2]; int batches = (npairs + THREADS-1) / THREADS; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(batches); + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + grid_range_to_tiles_kernel( + grid_tile_idx, dressed_bas_ij, bas_ij_idx, grid_frac_ranges, + nimgs_x, nimgs_y, nimgs_z, mesh_x, mesh_y, mesh_z, npairs, nbas, head); + }).wait(); +#else grid_range_to_tiles_kernel<<>>( grid_tile_idx, dressed_bas_ij, bas_ij_idx, grid_frac_ranges, nimgs_x, nimgs_y, nimgs_z, mesh_x, mesh_y, mesh_z, npairs, nbas, head); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in grid_range_to_tiles: %s\n", cudaGetErrorString(err)); @@ -537,10 +590,21 @@ int bvk_ovlp_mask_estimation(int8_t *ovlp_mask, PBCIntEnvVars *envs, { int nbas = envs->nbas; int bvk_nbas = nbas * envs->bvk_ncells; +#ifdef USE_SYCL + sycl::range<2> threads(16, 16); + sycl::range<2> grids((bvk_nbas + 15) / 16, (nbas + 15) / 16); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<2>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + ovlp_mask_estimation_kernel( + ovlp_mask, dev_envs, img_coords, nimgs, log_cutoff); + }).wait(); +#else dim3 threads(16, 16); dim3 blocks((bvk_nbas + 15) / 16, (nbas + 15) / 16); ovlp_mask_estimation_kernel<<>>( ovlp_mask, *envs, img_coords, nimgs, log_cutoff); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in bvk_ovlp_mask_estimation: %s\n", cudaGetErrorString(err)); @@ -554,8 +618,19 @@ int supmol_non_trivial_pairs(int64_t *supmol_bas_ij, int64_t *bas_ij_idx, { cudaMemset(head, 0, sizeof(int)); int blocks = (npairs + THREADS-1)/THREADS; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(blocks); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + supmol_non_trivial_pairs_kernel( + supmol_bas_ij, bas_ij_idx, dev_envs, npairs, log_cutoff, is_mgga, head); + }).wait(); +#else supmol_non_trivial_pairs_kernel<<>>( supmol_bas_ij, bas_ij_idx, *envs, npairs, log_cutoff, is_mgga, head); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in bvk_ovlp_mask_estimation: %s\n", cudaGetErrorString(err)); @@ -569,9 +644,21 @@ int estimate_aft_Ecut(float *Ecut, int64_t *bas_ij_idx, PBCIntEnvVars *envs, float log_cutoff, float Ecut_max, int is_mgga) { int blocks = (npairs + THREADS-1)/THREADS; +#ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> grids(blocks); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for + (sycl::nd_range<1>(grids * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + estimate_aft_Ecut_kernel( + Ecut, bas_ij_idx, dev_envs, img_coords, nimgs, npairs, log_cutoff, + Ecut_max, is_mgga); + }).wait(); +#else estimate_aft_Ecut_kernel<<>>( Ecut, bas_ij_idx, *envs, img_coords, nimgs, npairs, log_cutoff, Ecut_max, is_mgga); +#endif cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in raw_ovlp_mask: %s\n", cudaGetErrorString(err)); diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/utils.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/utils.cuh index 808904ea1..d9d3474ca 100644 --- a/gpu4pyscf/lib/multigrid/multigrid_v3/utils.cuh +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/utils.cuh @@ -21,6 +21,9 @@ __host__ __device__ T distance_squared(const T x, const T y, const T z) { return x * x + y * y + z * z; } +#ifdef USE_SYCL +inline +#endif __device__ __forceinline__ void multiply(double aR, double aI, double bR, double bI, double &cR, double &cI) { @@ -30,9 +33,15 @@ void multiply(double aR, double aI, double bR, double bI, double &cR, double &cI cI = outI; } +#ifdef USE_SYCL +inline +#endif __device__ __forceinline__ double reduce(double val, double *swap, int thread_id) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); +#endif constexpr int WARP_SIZE = 32; for (int offset = 16; offset > 0; offset >>= 1) { val += __shfl_down_sync(0xffffffff, val, offset); diff --git a/gpu4pyscf/lib/onemkl_helper/CMakeLists.txt b/gpu4pyscf/lib/onemkl_helper/CMakeLists.txt new file mode 100644 index 000000000..06ecd298c --- /dev/null +++ b/gpu4pyscf/lib/onemkl_helper/CMakeLists.txt @@ -0,0 +1,28 @@ +# gpu4pyscf is a plugin to use Nvidia GPU in PySCF package +# +# Copyright (C) 2022 Qiming Sun +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +add_library(onemkl_helper SHARED onemkl_lapack.cpp) + +set_target_properties(onemkl_helper PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} + CXX_STANDARD 17 + CXX_STANDARD_REQUIRED YES +) + +find_package(MKL REQUIRED CONFIG PATHS "$ENV{MKLROOT}") + +target_link_libraries(onemkl_helper PRIVATE sycl_compat gsycl MKL::MKL_SYCL) diff --git a/gpu4pyscf/lib/onemkl_helper/onemkl_lapack.cpp b/gpu4pyscf/lib/onemkl_helper/onemkl_lapack.cpp new file mode 100644 index 000000000..3c5d3ed5e --- /dev/null +++ b/gpu4pyscf/lib/onemkl_helper/onemkl_lapack.cpp @@ -0,0 +1,212 @@ +#include +#include +#include +#include + + +extern "C" void onemkl_trsm(double* a, double* b, + int m, int n, int lda, int ldb, + int lower, int trans, int unit_diagonal) { + auto queue = *sycl_get_queue(); + + oneapi::mkl::uplo uplo = lower ? oneapi::mkl::uplo::L : oneapi::mkl::uplo::U; + oneapi::mkl::transpose transA = trans ? oneapi::mkl::transpose::T : oneapi::mkl::transpose::N; + oneapi::mkl::diag diag = unit_diagonal ? oneapi::mkl::diag::U : oneapi::mkl::diag::N; + double alpha = 1.0; + + // in-place + auto e = oneapi::mkl::blas::column_major::trsm(queue, + oneapi::mkl::side::left, + uplo, + transA, + diag, + m, n, alpha, + a, lda, b, ldb); + e.wait(); +} + +extern "C" void onemkl_strsm(float* a, float* b, + int m, int n, int lda, int ldb, + int lower, int trans, int unit_diagonal) { + auto queue = *sycl_get_queue(); + + oneapi::mkl::uplo uplo = lower ? oneapi::mkl::uplo::L : oneapi::mkl::uplo::U; + oneapi::mkl::transpose transA = trans ? oneapi::mkl::transpose::T : oneapi::mkl::transpose::N; + oneapi::mkl::diag diag = unit_diagonal ? oneapi::mkl::diag::U : oneapi::mkl::diag::N; + float alpha = 1.0f; + + // in-place + auto e = oneapi::mkl::blas::column_major::trsm(queue, + oneapi::mkl::side::left, + uplo, + transA, + diag, + m, n, alpha, + a, lda, b, ldb); + e.wait(); +} + +extern "C" void onemkl_dsygvd_scratchpad_size(int itype, + int n, + int lda, + int ldb, + int* scratch_size) { + try { + auto queue = *sycl_get_queue(); + *scratch_size = oneapi::mkl::lapack::sygvd_scratchpad_size(queue, + itype, + oneapi::mkl::job::vec, + oneapi::mkl::uplo::lower, + n, lda, ldb); + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception: " << e.what() << std::endl; + } +} +extern "C" void onemkl_zhegvd_scratchpad_size(int itype, + int n, + int lda, + int ldb, + int* scratch_size) { + try { + auto queue = *sycl_get_queue(); + *scratch_size = oneapi::mkl::lapack::hegvd_scratchpad_size>(queue, + itype, + oneapi::mkl::job::vec, + oneapi::mkl::uplo::lower, + n, lda, ldb); + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception: " << e.what() << std::endl; + } +} + +extern "C" void onemkl_dsygvd(int itype, + int n, + double* A, + int lda, + double* B, + int ldb, + double* w, + double* scratchpad, + int scratchpad_size) { + try { + auto queue = *sycl_get_queue(); + auto e = oneapi::mkl::lapack::sygvd(queue, + itype, + oneapi::mkl::job::vec, + oneapi::mkl::uplo::lower, + n, + A, lda, + B, ldb, + w, scratchpad, scratchpad_size); + e.wait(); + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception: " << e.what() << std::endl; + } +} +extern "C" void onemkl_zhegvd(int itype, + int n, + std::complex* A, + int lda, + std::complex* B, + int ldb, + double* w, + std::complex* scratchpad, + int scratchpad_size) { + try { + auto queue = *sycl_get_queue(); + auto e = oneapi::mkl::lapack::hegvd(queue, + itype, + oneapi::mkl::job::vec, + oneapi::mkl::uplo::lower, + n, + A, lda, + B, ldb, + w, scratchpad, scratchpad_size); + e.wait(); + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception: " << e.what() << std::endl; + } +} + + + +extern "C" int64_t onemkl_dpotrf_scratchpad_size(int n, + int lda) { + try { + auto queue = *sycl_get_queue(); + auto scratch_size = oneapi::mkl::lapack::potrf_scratchpad_size(queue, + oneapi::mkl::uplo::upper, + n, lda); + return scratch_size; + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception: " << e.what() << std::endl; + } +} +extern "C" int64_t onemkl_zpotrf_scratchpad_size(int n, + int lda) { + try { + auto queue = *sycl_get_queue(); + auto scratch_size = oneapi::mkl::lapack::potrf_scratchpad_size>(queue, + oneapi::mkl::uplo::upper, + n, lda); + return scratch_size; + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception: " << e.what() << std::endl; + } +} + + +// Returns 0 on success, non-zero if the factorization failed. cuSOLVER +// reports this through dev_info and cusolver.py raises LinAlgError on it; +// callers in df.py rely on that to fall back to an eigendecomposition when +// j2c is singular, so the status has to survive the wrapper. +extern "C" int onemkl_dpotrf(int n, + double* A, + int lda, + double* scratchpad, + int scratchpad_size) { + try { + auto queue = *sycl_get_queue(); + auto e = oneapi::mkl::lapack::potrf(queue, + oneapi::mkl::uplo::upper, + n, + A, lda, + scratchpad, scratchpad_size); + e.wait(); + } catch (oneapi::mkl::lapack::exception const& e) { + // info() > 0 means the leading minor of that order is not positive + // definite -- the expected signal for a singular input, not a bug. + return e.info() ? static_cast(e.info()) : 1; + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception in dpotrf: " << e.what() << std::endl; + return 1; + } + return 0; +} +// Returns 0 on success, non-zero if the factorization failed. cuSOLVER +// reports this through dev_info and cusolver.py raises LinAlgError on it; +// callers in df.py rely on that to fall back to an eigendecomposition when +// j2c is singular, so the status has to survive the wrapper. +extern "C" int onemkl_zpotrf(int n, + std::complex* A, + int lda, + std::complex* scratchpad, + int scratchpad_size) { + try { + auto queue = *sycl_get_queue(); + auto e = oneapi::mkl::lapack::potrf(queue, + oneapi::mkl::uplo::upper, + n, + A, lda, + scratchpad, scratchpad_size); + e.wait(); + } catch (oneapi::mkl::lapack::exception const& e) { + // info() > 0 means the leading minor of that order is not positive + // definite -- the expected signal for a singular input, not a bug. + return e.info() ? static_cast(e.info()) : 1; + } catch (sycl::exception const& e) { + std::cerr << "SYCL exception in zpotrf: " << e.what() << std::endl; + return 1; + } + return 0; +} diff --git a/gpu4pyscf/lib/onemkl_lapack.py b/gpu4pyscf/lib/onemkl_lapack.py new file mode 100644 index 000000000..69c574416 --- /dev/null +++ b/gpu4pyscf/lib/onemkl_lapack.py @@ -0,0 +1,246 @@ +# Copyright 2021-2024 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""oneMKL-backed stand-in for gpu4pyscf.lib.cusolver on the SYCL backend. + +cusolver.py calls cuSOLVER through `cupy_backends.cuda.libs.cusolver`, which +does not exist under SYCL. gpu4pyscf/cupy/__init__.py therefore redirects both +`cupy_backends.cuda.libs.cusolver` and `gpu4pyscf.lib.cusolver` to this module, +so anything importing either name (x2c/x2c.py, lib/cupy_helper.py) reaches +oneMKL instead. + +This module must mirror cusolver.py's public surface exactly -- `eigh`, +`cholesky`, `LinAlgError` and `MAX_EIGH_DIM` -- with the same signatures and +the same failure behaviour, since callers are shared, unmodified upstream code. + +Only the routines with no dpnp equivalent are routed through libonemkl_helper: +`dsygvd`/`zhegvd` for the generalized problem Hx = (lambda)Sx (dpnp.linalg.eigh +takes a single matrix) and `trsm`. `potrf` is kept here too so that Cholesky +failures raise this module's LinAlgError rather than dpnp's, matching what the +CUDA path does. +""" + +import numpy as np +import dpnp +import dpctl +import ctypes +import os + +# workspace size (lwork) provided by the cusolver*_bufferSize is an 32-bit +# integer. For arrays above this dimension, the workspace size would overflow. +MAX_EIGH_DIM = 23150 + +CUSOLVER_EIG_TYPE_1 = 1 + + +_lib_dir = os.path.dirname(os.path.abspath(__file__)) +libonemkl = ctypes.CDLL(os.path.join(_lib_dir, 'libonemkl_helper.so')) + +libonemkl.onemkl_dsygvd_scratchpad_size.argtypes = [ + ctypes.c_int, # itype + ctypes.c_int, # n + ctypes.c_int, # lda + ctypes.c_int, # ldb + ctypes.c_void_p # *scratchpad_size +] +libonemkl.onemkl_zhegvd_scratchpad_size.argtypes = [ + ctypes.c_int, # itype + ctypes.c_int, # n + ctypes.c_int, # lda + ctypes.c_int, # ldb + ctypes.c_void_p # *scratchpad_size +] + + +libonemkl.onemkl_dsygvd.argtypes = [ + ctypes.c_int, # itype + ctypes.c_int, # n + ctypes.c_void_p, # *A + ctypes.c_int, # lda + ctypes.c_void_p, # *B + ctypes.c_int, # ldb + ctypes.c_void_p, # *w + ctypes.c_void_p, # *scratchpad + ctypes.c_int # scratchpad_size +] +libonemkl.onemkl_zhegvd.argtypes = [ + ctypes.c_int, # itype + ctypes.c_int, # n + ctypes.c_void_p, # *A + ctypes.c_int, # lda + ctypes.c_void_p, # *B + ctypes.c_int, # ldb + ctypes.c_void_p, # *w + ctypes.c_void_p, # *scratchpad + ctypes.c_int # scratchpad_size +] + + +libonemkl.onemkl_dpotrf_scratchpad_size.argtypes = [ + ctypes.c_int, # n + ctypes.c_int # lda +] +libonemkl.onemkl_dpotrf_scratchpad_size.restype = ctypes.c_int64 +libonemkl.onemkl_dpotrf.restype = ctypes.c_int +libonemkl.onemkl_zpotrf.restype = ctypes.c_int +libonemkl.onemkl_zpotrf_scratchpad_size.argtypes = [ + ctypes.c_int, # n + ctypes.c_int # lda +] +libonemkl.onemkl_zpotrf_scratchpad_size.restype = ctypes.c_int64 + +libonemkl.onemkl_dpotrf.argtypes = [ + ctypes.c_int, # n + ctypes.c_void_p, # *A + ctypes.c_int, # lda + ctypes.c_void_p, # *scratchpad + ctypes.c_int # scratchpad_size +] +libonemkl.onemkl_zpotrf.argtypes = [ + ctypes.c_int, # n + ctypes.c_void_p, # *A + ctypes.c_int, # lda + ctypes.c_void_p, # *scratchpad + ctypes.c_int # scratchpad_size +] + +_buffersize = {} +def eigh(h, s, overwrite=False): + """ + Solve the generalized eigenvalue problem Hx = λ Sx using oneMKL. + """ + assert h.dtype == s.dtype + assert h.dtype in (np.float64, np.complex128) + n = h.shape[0] + if h.dtype == np.complex128 and h.flags.c_contiguous: + # zhegvd requires the matrices in F-order. For hermitian matrices, + # .T.copy() is equivalent to .conj() + A = h.conj() + B = s.conj() + elif overwrite: + A = h + B = s + else: + A = h.copy() + B = s.copy() + + # Create buffers for A, B, and w + # https://github.com/IntelPython/dpctl/issues/888 + w = dpnp.zeros(n) + + # TODO: reuse workspace + if (h.dtype, n) in _buffersize: + lwork = _buffersize[h.dtype, n] + else: + lwork = ctypes.c_int(0) + if h.dtype == np.float64: + fn = libonemkl.onemkl_dsygvd_scratchpad_size + else: + fn = libonemkl.onemkl_zhegvd_scratchpad_size + fn( + CUSOLVER_EIG_TYPE_1, + n, + n, + n, + ctypes.byref(lwork) + ) + lwork = lwork.value + _buffersize[h.dtype, n] = lwork + + if h.dtype == np.float64: + fn = libonemkl.onemkl_dsygvd + else: + fn = libonemkl.onemkl_zhegvd + #Allocate work-space + work_buf = dpnp.empty((lwork,), dtype=h.dtype) + fn(CUSOLVER_EIG_TYPE_1, + n, + ctypes.cast(A.data.ptr, ctypes.c_void_p), + n, + ctypes.cast(B.data.ptr, ctypes.c_void_p), + n, + ctypes.cast(w.data.ptr, ctypes.c_void_p), + ctypes.cast(work_buf.data.ptr, ctypes.c_void_p), + lwork) + return w, A.T + +# The oneMKL potrf path below is retained but unused: all production callers +# (df.py, df_jk.py, grad/rhf.py, via cupy_helper) resolve cholesky to +# dpnp_helper.cholesky, which wraps dpnp.linalg.cholesky directly. Only +# lib/tests/test_cusolver.py imports cholesky from this module, and that test +# is satisfied by the dpnp delegation below. Kept commented rather than +# deleted in case oneMKL potrf is ever wanted for performance; the C wrappers +# onemkl_dpotrf/onemkl_zpotrf still exist in onemkl_helper/onemkl_lapack.cpp. +# +# def cholesky(A): +# """ +# Compute the Cholesky decomposition of a Hermitian positive-definite matrix. +# +# Args: +# A: Hermitian positive-definite matrix +# +# Returns: +# Lower triangular matrix L such that A = L * L.T +# """ +# n = len(A) +# # cusolver.py transposes an F-contiguous input and copies to C order; +# # do the same rather than asserting, so callers passing a transposed +# # view behave identically on both backends. +# if A.flags.f_contiguous: +# A = A.T +# x = A.copy(order='C') +# if A.dtype == np.float64: +# potrf = libonemkl.onemkl_dpotrf +# potrf_bufferSize = libonemkl.onemkl_dpotrf_scratchpad_size +# else: +# potrf = libonemkl.onemkl_zpotrf +# potrf_bufferSize = libonemkl.onemkl_zpotrf_scratchpad_size +# scratchpad_size = potrf_bufferSize(n, n) +# scratchpad = dpnp.empty(scratchpad_size, dtype=A.dtype) +# info = potrf(n, +# ctypes.cast(x.data.ptr, ctypes.c_void_p), +# n, +# ctypes.cast(scratchpad.data.ptr, ctypes.c_void_p), +# scratchpad_size) +# # cusolver.py raises on a non-zero dev_info; df.py and grad/rhf.py catch +# # RuntimeError to fall back to an eigendecomposition for singular j2c. +# if info != 0: +# raise LinAlgError('failed to perform Cholesky Decomposition') +# x = dpnp.tril(x, k=0) +# return x + +def cholesky(A): + """ + Cholesky decomposition of a Hermitian positive-definite matrix. + + Delegates to dpnp so the SYCL backend has a single implementation, and + re-raises as this module's LinAlgError (a RuntimeError, matching + cusolver.LinAlgError) because df.py, df_jk.py and grad/rhf.py catch + RuntimeError to fall back to an eigendecomposition on a singular j2c. + dpnp.linalg.LinAlgError derives from ValueError and would slip past them. + + Args: + A: Hermitian positive-definite matrix + + Returns: + Lower triangular matrix L such that A = L @ L.T + """ + try: + return dpnp.linalg.cholesky(A) + except dpnp.linalg.LinAlgError as e: + raise LinAlgError(str(e)) from e + +class LinAlgError(RuntimeError): + """Mirrors cusolver.LinAlgError, which also derives from RuntimeError.""" + pass diff --git a/gpu4pyscf/lib/pbc/CMakeLists.txt b/gpu4pyscf/lib/pbc/CMakeLists.txt index e02db8a24..6c2376caf 100644 --- a/gpu4pyscf/lib/pbc/CMakeLists.txt +++ b/gpu4pyscf/lib/pbc/CMakeLists.txt @@ -1,6 +1,4 @@ -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128") - -add_library(pbc SHARED +set(GPU_SRCS pbc_driver.cu ft_ao.cu ft_ao_ip1.cu contract_ft_ao.cu fill_int3c2e.cu contract_int3c2e.cu ejk_int3c2e_ip1.cu @@ -12,12 +10,25 @@ add_library(pbc SHARED rys_contract_j.cu unrolled_rys_k.cu nr_eval_gto.cu - sorting.c fill_triu.cu + sorting.c ) +add_library(pbc SHARED ${GPU_SRCS}) + set_target_properties(pbc PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_SEPARABLE_COMPILATION ON) + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + +if (USE_SYCL) + # icx rejects the project-wide -std=gnu99 (CMAKE_C_STANDARD 99) together + # with -fsycl, so plain-C sources are compiled as C++ here, matching + # gvhf-md/md_pairdata.c. sorting.c carries an extern "C" guard. + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + target_compile_options(pbc PRIVATE -x c++ -nocudainc -nocudalib) + target_link_libraries(pbc PRIVATE sycl_compat gsycl) +else() + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v") + set_target_properties(pbc PROPERTIES CUDA_SEPARABLE_COMPILATION ON) +endif() -target_link_libraries(pbc OpenMP::OpenMP_C) +target_link_libraries(pbc PRIVATE OpenMP::OpenMP_C) diff --git a/gpu4pyscf/lib/pbc/contract_ft_ao.cu b/gpu4pyscf/lib/pbc/contract_ft_ao.cu index 125ba16c0..7c1816ac8 100644 --- a/gpu4pyscf/lib/pbc/contract_ft_ao.cu +++ b/gpu4pyscf/lib/pbc/contract_ft_ao.cu @@ -30,21 +30,67 @@ #define OF_COMPLEX 2 #define POOL_SIZE 65536 + +#ifdef USE_SYCL +#define KERNEL_SETUP() \ + int thread_id = item.get_local_id(0); \ + auto thread_block = item.get_group(); \ + int &sp_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &img_max = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int (&img_counts)[sp_threads] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + double *shared_memory = reinterpret_cast(shm_mem); + +#define KERNEL_ARGS() \ + double *out, double *vG, \ + PBCIntEnvVars envs, int *shl_pair_offsets, \ + uint32_t *bas_ij_idx, int *img_idx, uint32_t *img_offsets, \ + int *gout_stride_lookup, double *Gv, int nGv, \ + int nbatches_shl_pair, int compressing, int *head, \ + sycl::nd_item<1> &item, std::byte *shm_mem + +#else // USE_SYCL + +#define KERNEL_SETUP() \ + int thread_id = threadIdx.x; \ + __shared__ int sp_block_id; \ + __shared__ int shl_pair0, shl_pair1; \ + __shared__ int li, lj; \ + __shared__ int iprim, jprim; \ + __shared__ int nao; \ + __shared__ int gout_stride, nsp_per_block; \ + __shared__ int img_max; \ + __shared__ int img_counts[sp_threads]; \ + extern __shared__ double shared_memory[]; + +#define KERNEL_ARGS() \ + double *out, double *vG, \ + PBCIntEnvVars envs, int *shl_pair_offsets, \ + uint32_t *bas_ij_idx, int *img_idx, uint32_t *img_offsets, \ + int *gout_stride_lookup, double *Gv, int nGv, \ + int nbatches_shl_pair, int compressing, int *head + +#endif // USE_SYCL + + __global__ static -void ft_aopair_kernel(double *out, double *vG, - PBCIntEnvVars envs, int *shl_pair_offsets, - uint32_t *bas_ij_idx, int *img_idx, uint32_t *img_offsets, - int *gout_stride_lookup, double *Gv, int nGv, - int nbatches_shl_pair, int compressing, int *head) +void ft_aopair_kernel(KERNEL_ARGS()) { constexpr int nGv_per_block = 16; - constexpr unsigned mask = (1u << nGv_per_block) - 1; constexpr int sp_threads = THREADS / nGv_per_block; + constexpr unsigned mask = (1u << nGv_per_block) - 1; constexpr unsigned sp_mask = (1u << sp_threads) - 1; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int Gv_id_in_block = thread_id % nGv_per_block; int t_id = thread_id / nGv_per_block; - __shared__ int sp_block_id; while (1) { if (thread_id == 0) { sp_block_id = atomicAdd(head, 1); @@ -59,11 +105,6 @@ while (1) { int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj; - __shared__ int iprim, jprim; - __shared__ int nao; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -91,7 +132,6 @@ while (1) { int stride_j = li + 1; int g_size = stride_j * (lj + 1); int gx_len = g_size * nGsp_per_block; - extern __shared__ double shared_memory[]; double *gxR = shared_memory + nGv_per_block * sp_id + Gv_id_in_block; double *gxI = gxR + gx_len; double *gyR = gxR + gx_len*2; @@ -113,12 +153,22 @@ while (1) { int jsh = bas_ij % bvk_nbas; int img0 = img_offsets[pair_ij]; int img1 = img_offsets[pair_ij+1]; - __shared__ int img_max; - __shared__ int img_counts[sp_threads]; if (Gv_id_in_block == 0) { img_counts[t_id] = img1 - img0; } __syncthreads(); +#ifdef USE_SYCL + // See ft_ao.cu for why a sub-group shuffle restricted to a lane + // subset (thread_id < sp_threads) is UB in SYCL when sp_threads + // doesn't match the HW sub-group width -- scan serially instead. + if (thread_id == 0) { + int count = img_counts[0]; + for (int w = 1; w < sp_threads; ++w) { + count = max(count, img_counts[w]); + } + img_max = count; + } +#else if (thread_id < sp_threads) { int count = img_counts[thread_id]; for (int offset = sp_threads/2; offset > 0; offset /= 2) { @@ -128,6 +178,7 @@ while (1) { img_max = count; } } +#endif __syncthreads(); int expi = bas[ish*BAS_SLOTS+PTR_EXP]; @@ -319,20 +370,15 @@ while (1) { } __global__ static -void ft_pdotp_kernel(double *out, double *vG, - PBCIntEnvVars envs, int *shl_pair_offsets, - uint32_t *bas_ij_idx, int *img_idx, uint32_t *img_offsets, - int *gout_stride_lookup, double *Gv, int nGv, - int nbatches_shl_pair, int compressing, int *head) +void ft_pdotp_kernel(KERNEL_ARGS()) { constexpr int nGv_per_block = 16; - constexpr unsigned mask = (1u << nGv_per_block) - 1; constexpr int sp_threads = THREADS / nGv_per_block; + constexpr unsigned mask = (1u << nGv_per_block) - 1; constexpr unsigned sp_mask = (1u << sp_threads) - 1; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int Gv_id_in_block = thread_id % nGv_per_block; int t_id = thread_id / nGv_per_block; - __shared__ int sp_block_id; while (1) { if (thread_id == 0) { sp_block_id = atomicAdd(head, 1); @@ -347,11 +393,6 @@ while (1) { int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj; - __shared__ int iprim, jprim; - __shared__ int nao; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -379,7 +420,6 @@ while (1) { int stride_j = li + 2; int g_size = stride_j * (lj + 2); int gx_len = g_size * nGsp_per_block; - extern __shared__ double shared_memory[]; double *gxR = shared_memory + nGv_per_block * sp_id + Gv_id_in_block; double *gxI = gxR + gx_len; double *gyR = gxR + gx_len*2; @@ -401,12 +441,22 @@ while (1) { int jsh = bas_ij % bvk_nbas; int img0 = img_offsets[pair_ij]; int img1 = img_offsets[pair_ij+1]; - __shared__ int img_max; - __shared__ int img_counts[sp_threads]; if (Gv_id_in_block == 0) { img_counts[t_id] = img1 - img0; } __syncthreads(); +#ifdef USE_SYCL + // See ft_ao.cu for why a sub-group shuffle restricted to a lane + // subset (thread_id < sp_threads) is UB in SYCL when sp_threads + // doesn't match the HW sub-group width -- scan serially instead. + if (thread_id == 0) { + int count = img_counts[0]; + for (int w = 1; w < sp_threads; ++w) { + count = max(count, img_counts[w]); + } + img_max = count; + } +#else if (thread_id < sp_threads) { int count = img_counts[thread_id]; for (int offset = sp_threads/2; offset > 0; offset /= 2) { @@ -416,6 +466,7 @@ while (1) { img_max = count; } } +#endif __syncthreads(); int expi = bas[ish*BAS_SLOTS+PTR_EXP]; @@ -675,11 +726,23 @@ int contract_ft_aopair(double *out, double *vG, PBCIntEnvVars *envs, int *head, int *gout_stride_lookup, double *grids, int ngrids, int compressing) { - cudaFuncSetAttribute(ft_aopair_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + cudaMemset(head, 0, sizeof(int)); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); +#ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(workers * THREADS, THREADS), [=](auto item) { + ft_aopair_kernel( + out, vG, dev_envs, shl_pair_offsets, bas_ij_idx, img_idx, img_offsets, + gout_stride_lookup, grids, ngrids, nbatches_shl_pair, compressing, head, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else + cudaFuncSetAttribute(ft_aopair_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); ft_aopair_kernel<<>>( out, vG, *envs, shl_pair_offsets, bas_ij_idx, img_idx, img_offsets, gout_stride_lookup, grids, ngrids, nbatches_shl_pair, compressing, head); @@ -688,6 +751,7 @@ int contract_ft_aopair(double *out, double *vG, PBCIntEnvVars *envs, int *head, fprintf(stderr, "CUDA Error in ft_aopair_kernel: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -697,11 +761,23 @@ int contract_ft_pdotp(double *out, double *vG, PBCIntEnvVars *envs, int *head, int *gout_stride_lookup, double *grids, int ngrids, int compressing) { - cudaFuncSetAttribute(ft_pdotp_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + cudaMemset(head, 0, sizeof(int)); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); +#ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(workers * THREADS, THREADS), [=](auto item) { + ft_pdotp_kernel( + out, vG, dev_envs, shl_pair_offsets, bas_ij_idx, img_idx, img_offsets, + gout_stride_lookup, grids, ngrids, nbatches_shl_pair, compressing, head, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else + cudaFuncSetAttribute(ft_pdotp_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); ft_pdotp_kernel<<>>( out, vG, *envs, shl_pair_offsets, bas_ij_idx, img_idx, img_offsets, gout_stride_lookup, grids, ngrids, nbatches_shl_pair, compressing, head); @@ -710,6 +786,10 @@ int contract_ft_pdotp(double *out, double *vG, PBCIntEnvVars *envs, int *head, fprintf(stderr, "CUDA Error in ft_pdotp_kernel: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } + +#undef KERNEL_SETUP +#undef KERNEL_ARGS diff --git a/gpu4pyscf/lib/pbc/contract_int3c2e.cu b/gpu4pyscf/lib/pbc/contract_int3c2e.cu index fb6db8ee4..fd14814b0 100644 --- a/gpu4pyscf/lib/pbc/contract_int3c2e.cu +++ b/gpu4pyscf/lib/pbc/contract_int3c2e.cu @@ -25,7 +25,11 @@ #include "int3c2e_create_tasks.cuh" #define NF_AUX_MAX 28 +#ifdef USE_SYCL +#define GOUT_WIDTH 30 +#else #define GOUT_WIDTH 29 +#endif // lattice sum over j and k for (ij|k) __global__ static @@ -35,15 +39,44 @@ void contract_int3c2e_dm_kernel(double *out, double *dm, double omega, PBCIntEnv int *img_idx, uint32_t *sp_img_offsets, int *gout_stride_lookup, int nauxbas, float *diffuse_exps, float *diffuse_coefs, float log_cutoff, - int *head, int sp_blocks) + int *head, int sp_blocks + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int thread_id = threadIdx.x; - img_pool += blockIdx.x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); - // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 - uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; - uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); - ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx.x * POOL_SIZE; - extern __shared__ double shared_memory[]; + #ifdef USE_SYCL + int blockIdx_x = item.get_group(0); + int threadIdx_x = item.get_local_id(0); + + auto thread_block = item.get_group(); + int &sp_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ck = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + int &num_ijk_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &num_sub_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_not_processed = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_tile_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int blockIdx_x = blockIdx.x; + int threadIdx_x = threadIdx.x; + __shared__ int sp_block_id, ksh_cell0; __shared__ int shl_pair0, shl_pair1; __shared__ int li, lj, lk, nroots; @@ -51,8 +84,19 @@ void contract_int3c2e_dm_kernel(double *out, double *dm, double omega, PBCIntEnv __shared__ int nao; __shared__ int gout_stride, nst_per_block; __shared__ int expk, ck; + + extern __shared__ double shared_memory[]; + __shared__ int num_ijk_tasks; __shared__ int num_sub_tasks, img_not_processed, img_tile_size; + #endif + + int thread_id = threadIdx_x; + img_pool += blockIdx_x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); + // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 + uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; + uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); + ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx_x * POOL_SIZE; while (1) { if (thread_id == 0) { int batch_id = atomicAdd(head, 1); @@ -292,15 +336,54 @@ void contract_int3c2e_auxvec_kernel(double *out, double *auxvec, double omega, P int *img_idx, uint32_t *sp_img_offsets, int *gout_stride_lookup, int nauxbas, float *diffuse_exps, float *diffuse_coefs, float log_cutoff, - int *head, int npairs_ij, int ksh_blocks) + int *head, int npairs_ij, int ksh_blocks + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int thread_id = threadIdx.x; - img_pool += blockIdx.x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); - // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 - uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; - uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); - ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx.x * POOL_SIZE; - extern __shared__ double shared_memory[]; + #ifdef USE_SYCL + int blockIdx_x = item.get_group(0); + int threadIdx_x = item.get_local_id(0); + + auto thread_block = item.get_group(); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ci = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &cj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &xi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &yi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &zi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &xjxi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &yjyi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &zjzi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &fac = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + int &num_ijk_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &num_sub_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_not_processed = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_tile_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int blockIdx_x = blockIdx.x; + int threadIdx_x = threadIdx.x; + __shared__ int pair_ij, ksh_block_id; __shared__ int ksh0_cell0, ksh1_cell0; __shared__ int ish, jsh, li, lj, lk, nroots; @@ -308,8 +391,20 @@ void contract_int3c2e_auxvec_kernel(double *out, double *auxvec, double omega, P __shared__ int gout_stride, nst_per_block; __shared__ int expi, expj, ci, cj; __shared__ double xi, yi, zi, xjxi, yjyi, zjzi; + __shared__ double fac; + __shared__ int num_ijk_tasks; __shared__ int num_sub_tasks, img_not_processed, img_tile_size; + + extern __shared__ double shared_memory[]; + #endif + + int thread_id = threadIdx_x; + img_pool += blockIdx_x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); + // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 + uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; + uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); + ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx_x * POOL_SIZE; while (1) { if (thread_id == 0) { int batch_id = atomicAdd(head, 1); @@ -355,15 +450,16 @@ while (1) { xjxi = rj[0] - xi; yjyi = rj[1] - yi; zjzi = rj[2] - zi; + int ish_cell0 = ish; + int jsh_cell0 = jsh % envs.nbas; + fac = PI_FAC; + if (ish_cell0 < jsh_cell0) { + fac = 0; + } gout_stride = gout_stride_lookup[lk*LMAX1*LMAX1+li*LMAX1+lj]; nst_per_block = THREADS / gout_stride; } __syncthreads(); - int ish_cell0 = ish; - int jsh_cell0 = jsh % envs.nbas; - if (ish_cell0 < jsh_cell0) { - continue; - } int gout_id = thread_id / nst_per_block; int st_id = thread_id - gout_id * nst_per_block; @@ -461,7 +557,7 @@ while (ksh0_cell0 < ksh1_cell0) { double rr_ij = xjLxi * xjLxi + yjLyi * yjLyi + zjLzi * zjLzi; double theta_ij = ai * aj_aij; double Kab = theta_ij * rr_ij; - double cicj = PI_FAC * env[ci+ip] * env[cj+jp]; + double cicj = fac * env[ci+ip] * env[cj+jp]; fac_ij = exp(-Kab) * cicj; } double xij = xjLxi * aj_aij + xi; @@ -552,11 +648,25 @@ int PBCcontract_int3c2e_dm(double *out, double *dm, double omega, PBCIntEnvVars int *img_idx, uint32_t *img_offsets, int *gout_stride_lookup, float *diffuse_exps, float *diffuse_coefs, float log_cutoff) { - cudaFuncSetAttribute(contract_int3c2e_dm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + cudaMemset(head, 0, sizeof(int)); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); +#ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(workers * THREADS, THREADS), [=](auto item) { + contract_int3c2e_dm_kernel( + out, dm, omega, dev_envs, pool, task_pool, bas_ij_idx, shl_pair_offsets, + img_idx, img_offsets, gout_stride_lookup, nauxbas, + diffuse_exps, diffuse_coefs, log_cutoff, + head, nbatches_shl_pair, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else + cudaFuncSetAttribute(contract_int3c2e_dm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); contract_int3c2e_dm_kernel<<>>( out, dm, omega, *envs, pool, task_pool, bas_ij_idx, shl_pair_offsets, img_idx, img_offsets, gout_stride_lookup, nauxbas, @@ -567,6 +677,7 @@ int PBCcontract_int3c2e_dm(double *out, double *dm, double omega, PBCIntEnvVars fprintf(stderr, "CUDA Error in contract_int3c2e_dm: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } @@ -577,11 +688,24 @@ int PBCcontract_int3c2e_auxvec(double *out, double *auxvec, double omega, PBCInt int *img_idx, uint32_t *img_offsets, int *gout_stride_lookup, float *diffuse_exps, float *diffuse_coefs, float log_cutoff) { - cudaFuncSetAttribute(contract_int3c2e_auxvec_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + cudaMemset(head, 0, sizeof(int)); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); +#ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(workers * THREADS, THREADS), [=](auto item) { + contract_int3c2e_auxvec_kernel(out, auxvec, omega, dev_envs, pool, task_pool, bas_ij_idx, ksh_offsets, + img_idx, img_offsets, gout_stride_lookup, nauxbas, + diffuse_exps, diffuse_coefs, log_cutoff, + head, npairs, nbatches_ksh, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else + cudaFuncSetAttribute(contract_int3c2e_auxvec_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); contract_int3c2e_auxvec_kernel<<>>( out, auxvec, omega, *envs, pool, task_pool, bas_ij_idx, ksh_offsets, img_idx, img_offsets, gout_stride_lookup, nauxbas, @@ -592,6 +716,7 @@ int PBCcontract_int3c2e_auxvec(double *out, double *auxvec, double omega, PBCInt fprintf(stderr, "CUDA Error in contract_int3c2e_auxvec: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/contract_int3c2e_pvp.cu b/gpu4pyscf/lib/pbc/contract_int3c2e_pvp.cu index 0d8e22015..c48ead583 100644 --- a/gpu4pyscf/lib/pbc/contract_int3c2e_pvp.cu +++ b/gpu4pyscf/lib/pbc/contract_int3c2e_pvp.cu @@ -33,15 +33,52 @@ void contract_int3c2e_pvp_auxvec_kernel(double *out, double *auxvec, double omeg int *img_idx, uint32_t *sp_img_offsets, int *gout_stride_lookup, int nauxbas, float *diffuse_exps, float *diffuse_coefs, float log_cutoff, - int *head, int npairs_ij, int ksh_blocks) + int *head, int npairs_ij, int ksh_blocks + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int thread_id = threadIdx.x; - img_pool += blockIdx.x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); - // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 - uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; - uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); - ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx.x * POOL_SIZE; - extern __shared__ double shared_memory[]; + #ifdef USE_SYCL + int blockIdx_x = item.get_group(0); + int threadIdx_x = item.get_local_id(0); + + auto thread_block = item.get_group(); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ci = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &cj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &xi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &yi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &zi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &xjxi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &yjyi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &zjzi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &num_ijk_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &num_sub_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_not_processed = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_tile_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int blockIdx_x = blockIdx.x; + int threadIdx_x = threadIdx.x; + __shared__ int pair_ij, ksh_block_id; __shared__ int ksh0_cell0, ksh1_cell0; __shared__ int ish, jsh, li, lj, lk, nroots; @@ -51,6 +88,16 @@ void contract_int3c2e_pvp_auxvec_kernel(double *out, double *auxvec, double omeg __shared__ double xi, yi, zi, xjxi, yjyi, zjzi; __shared__ int num_ijk_tasks; __shared__ int num_sub_tasks, img_not_processed, img_tile_size; + + extern __shared__ double shared_memory[]; + #endif + + int thread_id = threadIdx_x; + img_pool += blockIdx_x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); + // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 + uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; + uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); + ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx_x * POOL_SIZE; while (1) { if (thread_id == 0) { int batch_id = atomicAdd(head, 1); @@ -322,11 +369,25 @@ int PBCcontract_int3c2e_pvp_auxvec(double *out, double *auxvec, double omega, PB int *img_idx, uint32_t *img_offsets, int *gout_stride_lookup, float *diffuse_exps, float *diffuse_coefs, float log_cutoff) { - cudaFuncSetAttribute(contract_int3c2e_pvp_auxvec_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + cudaMemset(head, 0, sizeof(int)); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); +#ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(workers * THREADS, THREADS), [=](auto item) { + contract_int3c2e_pvp_auxvec_kernel( + out, auxvec, omega, dev_envs, pool, task_pool, bas_ij_idx, ksh_offsets, + img_idx, img_offsets, gout_stride_lookup, nauxbas, + diffuse_exps, diffuse_coefs, log_cutoff, + head, npairs, nbatches_ksh, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); +#else + cudaFuncSetAttribute(contract_int3c2e_pvp_auxvec_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); contract_int3c2e_pvp_auxvec_kernel<<>>( out, auxvec, omega, *envs, pool, task_pool, bas_ij_idx, ksh_offsets, img_idx, img_offsets, gout_stride_lookup, nauxbas, @@ -337,6 +398,7 @@ int PBCcontract_int3c2e_pvp_auxvec(double *out, double *auxvec, double omega, PB fprintf(stderr, "CUDA Error in contract_int3c2e_pvp_auxvec: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/create_tasks.cu b/gpu4pyscf/lib/pbc/create_tasks.cu index b3c8ce632..706747a0f 100644 --- a/gpu4pyscf/lib/pbc/create_tasks.cu +++ b/gpu4pyscf/lib/pbc/create_tasks.cu @@ -29,6 +29,9 @@ __device__ inline int mask_to_index(int keep, int *tmp_storage, int threads, int t_id) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); +#endif tmp_storage[t_id] = keep; __syncthreads(); for (int offset = 1; offset < threads; offset <<= 1) { @@ -54,8 +57,16 @@ void _fill_sr_vk_tasks(int &ntasks, int &pair_kl0, int64_t *bas_kl_idx, int *swap, JKMatrix& kmat, RysIntEnvVars& envs, BoundsInfo& bounds) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int thread_id = item.get_local_id(1) + item.get_local_range(1) * item.get_local_id(0); + int threads = item.get_local_range(1) * item.get_local_range(0); + int threadIdx_y = item.get_local_id(0); +#else int thread_id = threadIdx.x + blockDim.x * threadIdx.y; int threads = blockDim.x * blockDim.y; + int threadIdx_y = threadIdx.y; +#endif __syncthreads(); if (thread_id == 0) { ntasks = 0; @@ -182,7 +193,7 @@ void _fill_sr_vk_tasks(int &ntasks, int &pair_kl0, int64_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + thread_id < QUEUE_DEPTH && ntasks > 0) { + if (threadIdx_y == 0 && ntasks + thread_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+thread_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); @@ -198,8 +209,16 @@ void _fill_sr_ejk_tasks(int &ntasks, int &pair_kl0, int64_t *bas_kl_idx, int *swap, JKEnergy& jk, RysIntEnvVars& envs, BoundsInfo& bounds) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int thread_id = item.get_local_id(1) + item.get_local_range(1) * item.get_local_id(0); + int threads = item.get_local_range(1) * item.get_local_range(0); + int threadIdx_y = item.get_local_id(0); +#else int thread_id = threadIdx.x + blockDim.x * threadIdx.y; int threads = blockDim.x * blockDim.y; + int threadIdx_y = threadIdx.y; +#endif __syncthreads(); if (thread_id == 0) { ntasks = 0; @@ -330,7 +349,7 @@ void _fill_sr_ejk_tasks(int &ntasks, int &pair_kl0, int64_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + thread_id < QUEUE_DEPTH && ntasks > 0) { + if (threadIdx_y == 0 && ntasks + thread_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+thread_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); diff --git a/gpu4pyscf/lib/pbc/ejk_int3c2e_ip1.cu b/gpu4pyscf/lib/pbc/ejk_int3c2e_ip1.cu index 0caa436d7..76d9398b1 100644 --- a/gpu4pyscf/lib/pbc/ejk_int3c2e_ip1.cu +++ b/gpu4pyscf/lib/pbc/ejk_int3c2e_ip1.cu @@ -34,14 +34,64 @@ void ejk_int3c2e_ip1_kernel(double *ejk, double *ejk_aux, double *dm, double *de int *img_idx, uint32_t *sp_img_offsets, int *gout_stride_lookup, int *ao_pair_loc, int aux_offset, int nauxbas, int naux, float *diffuse_exps, float *diffuse_coefs, float log_cutoff, - int *head, int sp_blocks, int ksh_blocks) + int *head, int sp_blocks, int ksh_blocks + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int thread_id = threadIdx.x; +#ifdef USE_SYCL + int threadIdx_x = item.get_local_id(0); + int blockIdx_x = item.get_group(0); + + auto thread_block = item.get_group(); + int &sp_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + + int &num_ijk_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &num_sub_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_not_processed = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_tile_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + __shared__ int sp_block_id, ksh_block_id; - uint32_t *img_pool = pool + blockIdx.x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); + __shared__ int ksh0_cell0, ksh1_cell0; + __shared__ int shl_pair0, shl_pair1; + __shared__ int li, lj, lk, nroots, nf; + __shared__ int iprim, jprim, kprim; + __shared__ int nao; + __shared__ int g_size, gout_stride, nst_per_block; + + extern __shared__ double shared_memory[]; + __shared__ int num_ijk_tasks; + __shared__ int num_sub_tasks, img_not_processed, img_tile_size; +#endif + + int thread_id = threadIdx_x; + uint32_t *img_pool = pool + blockIdx_x * POOL_SIZE * (MAX_IMGS_PER_TASK+2); uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); - ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx.x * POOL_SIZE; + ShellTripletTaskInfo *ijk_tasks_info = task_pool + blockIdx_x * POOL_SIZE; while (1) { if (thread_id == 0) { int batch_id = atomicAdd(head, 1); @@ -58,12 +108,6 @@ while (1) { double *env = envs.env; double *img_coords = envs.img_coords; int nimgs = envs.nimgs; - __shared__ int ksh0_cell0, ksh1_cell0; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, lk, nroots, nf; - __shared__ int iprim, jprim, kprim; - __shared__ int nao; - __shared__ int g_size, gout_stride, nst_per_block; if (thread_id == 0) { int bvk_nbas = envs.nbas * ncells; shl_pair0 = shl_pair_offsets[sp_block_id]; @@ -99,7 +143,6 @@ while (1) { int gout_id = thread_id / nst_per_block; int st_id = thread_id - gout_id * nst_per_block; - extern __shared__ double shared_memory[]; double *rjri = shared_memory + st_id; double *Rpq = shared_memory + nst_per_block * 3 + st_id; double *gx = shared_memory + nst_per_block * 6 + st_id; @@ -108,7 +151,6 @@ while (1) { int idx_j = lex_xyz_offset(lj); int idx_k = lex_xyz_offset(lk); - __shared__ int num_ijk_tasks; if (thread_id == 0) { int nshl_pairs = shl_pair1 - shl_pair0; int nksh = ksh1_cell0 - ksh0_cell0; @@ -122,7 +164,6 @@ while (1) { while (num_ijk_tasks > 0) { _filter_jk_images(img_pool, rem_task_idx, num_ijk_tasks, ijk_tasks_info, envs, img_idx); - __shared__ int num_sub_tasks, img_not_processed, img_tile_size; if (thread_id == 0) { img_tile_size = 8; img_not_processed = MAX_IMGS_PER_TASK; @@ -391,11 +432,27 @@ int PBCsr_ejk_int3c2e_ip1(double *ejk, double*ejk_aux, double *dm, double *densi int *ao_pair_loc, int aux_offset, int nauxbas, int naux, float *diffuse_exps, float *diffuse_coefs, float log_cutoff) { - cudaFuncSetAttribute(ejk_int3c2e_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + cudaMemset(head, 0, sizeof(int)); + cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(workers * THREADS, THREADS), [=](auto item) { + ejk_int3c2e_ip1_kernel( + ejk, ejk_aux, dm, density_auxvec, omega, dev_envs, pool, task_pool, + bas_ij_idx, shl_pair_offsets, ksh_offsets, img_idx, img_offsets, + gout_stride_lookup, ao_pair_loc, aux_offset, nauxbas, naux, + diffuse_exps, diffuse_coefs, log_cutoff, + head, nbatches_shl_pair, nbatches_ksh, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + cudaFuncSetAttribute(ejk_int3c2e_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); ejk_int3c2e_ip1_kernel<<>>( ejk, ejk_aux, dm, density_auxvec, omega, *envs, pool, task_pool, bas_ij_idx, shl_pair_offsets, ksh_offsets, img_idx, img_offsets, @@ -407,6 +464,7 @@ int PBCsr_ejk_int3c2e_ip1(double *ejk, double*ejk_aux, double *dm, double *densi fprintf(stderr, "CUDA Error in ejk_int3c2e_ip1: %s\n", cudaGetErrorString(err)); return 1; } +#endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/fill_int2c2e.cu b/gpu4pyscf/lib/pbc/fill_int2c2e.cu index 1ef682f1c..dcf4e2ced 100644 --- a/gpu4pyscf/lib/pbc/fill_int2c2e.cu +++ b/gpu4pyscf/lib/pbc/fill_int2c2e.cu @@ -33,8 +33,30 @@ __global__ static void pbc_int2c2e_kernel(double *out, PBCIntEnvVars envs, double omega, double lr_factor, double sr_factor, int *shl_pair_offsets, uint32_t *bas_ij_idx, - int *gout_stride_lookup) + int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<1> &item, char *shm_mem + #endif + ) { + #ifdef USE_SYCL + int sp_block_id = item.get_group(0); + int thread_id = item.get_local_id(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nbas = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else int sp_block_id = blockIdx.x; int thread_id = threadIdx.x; extern __shared__ double shared_memory[]; @@ -42,6 +64,7 @@ void pbc_int2c2e_kernel(double *out, PBCIntEnvVars envs, __shared__ int nbas; __shared__ int li, lj, nroots, nao, iprim, jprim; __shared__ int gout_stride; + #endif int *bas = envs.bas; double *env = envs.env; @@ -254,6 +277,19 @@ int fill_int2c2e(double *out, PBCIntEnvVars *envs, int nbatches_shl_pair, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *gout_stride_lookup) { + #ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + pbc_int2c2e_kernel(out, dev_envs, omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, gout_stride_lookup, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(pbc_int2c2e_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); pbc_int2c2e_kernel<<>>( out, *envs, omega, lr_factor, sr_factor, @@ -263,6 +299,7 @@ int fill_int2c2e(double *out, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in int2c2e kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/fill_int2c2e_ip1.cu b/gpu4pyscf/lib/pbc/fill_int2c2e_ip1.cu index d0941dc02..a7abb3a6a 100644 --- a/gpu4pyscf/lib/pbc/fill_int2c2e_ip1.cu +++ b/gpu4pyscf/lib/pbc/fill_int2c2e_ip1.cu @@ -33,17 +33,44 @@ __global__ static void pbc_int2c2e_ip1_kernel(double *out, PBCIntEnvVars envs, double omega, double lr_factor, double sr_factor, int *shl_pair_offsets, uint32_t *bas_ij_idx, - int *gout_stride_lookup) + int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<1> &item, char *shm_mem + #endif + ) { + #ifdef USE_SYCL + int sp_block_id = item.get_group(0); + int thread_id = item.get_local_id(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nbas = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else int sp_block_id = blockIdx.x; int thread_id = threadIdx.x; - int *bas = envs.bas; - double *env = envs.env; - double *img_coords = envs.img_coords; + __shared__ int shl_pair0, shl_pair1; __shared__ int nbas; __shared__ int li, lj, nroots, nao, iprim, jprim; __shared__ int gout_stride; + + extern __shared__ double shared_memory[]; + #endif + + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -73,12 +100,11 @@ void pbc_int2c2e_ip1_kernel(double *out, PBCIntEnvVars envs, int stride_j = li + 2; int g_size = stride_j * (lj + 1); int gx_len = g_size * nsp_per_block; - extern __shared__ double shared_memory[]; double *rw = shared_memory + sp_id; double *gx = shared_memory + nsp_per_block * nroots*2 + sp_id; double *Rpq = shared_memory + nsp_per_block * (g_size*3+nroots*2) + sp_id; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); double goutx[GOUT_IP_WIDTH]; double gouty[GOUT_IP_WIDTH]; double goutz[GOUT_IP_WIDTH]; @@ -270,17 +296,44 @@ __global__ static void e_int2c2e_ip1_kernel(double *out, double *dm, PBCIntEnvVars envs, double omega, double lr_factor, double sr_factor, int *shl_pair_offsets, uint32_t *bas_ij_idx, - int *gout_stride_lookup) + int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<1> &item, char *shm_mem + #endif + ) { + #ifdef USE_SYCL + int sp_block_id = item.get_group(0); + int thread_id = item.get_local_id(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nbas = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else int sp_block_id = blockIdx.x; int thread_id = threadIdx.x; - int *bas = envs.bas; - double *env = envs.env; - double *img_coords = envs.img_coords; + __shared__ int shl_pair0, shl_pair1; __shared__ int nbas; __shared__ int li, lj, nroots, nao, iprim, jprim; __shared__ int gout_stride; + + extern __shared__ double shared_memory[]; + #endif + + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -311,12 +364,11 @@ void e_int2c2e_ip1_kernel(double *out, double *dm, PBCIntEnvVars envs, int j_1 = stride_j*nsp_per_block; int g_size = stride_j * (lj + 2); int gx_len = g_size * nsp_per_block; - extern __shared__ double shared_memory[]; double *rw = shared_memory + sp_id; double *gx = shared_memory + nsp_per_block * nroots*2 + sp_id; double *Rpq = shared_memory + nsp_per_block * (g_size*3+nroots*2) + sp_id; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { double v_ix = 0; @@ -496,6 +548,19 @@ int fill_int2c2e_ip1(double *out, PBCIntEnvVars *envs, int nbatches_shl_pair, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *gout_stride_lookup) { + #ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + pbc_int2c2e_ip1_kernel(out, dev_envs, omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, gout_stride_lookup, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(pbc_int2c2e_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); pbc_int2c2e_ip1_kernel<<>>( out, *envs, omega, lr_factor, sr_factor, @@ -505,6 +570,7 @@ int fill_int2c2e_ip1(double *out, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in int2c2e_ip1 kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -513,6 +579,19 @@ int e_int2c2e_ip1(double *out, double *dm, PBCIntEnvVars *envs, int nbatches_shl_pair, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *gout_stride_lookup) { + #ifdef USE_SYCL + sycl::range<1> threads(THREADS); + sycl::range<1> blocks(nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + e_int2c2e_ip1_kernel(out, dm, dev_envs, omega, lr_factor, sr_factor, + shl_pair_offsets, bas_ij_idx, gout_stride_lookup, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(e_int2c2e_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); e_int2c2e_ip1_kernel<<>>( out, dm, *envs, omega, lr_factor, sr_factor, @@ -522,6 +601,7 @@ int e_int2c2e_ip1(double *out, double *dm, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in int2c2e_ip1 kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/fill_int3c2e.cu b/gpu4pyscf/lib/pbc/fill_int3c2e.cu index d971afe23..f9055bc88 100644 --- a/gpu4pyscf/lib/pbc/fill_int3c2e.cu +++ b/gpu4pyscf/lib/pbc/fill_int3c2e.cu @@ -39,25 +39,61 @@ void pbc_int3c2e_latsum23_kernel(double *out, double omega, PBCIntEnvVars envs, int ao_pair_offset, int aux_offset, int nauxbas, int naux, int to_sph, float *diffuse_exps, float *diffuse_coefs, float log_cutoff, - int *head, int nbatches_shl_pair, int nbatches_ksh) + int *head, int nbatches_shl_pair, int nbatches_ksh + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { + #ifdef USE_SYCL + int thread_id = item.get_local_id(0); + int worker_id = item.get_group(0); + + auto thread_block = item.get_group(); + int &sp_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh_block_id = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh0_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ksh1_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lk = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nroots = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nf = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &kprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &g_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nst_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &num_ijk_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &num_sub_tasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_not_processed = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_tile_size = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else int thread_id = threadIdx.x; int worker_id = blockIdx.x; - c2s_pool += worker_id * (THREADS*GOUT_WIDTH); - img_pool += worker_id * POOL_SIZE * (MAX_IMGS_PER_TASK+2); - // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 - uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; - uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); - ShellTripletTaskInfo *ijk_tasks_info = task_pool + worker_id * POOL_SIZE; - extern __shared__ double shared_memory[]; + + __shared__ int sp_block_id, ksh_block_id; __shared__ int ksh0_cell0, ksh1_cell0; __shared__ int shl_pair0, shl_pair1; __shared__ int li, lj, lk, nroots, nf; __shared__ int iprim, jprim, kprim; __shared__ int g_size, gout_stride, nst_per_block; + extern __shared__ double shared_memory[]; __shared__ int num_ijk_tasks; __shared__ int num_sub_tasks, img_not_processed, img_tile_size; - __shared__ int sp_block_id, ksh_block_id; + #endif + + c2s_pool += worker_id * (THREADS*GOUT_WIDTH); + img_pool += worker_id * POOL_SIZE * (MAX_IMGS_PER_TASK+2); + // rem_task_idx stores the Id of the ijk tasks which has remaining_imgs > 0 + uint32_t *rem_task_idx = img_pool + POOL_SIZE * MAX_IMGS_PER_TASK; + uint32_t *sub_task_idx = img_pool + POOL_SIZE *(MAX_IMGS_PER_TASK+1); + ShellTripletTaskInfo *ijk_tasks_info = task_pool + worker_id * POOL_SIZE; while (1) { if (thread_id == 0) { int batch_id = atomicAdd(head, 1); @@ -967,7 +1003,13 @@ void ovlp_img_counts_kernel(int *img_counts, PBCIntEnvVars envs, float *exps, float *log_coef, float log_cutoff, int permutation_symmetry) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int bas_ij = item.get_global_id(0); + #else int bas_ij = blockIdx.x * blockDim.x + threadIdx.x; + #endif + int bvk_nbas = envs.bvk_ncells * envs.nbas; int ish = bas_ij / bvk_nbas; int jsh = bas_ij - bvk_nbas * ish; @@ -1033,7 +1075,13 @@ __global__ static void ovlp_img_idx_kernel(int *img_idx, uint32_t *img_offsets, uint32_t *bas_ij_idx, int npairs, PBCIntEnvVars envs, float *exps, float *log_coef, float log_cutoff) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int pair_id = item.get_global_id(0); + #else int pair_id = blockIdx.x * blockDim.x + threadIdx.x; + #endif + if (pair_id >= npairs) { return; } @@ -1105,11 +1153,12 @@ int PBCsr_int3c2e_latsum23(double *out, double omega, PBCIntEnvVars *envs, uint3 int aux_offset, int nauxbas, int naux, int to_sph, float *diffuse_exps, float *diffuse_coefs, float log_cutoff) { - cudaFuncSetAttribute(pbc_int3c2e_latsum23_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + cudaMemset(head, 0, sizeof(int)); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); + #ifndef USE_SYCL + cudaFuncSetAttribute(pbc_int3c2e_latsum23_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); pbc_int3c2e_latsum23_kernel<<>>( out, omega, *envs, pool, task_pool, c2s_pool, shm_size, bas_ij_idx, shl_pair_offsets, ksh_offsets, img_idx, sp_img_offsets, @@ -1122,6 +1171,22 @@ int PBCsr_int3c2e_latsum23(double *out, double omega, PBCIntEnvVars *envs, uint3 fprintf(stderr, "CUDA Error in fill_int3c2e: %s\n", cudaGetErrorString(err)); return 1; } + #else + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<1>(workers * THREADS, THREADS), [=](auto item) { + pbc_int3c2e_latsum23_kernel( + out, omega, dev_envs, pool, task_pool, c2s_pool, shm_size, + bas_ij_idx, shl_pair_offsets, ksh_offsets, img_idx, sp_img_offsets, + gout_stride_lookup, ao_pair_loc, + ao_pair_offset, aux_offset, nauxbas, naux, to_sph, + diffuse_exps, diffuse_coefs, log_cutoff, + head, nbatches_shl_pair, nbatches_ksh, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #endif return 0; } @@ -1132,6 +1197,12 @@ int bvk_ovlp_img_counts(int *img_counts, PBCIntEnvVars *envs, constexpr int threads = 512; int bvk_nbas = envs->nbas * envs->bvk_ncells; int nbatches = (envs->nbas * bvk_nbas + threads-1) / threads; + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<1>(nbatches * threads, threads), [=](auto item) { + ovlp_img_counts_kernel(img_counts, dev_envs, exps, log_coef, log_cutoff, permutation_symmetry); + }); + #else ovlp_img_counts_kernel<<>>( img_counts, *envs, exps, log_coef, log_cutoff, permutation_symmetry); cudaError_t err = cudaGetLastError(); @@ -1139,6 +1210,7 @@ int bvk_ovlp_img_counts(int *img_counts, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in bvk_ovlp_img_counts: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -1147,6 +1219,12 @@ int bvk_ovlp_img_idx(int *img_idx, uint32_t *img_offsets, uint32_t *bas_ij_idx, { constexpr int threads = 512; int blocks = (npairs + threads-1) / threads; + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + ovlp_img_idx_kernel(img_idx, img_offsets, bas_ij_idx, npairs, dev_envs, exps, log_coef, log_cutoff); + }); + #else ovlp_img_idx_kernel<<>>( img_idx, img_offsets, bas_ij_idx, npairs, *envs, exps, log_coef, log_cutoff); cudaError_t err = cudaGetLastError(); @@ -1154,6 +1232,7 @@ int bvk_ovlp_img_idx(int *img_idx, uint32_t *img_offsets, uint32_t *bas_ij_idx, fprintf(stderr, "CUDA Error in bvk_ovlp_img_idx: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/fill_triu.cu b/gpu4pyscf/lib/pbc/fill_triu.cu index 3d5ccb31e..5305aa864 100644 --- a/gpu4pyscf/lib/pbc/fill_triu.cu +++ b/gpu4pyscf/lib/pbc/fill_triu.cu @@ -10,12 +10,26 @@ __global__ static void fill_indexed_triu_kernel(double *out, int *tril_idx, int *ki_idx, int npairs, int nao, int naux) { - int pair_id = blockIdx.x * BLOCK_SIZE + threadIdx.y; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + int blockIdx_y = item.get_group(0); + int blockDim_x = item.get_local_range(1); + #else + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int blockDim_x = blockDim.x; + #endif + int pair_id = blockIdx_x * BLOCK_SIZE + threadIdx_y; if (pair_id >= npairs) { return; } int pair_ij = tril_idx[pair_id]; - int kp = blockIdx.y; + int kp = blockIdx_y; size_t Nao = nao; size_t Naux = naux; int ij = pair_ij + kp * Nao * Nao; @@ -25,7 +39,7 @@ void fill_indexed_triu_kernel(double *out, int *tril_idx, int *ki_idx, int ji = (ki * nao + j) * Nao + i; if (ji == ij) return; - for (int aux_id = threadIdx.x; aux_id < naux; aux_id += blockDim.x) { + for (int aux_id = threadIdx_x; aux_id < naux; aux_id += blockDim_x) { out[ji*Naux+aux_id] = out[ij*Naux+aux_id]; } } @@ -34,7 +48,17 @@ __global__ static void fill_bvk_triu_kernel(double *out, int *pair_address, int *conj_mapping, int bvk_ncells, int nao, int naux) { - int ij = pair_address[blockIdx.x]; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int threadIdx_x = item.get_local_id(0); + int blockIdx_x = item.get_group(0); + int blockDim_x = item.get_local_range(0); + #else + int threadIdx_x = threadIdx.x; + int blockIdx_x = blockIdx.x; + int blockDim_x = blockDim.x; + #endif + int ij = pair_address[blockIdx_x]; int r = ij / nao; int j = ij - nao * r; int i = r / bvk_ncells; @@ -44,7 +68,7 @@ void fill_bvk_triu_kernel(double *out, int *pair_address, int *conj_mapping, if (ji == ij) return; size_t Naux = naux; - for (int aux_id = threadIdx.x; aux_id < naux; aux_id += blockDim.x) { + for (int aux_id = threadIdx_x; aux_id < naux; aux_id += blockDim_x) { out[ji*Naux+aux_id] = out[ij*Naux+aux_id]; } } @@ -53,7 +77,12 @@ __global__ static void fill_bvk_triu_naux1_kernel(double *out, int *pair_address, int *conj_mapping, int npairs, int bvk_ncells, int nao) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int pair_id = item.get_global_id(1); + #else int pair_id = blockIdx.x * blockDim.x + threadIdx.x; + #endif if (pair_id >= npairs) return; int ij = pair_address[pair_id]; int r = ij / nao; @@ -69,8 +98,14 @@ void fill_bvk_triu_naux1_kernel(double *out, int *pair_address, int *conj_mappin __global__ static void fill_bvk_triu_axis0_kernel(double *out, int *conj_mapping, int bvk_ncells, int nao) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int j = item.get_global_id(1); + int i = item.get_global_id(0); + #else int j = blockIdx.x * blockDim.x + threadIdx.x; int i = blockIdx.y * blockDim.y + threadIdx.y; + #endif if (i >= nao || j >= nao || i <= j) { return; } @@ -87,6 +122,13 @@ extern "C" { int fill_indexed_triu(double *out, int *tril_idx, int *ki_idx, int npairs, int nkpts, int nao, int naux) { + #ifdef USE_SYCL + sycl::range<2> threads(BLOCK_SIZE, 32); + sycl::range<2> blocks(nkpts, (npairs+BLOCK_SIZE-1)/BLOCK_SIZE); + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + fill_indexed_triu_kernel(out, tril_idx, ki_idx, npairs, nao, naux); + }); + #else dim3 threads(32, BLOCK_SIZE); dim3 blocks((npairs+BLOCK_SIZE-1)/BLOCK_SIZE, nkpts); fill_indexed_triu_kernel<<>>( @@ -96,12 +138,25 @@ int fill_indexed_triu(double *out, int *tril_idx, int *ki_idx, fprintf(stderr, "CUDA Error in fill_indexed_triu: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } int fill_bvk_triu(double *out, int *pair_address, int *conj_mapping, int npairs, int bvk_ncells, int nao, int naux) { + #ifdef USE_SYCL + if (naux == 1) { + int blocks = (npairs+255)/256; + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * 256, 256), [=](auto item) { + fill_bvk_triu_naux1_kernel(out, pair_address, conj_mapping, npairs, bvk_ncells, nao); + }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<1>(npairs * 256, 256), [=](auto item) { + fill_bvk_triu_kernel(out, pair_address, conj_mapping, bvk_ncells, nao, naux); + }); + } + #else if (naux == 1) { dim3 blocks((npairs+255)/256); fill_bvk_triu_naux1_kernel<<>>( @@ -115,13 +170,21 @@ int fill_bvk_triu(double *out, int *pair_address, int *conj_mapping, fprintf(stderr, "CUDA Error in fill_bvk_triu: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } int fill_bvk_triu_axis0(double *out, int *conj_mapping, int nao, int bvk_ncells) { - dim3 threads(BLOCK_SIZE, BLOCK_SIZE); int nao_b = (nao + BLOCK_SIZE-1) / BLOCK_SIZE; + #ifdef USE_SYCL + sycl::range<2> threads(BLOCK_SIZE, BLOCK_SIZE); + sycl::range<2> blocks(nao_b, nao_b); + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + fill_bvk_triu_axis0_kernel(out, conj_mapping, bvk_ncells, nao); + }); + #else + dim3 threads(BLOCK_SIZE, BLOCK_SIZE); dim3 blocks(nao_b, nao_b); fill_bvk_triu_axis0_kernel<<>>(out, conj_mapping, bvk_ncells, nao); cudaError_t err = cudaGetLastError(); @@ -129,6 +192,7 @@ int fill_bvk_triu_axis0(double *out, int *conj_mapping, int nao, int bvk_ncells) fprintf(stderr, "CUDA Error in fill_bvk_triu_axis0: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/ft_ao.cu b/gpu4pyscf/lib/pbc/ft_ao.cu index b4a09ba0f..1b1fa334e 100644 --- a/gpu4pyscf/lib/pbc/ft_ao.cu +++ b/gpu4pyscf/lib/pbc/ft_ao.cu @@ -1,5 +1,5 @@ /* - * Copyright 2024-2026 The PySCF Developers. All Rights Reserved. + * Copyright 2024 The PySCF Developers. All Rights Reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -29,9 +29,11 @@ #define WARP_SIZE 32 #endif #define WARPS 8 -#define THREADS 256 -#define NG_PER_BLOCK WARP_SIZE #define FT_AO_THREADS (WARP_SIZE*4) +// One shell per block (nsh_per_block == 1): every thread in the block then +// sees the same shell's iprim, so the primitive loop's __syncthreads() trip +// count is uniform without needing a per-block max-iprim workaround. +#define NG_PER_BLOCK FT_AO_THREADS #define GOUT_WIDTH 30 // pi^1.5 #define OVERLAP_FAC 5.56832799683170787 @@ -43,22 +45,42 @@ __global__ static void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) { + int nsh_per_block = FT_AO_THREADS / NG_PER_BLOCK; + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + + int sh_block_id = item.get_group_range(0) - item.get_group(0) - 1; + int Gv_block_id = item.get_group(1); + int sh_id_in_block = item.get_local_id(0); + int Gv_id_in_block = item.get_local_id(1); + + double (&g)[(AUXL+1)*FT_AO_THREADS * 6] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + #else int sh_block_id = gridDim.y - blockIdx.y - 1; int Gv_block_id = blockIdx.x; - int nsh_per_block = FT_AO_THREADS / NG_PER_BLOCK; int sh_id_in_block = threadIdx.y; int Gv_id_in_block = threadIdx.x; + + __shared__ double g[(AUXL+1)*FT_AO_THREADS * 6]; + #endif + int sh_id = sh_block_id * nsh_per_block + sh_id_in_block; - if (sh_id >= envs.nbas) { - return; - } + // A work-item whose shell index falls outside envs.nbas cannot return + // here: every __syncthreads() below is a real SYCL group_barrier, which + // -- unlike CUDA's warp-retirement semantics -- requires every work-item + // in the group to reach it. Clamp to a valid shell instead so out-of- + // range lanes take the identical control-flow path (and therefore the + // same barrier count) as their neighbours; the final write-out below is + // masked so the clamped, discarded computation never reaches memory. + int valid = sh_id < envs.nbas; + int sh_id_clamped = valid ? sh_id : envs.nbas - 1; int *atm = envs.atm; int *bas = envs.bas; double *env = envs.env; - int li = bas[sh_id*BAS_SLOTS+ANG_OF]; + int li = bas[sh_id_clamped*BAS_SLOTS+ANG_OF]; int nfi = c_nf[li]; - int iprim = bas[sh_id*BAS_SLOTS+NPRIM_OF]; + int iprim = bas[sh_id_clamped*BAS_SLOTS+NPRIM_OF]; int Gv_id = Gv_block_id * NG_PER_BLOCK + Gv_id_in_block; double kx = 0; double ky = 0; @@ -71,14 +93,13 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) double kk = kx * kx + ky * ky + kz * kz; int gx_len = (AUXL+1) * FT_AO_THREADS; - __shared__ double g[(AUXL+1)*FT_AO_THREADS * 6]; double *gxR = g + (AUXL+1) * NG_PER_BLOCK * sh_id_in_block + Gv_id_in_block; double *gxI = gxR + gx_len; double *gyR = gxR + gx_len*2; double *gyI = gxR + gx_len*3; double *gzR = gxR + gx_len*4; double *gzI = gxR + gx_len*5; - int *idx = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx = _c_cartesian_lexical_xyz + lex_xyz_offset(li); constexpr int aux_nf = (AUXL+1)*(AUXL+2)/2; double goutR[aux_nf]; @@ -95,9 +116,9 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) double s0zR, s1zR, s2zR; double s0zI, s1zI, s2zI; - int ia = bas[sh_id*BAS_SLOTS+ATOM_OF]; - double *expi = env + bas[sh_id*BAS_SLOTS+PTR_EXP]; - double *ci = env + bas[sh_id*BAS_SLOTS+PTR_COEFF]; + int ia = bas[sh_id_clamped*BAS_SLOTS+ATOM_OF]; + double *expi = env + bas[sh_id_clamped*BAS_SLOTS+PTR_EXP]; + double *ci = env + bas[sh_id_clamped*BAS_SLOTS+PTR_COEFF]; double *ri = env + atm[ia*ATM_SLOTS+PTR_COORD]; for (int ip = 0; ip < iprim; ++ip) { __syncthreads(); @@ -187,9 +208,9 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) } } - if (Gv_id < nGv) { + if (valid && Gv_id < nGv) { size_t stride = (size_t)nGv * OF_COMPLEX; - double *aft_tensor = out + ((size_t)envs.ao_loc[sh_id] * nGv + Gv_id) * OF_COMPLEX; + double *aft_tensor = out + ((size_t)envs.ao_loc[sh_id_clamped] * nGv + Gv_id) * OF_COMPLEX; #pragma unroll for (int n = 0; n < aux_nf; ++n) { if (n >= nfi) break; @@ -203,36 +224,61 @@ __global__ static void ft_aopair_kernel(double *out, PBCIntEnvVars envs, double *pool, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *img_idx, uint32_t *img_offsets, int *gout_stride_lookup, int *ao_pair_loc, int ao_pair_offset, - double *Gv, int nGv, int *ao_loc, int compressing, int to_sph, - int *head, int nbatches_shl_pair) + double *Gv, int nGv, int *ao_loc, int compressing, int to_sph + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { - constexpr int nGv_per_block = WARP_SIZE; - int thread_id = threadIdx.x; - int Gv_id_in_block = thread_id % nGv_per_block; - int warp_id = thread_id / nGv_per_block; - __shared__ int Gv_block_id, sp_block_id; - double *c2s_pool = pool + blockIdx.x * POOL_SIZE; -while (1) { - if (thread_id == 0) { - int batch_id = atomicAdd(head, 1); - Gv_block_id = batch_id / nbatches_shl_pair; - sp_block_id = batch_id - Gv_block_id * nbatches_shl_pair; - } - __syncthreads(); - if (Gv_block_id * nGv_per_block >= nGv) { - return; - } + #ifdef USE_SYCL + int sp_block_id = item.get_group_range(1) - item.get_group(1) - 1; + int Gv_block_id = item.get_group(0); + int Gv_id_in_block = item.get_local_id(1); + int warp_id = item.get_local_id(0); + int blockDim_y = item.get_local_range(0); + + auto thread_block = item.get_group(); + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nao = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &img_max = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int (&img_counts)[WARPS] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double *shared_memory = reinterpret_cast(shm_mem); + #else + int sp_block_id = gridDim.x - blockIdx.x - 1; + int Gv_block_id = blockIdx.y; + int Gv_id_in_block = threadIdx.x; + int warp_id = threadIdx.y; + int blockDim_y = blockDim.y; - int ncells = envs.bvk_ncells; - int bvk_nbas = envs.nbas * ncells; - int *bas = envs.bas; - double *env = envs.env; - double *img_coords = envs.img_coords; __shared__ int shl_pair0, shl_pair1; __shared__ int li, lj; __shared__ int iprim, jprim; __shared__ int nao; __shared__ int gout_stride, nsp_per_block; + __shared__ int img_max; + __shared__ int img_counts[WARPS]; + + extern __shared__ double shared_memory[]; + #endif + + // ft_aopair_kernel's grid-points-per-block is independent of ft_ao_bdiv_kernel's + // NG_PER_BLOCK (bumped to FT_AO_THREADS for that kernel's divergence fix); upstream + // hardcodes WARP_SIZE here and sizes shm_size/grid dims on the host to match. + constexpr int nGv_per_block = WARP_SIZE; + int thread_id = Gv_id_in_block + nGv_per_block * warp_id; + int ncells = envs.bvk_ncells; + int bvk_nbas = envs.nbas * ncells; + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -247,7 +293,7 @@ while (1) { // cannot handle spherical integrals nao = ao_loc[envs.nbas]; gout_stride = gout_stride_lookup[li*LMAX1+lj]; - nsp_per_block = WARPS / gout_stride; + nsp_per_block = blockDim_y / gout_stride; } __syncthreads(); int nGsp_per_block = nGv_per_block * nsp_per_block; @@ -260,7 +306,6 @@ while (1) { int stride_j = li + 1; int g_size = stride_j * (lj + 1); int gx_len = g_size * nGsp_per_block; - extern __shared__ double shared_memory[]; double *gxR = shared_memory + nGv_per_block * sp_id + Gv_id_in_block; double *gxI = gxR + gx_len; double *gyR = gxR + gx_len*2; @@ -277,6 +322,7 @@ while (1) { if (thread_id < nfj * 3) { idx_j[thread_id] = lex_xyz_address(lj, thread_id) * stride_j * nGsp_per_block; } + double *c2s_pool = pool + get_smid() * POOL_SIZE; int Gv_id = Gv_block_id * nGv_per_block + Gv_id_in_block; double kx = 0; @@ -300,12 +346,27 @@ while (1) { int jsh = bas_ij % bvk_nbas; int img0 = img_offsets[pair_ij]; int img1 = img_offsets[pair_ij+1]; - __shared__ int img_max; - __shared__ int img_counts[WARPS]; if (Gv_id_in_block == 0) { img_counts[warp_id] = img1 - img0; } __syncthreads(); +#ifdef USE_SYCL + // A sub-group shuffle reduction restricted to thread_id < WARPS is UB in + // SYCL: sub-group collectives require every lane of the *hardware* + // sub-group to participate uniformly, but this HW's sub-group width + // (16 on Intel Data Center GPU Max) doesn't match WARPS (8), so only + // part of the sub-group would call shift_group_left. CUDA's masked + // __shfl_down_sync tolerates this (warp is a fixed 32 lanes and the + // mask exactly matches the active lanes), so keep that path for CUDA + // and just scan img_counts[] serially here instead. + if (thread_id == 0) { + int count = img_counts[0]; + for (int w = 1; w < WARPS; ++w) { + count = max(count, img_counts[w]); + } + img_max = count; + } +#else if (thread_id < WARPS) { int count = img_counts[thread_id]; unsigned mask = (1u << WARPS) - 1; @@ -316,6 +377,7 @@ while (1) { img_max = count; } } +#endif __syncthreads(); int expi = bas[ish*BAS_SLOTS+PTR_EXP]; @@ -435,7 +497,7 @@ while (1) { } } __syncthreads(); - if (pair_idx < shl_pair1 && img < img1 && Gv_id < nGv) { + if (pair_idx < shl_pair1 && img < img1) { float div_nfi = c_div_nf[li]; #pragma unroll for (int n = 0; n < GOUT_WIDTH; ++n) { @@ -1149,15 +1211,23 @@ while (1) { } } } -} extern "C" { int build_ft_ao(double *out, RysIntEnvVars *envs, int ngrids, double *grids, int nbas) { int nsh_per_block = FT_AO_THREADS/NG_PER_BLOCK; - dim3 threads(NG_PER_BLOCK, nsh_per_block); int nbatches_grids = (ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK; int nbatches_shls = (nbas + nsh_per_block - 1) / nsh_per_block; + + #ifdef USE_SYCL + sycl::range<2> threads(nsh_per_block, NG_PER_BLOCK); + sycl::range<2> blocks(nbatches_shls, nbatches_grids); + auto dev_envs = *envs; + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ft_ao_bdiv_kernel(out, dev_envs, ngrids, grids); + }); + #else + dim3 threads(NG_PER_BLOCK, nsh_per_block); dim3 blocks(nbatches_grids, nbatches_shls); ft_ao_bdiv_kernel<<>>(out, *envs, ngrids, grids); cudaError_t err = cudaGetLastError(); @@ -1165,29 +1235,51 @@ int build_ft_ao(double *out, RysIntEnvVars *envs, int ngrids, double *grids, int fprintf(stderr, "CUDA Error in ft_ao_bdiv_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } +// `head` is the persistent-worker task counter upstream's kernel consumes via +// atomicAdd. Both branches below map the grid statically instead, so the +// pointer is unused -- but it stays in the signature because the ctypes +// caller in pbc/df/ft_ao.py is upstream code and still passes it; dropping it +// shifts every later argument by one and aborts the process. int build_ft_aopair(double *out, PBCIntEnvVars *envs, double *pool, int *head, int shm_size, int nbatches_shl_pair, int *shl_pair_offsets, uint32_t *bas_ij_idx, int *img_idx, uint32_t *img_offsets, int *gout_stride_lookup, int *ao_pair_loc, int ao_pair_offset, double *grids, int ngrids, int *ao_loc, int compressing, int to_sph) { + (void)head; + constexpr int nGv_per_block = WARP_SIZE; + int Gv_batches = (ngrids + nGv_per_block - 1) / nGv_per_block; + #ifdef USE_SYCL + sycl::range<2> threads(WARPS, nGv_per_block); + sycl::range<2> blocks(Gv_batches, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ft_aopair_kernel(out, dev_envs, pool, shl_pair_offsets, bas_ij_idx, img_idx, img_offsets, + gout_stride_lookup, ao_pair_loc, ao_pair_offset, grids, ngrids, + ao_loc, compressing, to_sph, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(ft_aopair_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - cudaDeviceProp prop; - cudaGetDeviceProperties(&prop, 0); - int workers = prop.multiProcessorCount; - cudaMemset(head, 0, sizeof(int)); - ft_aopair_kernel<<>>( + dim3 threads(nGv_per_block, WARPS); + dim3 blocks(nbatches_shl_pair, Gv_batches); + ft_aopair_kernel<<>>( out, *envs, pool, shl_pair_offsets, bas_ij_idx, img_idx, img_offsets, gout_stride_lookup, ao_pair_loc, ao_pair_offset, grids, ngrids, - ao_loc, compressing, to_sph, head, nbatches_shl_pair); + ao_loc, compressing, to_sph); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "CUDA Error in ft_aopair_kernel: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/ft_ao_ip1.cu b/gpu4pyscf/lib/pbc/ft_ao_ip1.cu index dba22d248..839cecc9c 100644 --- a/gpu4pyscf/lib/pbc/ft_ao_ip1.cu +++ b/gpu4pyscf/lib/pbc/ft_ao_ip1.cu @@ -43,15 +43,31 @@ __global__ void ft_aopair_ejk_ip1_kernel(double *out, double *dm, double *vG, double *Gv, PBCIntEnvVars envs, int nGv, int shm_size, int *bas_ij_idx, int *bas_ij_img_idx, - int *shl_pair_offsets, int permutation_symmetry) + int *shl_pair_offsets, int permutation_symmetry + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { constexpr int nGv_per_block = NGV_PER_BLOCK; constexpr int threads = NGV_PER_BLOCK * NSP_PER_BLOCK; + #ifdef USE_SYCL + int sp_block_id = item.get_group(1); + int Gv_block_id = item.get_group(0); + int Gv_id_in_block = item.get_local_id(1); + int thread_id = item.get_local_id(1) + item.get_local_id(0) * item.get_local_range(1); + double *shared_memory = reinterpret_cast(shm_mem); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + #else int sp_block_id = blockIdx.x; int Gv_block_id = blockIdx.y; int Gv_id_in_block = threadIdx.x; - int thread_id = threadIdx.x + threadIdx.y * blockDim.x; + extern __shared__ double shared_memory[]; + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + #endif int shl_pair0 = shl_pair_offsets[sp_block_id]; int shl_pair1 = shl_pair_offsets[sp_block_id+1]; int bas_ij0 = bas_ij_idx[shl_pair0]; @@ -73,8 +89,8 @@ void ft_aopair_ejk_ip1_kernel(double *out, double *dm, double *vG, double *Gv, gout_stride *= 2; } int nsp_per_block = NSP_PER_BLOCK / gout_stride; - int gout_id = threadIdx.y % gout_stride; - int sp_id = threadIdx.y / gout_stride; + int gout_id = threadIdx_y % gout_stride; + int sp_id = threadIdx_y / gout_stride; int Gv_gout_id = Gv_id_in_block + nGv_per_block * gout_id; int nGv_gout = nGv_per_block * gout_stride; int lij = li + lj + 1; @@ -100,7 +116,6 @@ void ft_aopair_ejk_ip1_kernel(double *out, double *dm, double *vG, double *Gv, } double kk = kx * kx + ky * ky + kz * kz; - extern __shared__ double shared_memory[]; double *gxR = shared_memory + g_size * nGv_per_block * sp_id + Gv_id_in_block; double *gxI = gxR + gx_len*1; double *gyR = gxR + gx_len*2; @@ -108,8 +123,8 @@ void ft_aopair_ejk_ip1_kernel(double *out, double *dm, double *vG, double *Gv, double *gzR = gxR + gx_len*4; double *gzI = gxR + gx_len*5; double *rjri = shared_memory + gx_len * 6 + sp_id; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { __syncthreads(); @@ -362,15 +377,33 @@ void ft_aopair_strain_deriv_kernel(double *out, double *sigma, double *dm, double *vG, double *Gv, PBCIntEnvVars envs, int nGv, int shm_size, int *bas_ij_idx, int *bas_ij_img_idx, - int *shl_pair_offsets, int permutation_symmetry) + int *shl_pair_offsets, int permutation_symmetry + #ifdef USE_SYCL + , sycl::nd_item<2> &item, char *shm_mem + #endif + ) { constexpr int nGv_per_block = NGV_PER_BLOCK; constexpr int threads = NGV_PER_BLOCK * NSP_PER_BLOCK; + + #ifdef USE_SYCL + int sp_block_id = item.get_group(1); + int Gv_block_id = item.get_group(0); + int Gv_id_in_block = item.get_local_id(1); + int thread_id = item.get_local_id(1) + item.get_local_id(0) * item.get_local_range(1); + double *shared_memory = reinterpret_cast(shm_mem); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + #else int sp_block_id = blockIdx.x; int Gv_block_id = blockIdx.y; int Gv_id_in_block = threadIdx.x; - int thread_id = threadIdx.x + threadIdx.y * blockDim.x; + extern __shared__ double shared_memory[]; + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + #endif + int shl_pair0 = shl_pair_offsets[sp_block_id]; int shl_pair1 = shl_pair_offsets[sp_block_id+1]; int bas_ij0 = bas_ij_idx[shl_pair0]; @@ -392,8 +425,8 @@ void ft_aopair_strain_deriv_kernel(double *out, double *sigma, gout_stride *= 2; } int nsp_per_block = NSP_PER_BLOCK / gout_stride; - int gout_id = threadIdx.y % gout_stride; - int sp_id = threadIdx.y / gout_stride; + int gout_id = threadIdx_y % gout_stride; + int sp_id = threadIdx_y / gout_stride; int Gv_gout_id = Gv_id_in_block + nGv_per_block * gout_id; int nGv_gout = nGv_per_block * gout_stride; int lij = li + lj + 1; @@ -419,7 +452,6 @@ void ft_aopair_strain_deriv_kernel(double *out, double *sigma, } double kk = kx * kx + ky * ky + kz * kz; - extern __shared__ double shared_memory[]; double *gxR = shared_memory + g_size * nGv_per_block * sp_id + Gv_id_in_block; double *gxI = gxR + gx_len*1; double *gyR = gxR + gx_len*2; @@ -427,8 +459,8 @@ void ft_aopair_strain_deriv_kernel(double *out, double *sigma, double *gzR = gxR + gx_len*4; double *gzI = gxR + gx_len*5; double *rjri = shared_memory + gx_len * 6 + sp_id; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); double sigma_xx = 0; double sigma_xy = 0; @@ -743,9 +775,23 @@ int PBC_ft_aopair_ej_ip1(double *out, double *dm, double *vG, double *GvT, int *bas_ij_idx, int *bas_ij_img_idx, int *shl_pair_offsets, int permutation_symmetry) { + int Gv_batches = (ngrids + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + + #ifdef USE_SYCL + sycl::range<2> threads(NSP_PER_BLOCK, NGV_PER_BLOCK); + sycl::range<2> blocks(Gv_batches, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ft_aopair_ejk_ip1_kernel(out, dm, vG, GvT, dev_envs, ngrids, shm_size, + bas_ij_idx, bas_ij_img_idx, shl_pair_offsets, permutation_symmetry, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(ft_aopair_ejk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 threads(NGV_PER_BLOCK, NSP_PER_BLOCK); - int Gv_batches = (ngrids + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; dim3 blocks(nbatches_shl_pair, Gv_batches); ft_aopair_ejk_ip1_kernel<<>>( out, dm, vG, GvT, *envs, ngrids, shm_size, @@ -755,6 +801,7 @@ int PBC_ft_aopair_ej_ip1(double *out, double *dm, double *vG, double *GvT, fprintf(stderr, "CUDA Error in ft_aopair_ej_ip1: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -763,9 +810,23 @@ int PBC_ft_aopair_ek_ip1(double *out, double *dm_vG, double *GvT, PBCIntEnvVars int *bas_ij_idx, int *bas_ij_img_idx, int *shl_pair_offsets, int permutation_symmetry) { + int Gv_batches = (ngrids + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + + #ifdef USE_SYCL + sycl::range<2> threads(NSP_PER_BLOCK, NGV_PER_BLOCK); + sycl::range<2> blocks(Gv_batches, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ft_aopair_ejk_ip1_kernel(out, dm_vG, NULL, GvT, dev_envs, ngrids, shm_size, + bas_ij_idx, bas_ij_img_idx, shl_pair_offsets, permutation_symmetry, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(ft_aopair_ejk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 threads(NGV_PER_BLOCK, NSP_PER_BLOCK); - int Gv_batches = (ngrids + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; dim3 blocks(nbatches_shl_pair, Gv_batches); ft_aopair_ejk_ip1_kernel<<>>( out, dm_vG, NULL, GvT, *envs, ngrids, shm_size, @@ -775,6 +836,7 @@ int PBC_ft_aopair_ek_ip1(double *out, double *dm_vG, double *GvT, PBCIntEnvVars fprintf(stderr, "CUDA Error in ft_aopair_ek_ip1: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -785,8 +847,21 @@ int PBC_ft_aopair_ej_strain_deriv(double *out, double *sigma, double *dm, int permutation_symmetry) { cudaFuncSetAttribute(ft_aopair_strain_deriv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - dim3 threads(NGV_PER_BLOCK, NSP_PER_BLOCK); int Gv_batches = (ngrids + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + #ifdef USE_SYCL + sycl::range<2> threads(NSP_PER_BLOCK, NGV_PER_BLOCK); + sycl::range<2> blocks(Gv_batches, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ft_aopair_strain_deriv_kernel(out, sigma, dm, vG, GvT, dev_envs, ngrids, shm_size, + bas_ij_idx, bas_ij_img_idx, shl_pair_offsets, permutation_symmetry, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(NGV_PER_BLOCK, NSP_PER_BLOCK); dim3 blocks(nbatches_shl_pair, Gv_batches); ft_aopair_strain_deriv_kernel<<>>( out, sigma, dm, vG, GvT, *envs, ngrids, shm_size, @@ -796,6 +871,7 @@ int PBC_ft_aopair_ej_strain_deriv(double *out, double *sigma, double *dm, fprintf(stderr, "CUDA Error in ft_aopair_ej_strain_deriv: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -805,9 +881,23 @@ int PBC_ft_aopair_ek_strain_deriv(double *out, double *sigma, int *bas_ij_idx, int *bas_ij_img_idx, int *shl_pair_offsets, int permutation_symmetry) { + int Gv_batches = (ngrids + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + + #ifdef USE_SYCL + sycl::range<2> threads(NSP_PER_BLOCK, NGV_PER_BLOCK); + sycl::range<2> blocks(Gv_batches, nbatches_shl_pair); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + ft_aopair_strain_deriv_kernel(out, sigma, dm_vG, NULL, GvT, dev_envs, ngrids, shm_size, + bas_ij_idx, bas_ij_img_idx, shl_pair_offsets, permutation_symmetry, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else cudaFuncSetAttribute(ft_aopair_strain_deriv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); dim3 threads(NGV_PER_BLOCK, NSP_PER_BLOCK); - int Gv_batches = (ngrids + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; dim3 blocks(nbatches_shl_pair, Gv_batches); ft_aopair_strain_deriv_kernel<<>>( out, sigma, dm_vG, NULL, GvT, *envs, ngrids, shm_size, @@ -817,6 +907,7 @@ int PBC_ft_aopair_ek_strain_deriv(double *out, double *sigma, fprintf(stderr, "CUDA Error in ft_aopair_ek_strain_deriv: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/int3c2e.cuh b/gpu4pyscf/lib/pbc/int3c2e.cuh index 4e49d764d..1131246a1 100644 --- a/gpu4pyscf/lib/pbc/int3c2e.cuh +++ b/gpu4pyscf/lib/pbc/int3c2e.cuh @@ -17,6 +17,7 @@ #pragma once #include +#include // WARP_SIZE: compile-time constant used for shared-memory sizing. // `warpSize` (HIP/CUDA device-runtime built-in) is not constexpr, diff --git a/gpu4pyscf/lib/pbc/int3c2e_create_tasks.cuh b/gpu4pyscf/lib/pbc/int3c2e_create_tasks.cuh index 36727ff93..d9fcad9ce 100644 --- a/gpu4pyscf/lib/pbc/int3c2e_create_tasks.cuh +++ b/gpu4pyscf/lib/pbc/int3c2e_create_tasks.cuh @@ -57,6 +57,9 @@ typedef struct { __device__ inline int mask_to_index(int keep, int *tmp_storage, int threads, int t_id) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); +#endif tmp_storage[t_id] = keep; __syncthreads(); for (int offset = 1; offset < threads; offset <<= 1) { @@ -80,7 +83,12 @@ void initialize_ijk_tasks(uint32_t *img_pool, uint32_t *rem_task_idx, uint32_t *bas_ij_idx, int *img_idx, uint32_t *sp_img_offsets, float *diffuse_exps, float *diffuse_coefs, float log_cutoff) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + #else int thread_id = threadIdx.x; + #endif int ncells = envs.bvk_ncells; int bvk_nbas = envs.nbas * ncells; int *bas = envs.bas; @@ -189,8 +197,14 @@ __device__ inline void _filter_ijk_tasks(uint32_t *rem_task_idx, int& num_ijk_tasks, ShellTripletTaskInfo *ijk_tasks_info, int *swap) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int threads = item.get_local_range(0); + #else int thread_id = threadIdx.x; int threads = blockDim.x * blockDim.y; + #endif int tot_tasks = num_ijk_tasks; __syncthreads(); if (thread_id == 0) { @@ -223,8 +237,14 @@ void _select_sub_ijk(uint32_t *sub_task_idx, int &num_sub_tasks, uint32_t *rem_task_idx, int num_ijk_tasks, ShellTripletTaskInfo *ijk_tasks_info, int *swap) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int threads = item.get_local_range(0); + #else int thread_id = threadIdx.x; int threads = blockDim.x * blockDim.y; + #endif __syncthreads(); if (thread_id == 0) { num_sub_tasks = 0; @@ -264,8 +284,14 @@ void _filter_jk_images(uint32_t *img_pool, uint32_t *rem_task_idx, int num_ijk_tasks, ShellTripletTaskInfo *ijk_tasks_info, PBCIntEnvVars &envs, int *sp_img_idx) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int &task_head = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else int thread_id = threadIdx.x; __shared__ int task_head; +#endif if (thread_id == 0) { task_head = THREADS; } @@ -357,6 +383,9 @@ void _filter_jk_images(uint32_t *img_pool, uint32_t *rem_task_idx, __device__ inline int warp_max(int val) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); +#endif for (int offset = warpSize / 2; offset > 0; offset >>= 1) { val = max(val, __shfl_down_sync(0xffffffff, val, offset)); } @@ -366,9 +395,16 @@ int warp_max(int val) __device__ inline void block_max(int val, int& out) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int thread_id = item.get_local_id(0); + int (&buf)[WARPS] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else int thread_id = threadIdx.x; - val = warp_max(val); __shared__ int buf[WARPS]; +#endif + + val = warp_max(val); int lane = thread_id % warpSize; int warp_id = thread_id / warpSize; if (lane == 0) { @@ -386,4 +422,3 @@ void block_max(int val, int& out) } __syncthreads(); } - diff --git a/gpu4pyscf/lib/pbc/nr_eval_gto.cu b/gpu4pyscf/lib/pbc/nr_eval_gto.cu index d0daa548d..f196a1e6d 100644 --- a/gpu4pyscf/lib/pbc/nr_eval_gto.cu +++ b/gpu4pyscf/lib/pbc/nr_eval_gto.cu @@ -25,6 +25,18 @@ #define LMAX 4 #define THREADS 256 +// Abstracts CUDA/SYCL 2D thread-index setup for GTO eval kernels. Used 8x in this file. +#ifdef USE_SYCL +#define SETUP_GTO_KERNEL() \ + auto item = syclex::this_work_item::get_nd_item<2>(); \ + int grid_id = item.get_global_id(1); \ + int bas_id = item.get_group(0); +#else +#define SETUP_GTO_KERNEL() \ + int grid_id = blockIdx.x * blockDim.x + threadIdx.x; \ + int bas_id = blockIdx.y; +#endif + template __device__ __forceinline__ void _cart_gto_ip2(double gto[], double gx[], double gy[], double gz[], double a2, double rx, double ry, double rz) @@ -157,7 +169,12 @@ void _eval_cart_deriv1_strain_tensor( double xi, double yi, double zi, double rrcutoff, int *bas, int nimgs, int nao, int ngrids) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int bas_id = item.get_group(0); +#else int bas_id = blockIdx.y; +#endif int nprim = bas[NPRIM_OF+bas_id*BAS_SLOTS]; double *expi = env + bas[bas_id*BAS_SLOTS+PTR_EXP]; double *ci = env + bas[bas_id*BAS_SLOTS+PTR_COEFF]; @@ -212,11 +229,10 @@ __global__ static void _cart_deriv0_kernel(double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -312,11 +328,10 @@ __global__ static void _cart_deriv1_kernel(double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -535,11 +550,10 @@ __global__ static void _cart_ip2_kernel(double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -636,11 +650,10 @@ __global__ static void _sph_deriv0_kernel(double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -768,11 +781,10 @@ __global__ static void _sph_deriv1_kernel(double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -1021,11 +1033,10 @@ __global__ static void _sph_ip2_kernel(double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -1163,11 +1174,10 @@ static void _cart_deriv0_strain_tensor_kernel( double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -1450,11 +1460,10 @@ static void _cart_deriv1_strain_tensor_kernel( double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { - int grid_id = blockIdx.x * blockDim.x + threadIdx.x; + SETUP_GTO_KERNEL(); if (grid_id >= ngrids) { return; } - int bas_id = blockIdx.y; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; @@ -1497,6 +1506,40 @@ int PBCeval_gto_deriv(double *out, PBCIntEnvVars *envs, { constexpr int ngrids_per_block = THREADS; int threads = ngrids_per_block; + + #ifdef USE_SYCL + sycl::range<2> thread(1, threads); + sycl::range<2> blocks(nbas, (ngrids+ngrids_per_block-1)/ngrids_per_block); + auto dev_envs = *envs; + switch (deriv) { + case 0: + if (cart) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _cart_deriv0_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _sph_deriv0_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + } + break; + case 1: + if (cart) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _cart_deriv1_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _sph_deriv1_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + } + break; + case 2: + if (cart) { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _cart_deriv1_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _cart_ip2_kernel(out+4*nao*ngrids, dev_envs, grids, ngrids, nao, rcut); }); + } else { + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _sph_deriv1_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _sph_ip2_kernel(out+4*nao*ngrids, dev_envs, grids, ngrids, nao, rcut); }); + } + break; + default: + fprintf(stderr, "PBCeval_gto deriv = %d not supported\n", deriv); + return 1; + } + #else dim3 blocks((ngrids+ngrids_per_block-1)/ngrids_per_block, nbas); switch (deriv) { case 0: @@ -1531,6 +1574,7 @@ int PBCeval_gto_deriv(double *out, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in PBCeval_gto: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -1544,6 +1588,22 @@ int PBCeval_gto_strain_tensor(double *out, PBCIntEnvVars *envs, } constexpr int ngrids_per_block = THREADS; int threads = ngrids_per_block; + #ifdef USE_SYCL + sycl::range<2> thread(1, threads); + sycl::range<2> blocks(nbas, (ngrids+ngrids_per_block-1)/ngrids_per_block); + auto dev_envs = *envs; + switch (deriv) { + case 0: + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _cart_deriv0_strain_tensor_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + break; + case 1: + sycl_get_queue()->parallel_for(sycl::nd_range<2>(blocks * thread, thread), [=](auto item) { _cart_deriv1_strain_tensor_kernel(out, dev_envs, grids, ngrids, nao, rcut); }); + break; + default: + fprintf(stderr, "PBCeval_gto_strain_tensor deriv = %d not supported\n", deriv); + return 1; + } + #else dim3 blocks((ngrids+ngrids_per_block-1)/ngrids_per_block, nbas); switch (deriv) { case 0: @@ -1561,6 +1621,9 @@ int PBCeval_gto_strain_tensor(double *out, PBCIntEnvVars *envs, fprintf(stderr, "CUDA Error in PBCeval_gto_strain_tensor: %s\n", cudaGetErrorString(err)); return 1; } + #endif return 0; } } + +#undef SETUP_GTO_KERNEL diff --git a/gpu4pyscf/lib/pbc/overlap.cu b/gpu4pyscf/lib/pbc/overlap.cu index b04e052a2..171137537 100644 --- a/gpu4pyscf/lib/pbc/overlap.cu +++ b/gpu4pyscf/lib/pbc/overlap.cu @@ -29,10 +29,90 @@ #define GOUT_WIDTH_IP1 18 #define REMOTE_THRESHOLD 50 +// Abstracts CUDA/SYCL thread-index setup for 1D overlap kernels. Used 10x in this file. +#ifdef USE_SYCL +#define KERNEL_SETUP() \ + int sp_block_id = item.get_group(0); \ + int thread_id = item.get_local_id(0); \ + auto thread_block = item.get_group(); \ + int &shl_pair0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &shl_pair1 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &li = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &lj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &iprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &jprim = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + double *g = reinterpret_cast(shm_mem); +#else +#define KERNEL_SETUP() \ + int sp_block_id = blockIdx.x; \ + int thread_id = threadIdx.x; \ + __shared__ int shl_pair0, shl_pair1; \ + __shared__ int li, lj, iprim, jprim; \ + __shared__ int gout_stride, nsp_per_block; \ + extern __shared__ double g[]; +#endif + +// Abstracts 1D kernel launch for overlap integrals. Used 10x in this file. +// cudaFuncSetAttribute (where needed) must be placed outside this macro. +// +// The PBCIntEnvVars struct is owned by these macros, not passed by the caller: +// - SYCL: the raw host pointer `envs` MUST NOT be dereferenced on the device. +// We copy `*envs` into a host-local value `dev_envs` before submit(), and the +// [=] lambda captures that value. Passing `*envs` into the lambda instead would +// defer the deref to device execution and fault (host pointer not GPU-mapped). +// - CUDA: `*envs` is dereferenced host-side at launch and copied into kernel +// params by value, so it is passed inline with no local. +// +// Two shapes exist: +// LAUNCH_OVERLAP_KERNEL -> KERNEL(out, envs, ) +// LAUNCH_OVERLAP_KERNEL_DM -> KERNEL(out, dm, envs, ) +#ifdef USE_SYCL +#define LAUNCH_OVERLAP_KERNEL(KERNEL, nbatches_, out_, ...) { \ + auto dev_envs = *envs; \ + sycl_get_queue()->submit([&](sycl::handler &cgh) { \ + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); \ + cgh.parallel_for(sycl::nd_range<1>(nbatches_ * THREADS, THREADS), [=](auto item) { \ + KERNEL(out_, dev_envs, __VA_ARGS__, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); \ + }); \ + }); \ +} +#define LAUNCH_OVERLAP_KERNEL_DM(KERNEL, nbatches_, out_, dm_, ...) { \ + auto dev_envs = *envs; \ + sycl_get_queue()->submit([&](sycl::handler &cgh) { \ + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); \ + cgh.parallel_for(sycl::nd_range<1>(nbatches_ * THREADS, THREADS), [=](auto item) { \ + KERNEL(out_, dm_, dev_envs, __VA_ARGS__, item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); \ + }); \ + }); \ +} +#else +#define LAUNCH_OVERLAP_KERNEL(KERNEL, nbatches_, out_, ...) { \ + KERNEL<<>>(out_, *envs, __VA_ARGS__); \ + cudaError_t err = cudaGetLastError(); \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error in " #KERNEL ": %s\n", cudaGetErrorString(err)); \ + return 1; \ + } \ +} +#define LAUNCH_OVERLAP_KERNEL_DM(KERNEL, nbatches_, out_, dm_, ...) { \ + KERNEL<<>>(out_, dm_, *envs, __VA_ARGS__); \ + cudaError_t err = cudaGetLastError(); \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error in " #KERNEL ": %s\n", cudaGetErrorString(err)); \ + return 1; \ + } \ +} +#endif + __inline__ __device__ void vrr_hrr(double *gx, double *rjri, double ai, double aj, double cicj, int li, int lj, int gout_id, int gout_stride, int nsp_per_block) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + #endif int stride_j = li + 1; int g_size = (li + 1) * (lj + 1); int gx_len = g_size * nsp_per_block; @@ -78,17 +158,17 @@ void vrr_hrr(double *gx, double *rjri, double ai, double aj, double cicj, __global__ static void int1e_ovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -107,7 +187,6 @@ void int1e_ovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 1) * (lj + 1); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -218,17 +297,17 @@ void int1e_ovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, static __global__ void int1e_kin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -247,7 +326,6 @@ void int1e_kin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 3) * (lj + 1); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -378,17 +456,17 @@ void int1e_kin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, __global__ static void int1e_r2_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -407,7 +485,6 @@ void int1e_r2_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 1) * (lj + 3); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -540,17 +617,17 @@ void int1e_r2_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, __global__ static void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -569,7 +646,6 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 1) * (lj + 5); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -712,17 +788,17 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, __global__ static void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -741,7 +817,6 @@ void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 1) * (lj + 4); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -909,17 +984,17 @@ void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, __global__ static void int1e_r4_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -938,7 +1013,6 @@ void int1e_r4_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 1) * (lj + 6); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -1112,17 +1186,17 @@ void int1e_r4_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, static __global__ void int1e_ipovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -1141,7 +1215,6 @@ void int1e_ipovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 2) * (lj + 1); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -1273,17 +1346,17 @@ void int1e_ipovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, static __global__ void int1e_ipkin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) + int naoi, int naoj, size_t ij_offset + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int nbas = envs.cell0_nbas * envs.bvk_ncells; int *bas = envs.bas; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -1302,7 +1375,6 @@ void int1e_ipkin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int gout_id = thread_id / nsp_per_block; int g_size = (li + 4) * (lj + 1); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -1464,10 +1536,13 @@ void int1e_ipkin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, static __global__ void ovlp_strain_deriv_kernel(double *out, double *dm, PBCIntEnvVars envs, int *shl_pair_offsets, int *bas_ij_idx, - int *gout_stride_lookup) + int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int *bas = envs.bas; int cell0_nbas = envs.cell0_nbas; int nbas = envs.cell0_nbas * envs.bvk_ncells; @@ -1475,9 +1550,6 @@ void ovlp_strain_deriv_kernel(double *out, double *dm, PBCIntEnvVars envs, int nao = ao_loc[cell0_nbas]; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -1497,7 +1569,6 @@ void ovlp_strain_deriv_kernel(double *out, double *dm, PBCIntEnvVars envs, int g_size = (li + 2) * (lj + 1); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -1642,10 +1713,13 @@ void ovlp_strain_deriv_kernel(double *out, double *dm, PBCIntEnvVars envs, static __global__ void kin_strain_deriv_kernel(double *out, double *dm, PBCIntEnvVars envs, int *shl_pair_offsets, int *bas_ij_idx, - int *gout_stride_lookup) + int *gout_stride_lookup + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; + KERNEL_SETUP(); int *bas = envs.bas; int cell0_nbas = envs.cell0_nbas; int nbas = envs.cell0_nbas * envs.bvk_ncells; @@ -1653,9 +1727,6 @@ void kin_strain_deriv_kernel(double *out, double *dm, PBCIntEnvVars envs, int nao = ao_loc[cell0_nbas]; double *env = envs.env; double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; if (thread_id == 0) { shl_pair0 = shl_pair_offsets[sp_block_id]; shl_pair1 = shl_pair_offsets[sp_block_id+1]; @@ -1675,7 +1746,6 @@ void kin_strain_deriv_kernel(double *out, double *dm, PBCIntEnvVars envs, int g_size = (li + 4) * (lj + 1); int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; double *gx = g + sp_id; double *gy = g + gx_len + sp_id; double *gz = g + gx_len * 2 + sp_id; @@ -1855,7 +1925,12 @@ void ovlp_mask_estimation_kernel(int8_t *ovlp_mask, float *exps, float *log_coef PBCIntEnvVars envs, int hermi, float log_cutoff, double *bvkmesh_Ls, int ish0, int ish1, int jsh0, int jsh1) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + size_t pair_ij = item.get_global_id(0); + #else size_t pair_ij = blockIdx.x * (size_t)blockDim.x + threadIdx.x; + #endif int nish = ish1 - ish0; int njsh = jsh1 - jsh0; if (pair_ij >= (size_t)nish * njsh * envs.bvk_ncells) { @@ -1930,14 +2005,9 @@ int PBCint1e_ovlp(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_ovlp_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_ovlp_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_ovlp_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_ovlp kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -1947,14 +2017,9 @@ int PBCint1e_kin(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_kin_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_kin_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_kin_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_ovlp kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -1964,14 +2029,9 @@ int PBCint1e_r2_origi(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_r2_origi_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r2_origi_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_r2_origi_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r2_origi kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -1981,14 +2041,9 @@ int PBCint1e_r4_origi(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_r4_origi_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r4_origi_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_r4_origi_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r4_origi kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -1998,14 +2053,9 @@ int PBCint1e_r2_origi_ip2(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_r2_origi_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r2_origi_ip2_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_r2_origi_ip2_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r2_origi_ip2 kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -2015,14 +2065,9 @@ int PBCint1e_r4_origi_ip2(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_r4_origi_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r4_origi_ip2_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_r4_origi_ip2_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r4_origi_ip2 kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -2032,14 +2077,9 @@ int PBCint1e_ipovlp(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_ipovlp_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_ipovlp_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_ipovlp_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_ipovlp kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -2049,14 +2089,9 @@ int PBCint1e_ipkin(double *out, PBCIntEnvVars *envs, int shm_size, int naoi, int naoj, size_t ij_offset) { cudaFuncSetAttribute(int1e_ipkin_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_ipkin_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + LAUNCH_OVERLAP_KERNEL(int1e_ipkin_kernel, nbatches_shl_pair, + out, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_ipkin kernel: %s\n", cudaGetErrorString(err)); - return 1; - } return 0; } @@ -2065,13 +2100,8 @@ int PBCovlp_strain_deriv(double *out, double *dm, int *shl_pair_offsets, int *bas_ij_idx, int *gout_stride_lookup) { cudaFuncSetAttribute(ovlp_strain_deriv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - ovlp_strain_deriv_kernel<<>>( - out, dm, *envs, shl_pair_offsets, bas_ij_idx, gout_stride_lookup); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in ovlp_strain_deriv kernel: %s\n", cudaGetErrorString(err)); - return 1; - } + LAUNCH_OVERLAP_KERNEL_DM(ovlp_strain_deriv_kernel, nbatches_shl_pair, + out, dm, shl_pair_offsets, bas_ij_idx, gout_stride_lookup); return 0; } @@ -2080,13 +2110,8 @@ int PBCkin_strain_deriv(double *out, double *dm, int *shl_pair_offsets, int *bas_ij_idx, int *gout_stride_lookup) { cudaFuncSetAttribute(kin_strain_deriv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - kin_strain_deriv_kernel<<>>( - out, dm, *envs, shl_pair_offsets, bas_ij_idx, gout_stride_lookup); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in kin_strain_deriv kernel: %s\n", cudaGetErrorString(err)); - return 1; - } + LAUNCH_OVERLAP_KERNEL_DM(kin_strain_deriv_kernel, nbatches_shl_pair, + out, dm, shl_pair_offsets, bas_ij_idx, gout_stride_lookup); return 0; } @@ -2103,8 +2128,22 @@ void PBCovlp_mask_estimation(int8_t *ovlp_mask, float *exps, float *log_coeff, size_t njsh = jsh1 - jsh0; size_t npairs = nish * ncells * njsh; int nbatches = (npairs + 255) / 256; + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + cgh.parallel_for(sycl::nd_range<1>(nbatches * 256, 256), [=](auto item) { + ovlp_mask_estimation_kernel(ovlp_mask, exps, log_coeff, dev_envs, hermi, log_cutoff, bvkmesh_Ls, + ish0, ish1, jsh0, jsh1); + }); + }); + #else ovlp_mask_estimation_kernel<<>>( ovlp_mask, exps, log_coeff, *envs, hermi, log_cutoff, bvkmesh_Ls, ish0, ish1, jsh0, jsh1); + #endif } } + +#undef KERNEL_SETUP +#undef LAUNCH_OVERLAP_KERNEL +#undef LAUNCH_OVERLAP_KERNEL_DM diff --git a/gpu4pyscf/lib/pbc/pbc.cuh b/gpu4pyscf/lib/pbc/pbc.cuh index e487599df..164776a21 100644 --- a/gpu4pyscf/lib/pbc/pbc.cuh +++ b/gpu4pyscf/lib/pbc/pbc.cuh @@ -15,6 +15,8 @@ */ #pragma once +#include + #ifdef __CUDACC__ extern __constant__ int16_t c_pair_idx[]; extern __constant__ int c_pair_offsets[]; diff --git a/gpu4pyscf/lib/pbc/pbc_driver.cu b/gpu4pyscf/lib/pbc/pbc_driver.cu index 1eaaef8f6..6f9e859eb 100644 --- a/gpu4pyscf/lib/pbc/pbc_driver.cu +++ b/gpu4pyscf/lib/pbc/pbc_driver.cu @@ -14,6 +14,8 @@ * limitations under the License. */ +#include + __constant__ int c_nf[] = { 1, 3, diff --git a/gpu4pyscf/lib/pbc/rys_contract_j.cu b/gpu4pyscf/lib/pbc/rys_contract_j.cu index 846311c7d..7fdf3e284 100644 --- a/gpu4pyscf/lib/pbc/rys_contract_j.cu +++ b/gpu4pyscf/lib/pbc/rys_contract_j.cu @@ -28,6 +28,18 @@ #include "gvhf-rys/build_rys_gxyz.cuh" #include "pbc/create_tasks.cu" +#ifdef USE_SYCL +// libpbc's OWN gxyz offset table, defined in pbc/rys_contract_k.cu. +// +// This TU includes gvhf-rys/vhf.cuh, which declares libgvhf_rys's +// s_rys_gxyz_offset. Before the rename both libraries used the identical name +// `s_gxyz_offset`, so this kernel silently compiled against gvhf-rys's +// declaration and the dynamic linker then bound it to whichever library loaded +// first -- the exact cross-library aliasing being fixed. Declare libpbc's own +// symbol explicitly so this TU can never pick up the neighbouring library's. +extern SYCL_EXTERNAL sycl_device_global s_pbc_gxyz_offset; +#endif + #define GOUT_WIDTH1 81 __device__ static @@ -38,10 +50,19 @@ void _fill_sr_vj_tasks(int &ntasks, int &pair_kl0, int64_t *bas_kl_idx, float *q_cond_ij, float *q_cond_kl, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, float dm_penalty, - JKMatrix& jmat, RysIntEnvVars& envs, BoundsInfo& bounds) + JKMatrix& jmat, RysIntEnvVars& envs, BoundsInfo& bounds, + double *shared_memory) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int thread_id = item.get_local_id(1) + item.get_local_range(1) * item.get_local_id(0); + int threads = item.get_local_range(1) * item.get_local_range(0); + int threadIdx_y = item.get_local_id(0); +#else int thread_id = threadIdx.x + blockDim.x * threadIdx.y; int threads = blockDim.x * blockDim.y; + int threadIdx_y = threadIdx.y; +#endif __syncthreads(); if (thread_id == 0) { ntasks = 0; @@ -93,7 +114,6 @@ void _fill_sr_vj_tasks(int &ntasks, int &pair_kl0, int64_t *bas_kl_idx, float omega2 = omega * omega; float theta_ij = omega2 * aij / (aij + omega2); - extern __shared__ double shared_memory[]; int *swap = (int *)shared_memory; while (pair_kl0 < pair_kl1 && ntasks < QUEUE_DEPTH - 512) { @@ -167,13 +187,14 @@ void _fill_sr_vj_tasks(int &ntasks, int &pair_kl0, int64_t *bas_kl_idx, } __syncthreads(); } - if (threadIdx.y == 0 && ntasks + thread_id < QUEUE_DEPTH && ntasks > 0) { + if (threadIdx_y == 0 && ntasks + thread_id < QUEUE_DEPTH && ntasks > 0) { bas_kl_idx[ntasks+thread_id] = bas_kl_idx[ntasks-1]; } __syncthreads(); } // gout_pattern = ((li == 0) << 3) | ((lj == 0) << 2) | ((lk == 0) << 1) | (ll == 0); +template __global__ static void rys_j_kernel(RysIntEnvVars envs, JKMatrix jmat, BoundsInfo bounds, int64_t *pair_ij_mapping, int64_t *pair_kl_mapping, @@ -181,15 +202,62 @@ void rys_j_kernel(RysIntEnvVars envs, JKMatrix jmat, BoundsInfo bounds, int nimgs, int nimgs_uniq_pair, int nbas_cell0, int nao, float *q_cond_ij, float *q_cond_kl, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - float dm_penalty, int64_t *pool, int *head, - int gout_pattern, int reserved_shm_size) + float dm_penalty, + int64_t *pool, int *head, const GXYZOffset *p_gxyz_offsets, + int gout_pattern, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { // sq is short for shl_quartet + #ifdef USE_SYCL + int sq_id = item.get_local_id(1); + int nsq_per_block = item.get_local_range(1); + int gout_id = item.get_local_id(0); + int gout_stride = item.get_local_range(0); + int t_id = item.get_local_id(0) * item.get_local_range(1) + item.get_local_id(1); + int blockIdx_x = item.get_group(1); + + double *shared_memory = reinterpret_cast(shm_mem); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &cell_j = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + auto gxyz_offsets = s_pbc_gxyz_offset.get() + OFFSET; + #else int sq_id = threadIdx.x; int nsq_per_block = blockDim.x; int gout_id = threadIdx.y; int gout_stride = blockDim.y; int t_id = threadIdx.y * blockDim.x + threadIdx.x; + int blockIdx_x = blockIdx.x; + + extern __shared__ double shared_memory[]; + __shared__ int ntasks, pair_ij, pair_kl0; + __shared__ int cell_j, ish_cell0, jsh_cell0, i0, j0; + __shared__ double ri[3]; + __shared__ double rjri[3]; + __shared__ double aij_cache[2]; + __shared__ int expi; + __shared__ int expj; + + // c_gxyz_offset is a 256-entry __constant__; the launcher copies the + // chunk for this OFFSET into it before each launch, so no offset here. + const GXYZOffset *gxyz_offsets = p_gxyz_offsets; + #endif int li = bounds.li; int lj = bounds.lj; int lk = bounds.lk; @@ -199,7 +267,6 @@ void rys_j_kernel(RysIntEnvVars envs, JKMatrix jmat, BoundsInfo bounds, int stride_l = bounds.stride_l; int g_size = bounds.g_size; - extern __shared__ double shared_memory[]; double *rlrk = shared_memory + sq_id; double *Rpq = shared_memory + nsq_per_block * 3 + sq_id; double *gx = shared_memory + nsq_per_block * 6 + sq_id; @@ -231,8 +298,7 @@ void rys_j_kernel(RysIntEnvVars envs, JKMatrix jmat, BoundsInfo bounds, idx_l[t_id] = lex_xyz_address(ll, t_id) * stride_l * nsq_per_block; } - int64_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - __shared__ int ntasks, pair_ij, pair_kl0; + int64_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; while (1) { __syncthreads(); __syncthreads(); @@ -252,17 +318,11 @@ while (1) { _fill_sr_vj_tasks(ntasks, pair_kl0, bas_kl_idx, pair_ij, ish, jsh, pair_kl_mapping, supcell_shl, Ts_ij_lookup, nimgs, nbas_cell0, q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, - dm_penalty, jmat, envs, bounds); + dm_penalty, jmat, envs, bounds, shared_memory); if (ntasks == 0) { continue; } - __shared__ int cell_j, ish_cell0, jsh_cell0, i0, j0; - __shared__ double ri[3]; - __shared__ double rjri[3]; - __shared__ double aij_cache[2]; - __shared__ int expi; - __shared__ int expj; int *bas = envs.bas; double *env = envs.env; if (t_id == 0) { @@ -398,7 +458,7 @@ while (1) { if (task_id >= ntasks) { continue; } - GXYZOffset goff = c_gxyz_offset[gout_id]; + GXYZOffset goff = gxyz_offsets[gout_id]; int *addr_i = idx_i + goff.ioff*3; int *addr_j = idx_j + goff.joff*3; int *addr_k = idx_k + goff.koff*3; @@ -427,7 +487,7 @@ while (1) { __syncthreads(); if (task_id < ntasks) { - GXYZOffset goff = c_gxyz_offset[gout_id]; + GXYZOffset goff = gxyz_offsets[gout_id]; int ioff = goff.ioff; int joff = goff.joff; int koff = goff.koff; @@ -461,7 +521,7 @@ while (1) { } } -extern void RYS_make_gxyz_offset(GXYZOffset *gxyz_offset, BoundsInfo &bounds); +extern GXYZOffset *PBC_make_gxyz_offset(GXYZOffset *goff, BoundsInfo &bounds); extern void threads_scheme_for_k(int *scheme, BoundsInfo &bounds, int shm_size, int gout_stride_max); @@ -520,8 +580,8 @@ int PBC_build_j(double *vj, double *dm, int n_dm, int nao, cudaMemset(head, 0, sizeof(int)*3); if (1) { - GXYZOffset gxyz_offset[256*3]; - RYS_make_gxyz_offset(gxyz_offset, bounds); + GXYZOffset gxyz_offset[625]; + GXYZOffset* p_gxyz_offset = PBC_make_gxyz_offset(gxyz_offset, bounds); int n_tiles = ntiles_i * ntiles_j * ntiles_k * ntiles_l; int gout_pattern = (((li == 0) << 3) | ((lj == 0) << 2) | @@ -529,15 +589,34 @@ int PBC_build_j(double *vj, double *dm, int n_dm, int nao, ( ll == 0)); auto launch = [&](auto offset, int tile_chunk) { - checkCudaErrors( - cudaMemcpyToSymbol(c_gxyz_offset, gxyz_offset+offset, - tile_chunk*sizeof(GXYZOffset), - 0, cudaMemcpyHostToDevice)); + constexpr int OFFSET = decltype(offset)::value; int scheme[4]; threads_scheme_for_k(scheme, bounds, shm_size, tile_chunk); int buflen = scheme[2]; + int reserved_shm_size = scheme[3]; + + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(scheme[1], scheme[0]); + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_j_kernel(dev_envs, jmat, bounds, pair_ij_mapping, pair_kl_mapping, + supcell_shl, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, + q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, + dm_penalty, pool, head + OFFSET/256, p_gxyz_offset, + gout_pattern, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + checkCudaErrors( + cudaMemcpyToSymbol(c_gxyz_offset, gxyz_offset+OFFSET, + tile_chunk*sizeof(GXYZOffset), + 0, cudaMemcpyHostToDevice)); if (buflen > 48000) { - cudaFuncSetAttribute(rys_j_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + cudaFuncSetAttribute(rys_j_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, @@ -546,18 +625,18 @@ int PBC_build_j(double *vj, double *dm, int n_dm, int nao, } } dim3 threads(scheme[0], scheme[1]); - int reserved_shm_size = scheme[3]; - rys_j_kernel<<>>( + rys_j_kernel<<>>( *envs, jmat, bounds, pair_ij_mapping, pair_kl_mapping, supcell_shl, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, - dm_penalty, pool, head + offset/256, + dm_penalty, pool, head + OFFSET/256, p_gxyz_offset, gout_pattern, reserved_shm_size); + #endif }; - launch(0, 256); - if (n_tiles > 256) launch(256, min(256, n_tiles-256)); - if (n_tiles > 512) launch(512, min(256, n_tiles-512)); + launch(std::integral_constant{}, 256); + if (n_tiles > 256) launch(std::integral_constant{}, min(256, n_tiles-256)); + if (n_tiles > 512) launch(std::integral_constant{}, min(256, n_tiles-512)); } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { diff --git a/gpu4pyscf/lib/pbc/rys_contract_jk_ip1.cu b/gpu4pyscf/lib/pbc/rys_contract_jk_ip1.cu index 29b4a3d18..26c9e7230 100644 --- a/gpu4pyscf/lib/pbc/rys_contract_jk_ip1.cu +++ b/gpu4pyscf/lib/pbc/rys_contract_jk_ip1.cu @@ -36,17 +36,40 @@ void rys_ejk_ip1_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, float *q_cond_ij, float *q_cond_kl, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, int64_t *pool, double *dd_pool, int *head, - int reserved_shm_size) + int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double *shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - int threads = blockDim.x * blockDim.y; - int thread_id = threadIdx.x + blockDim.x * threadIdx.y; - int64_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - int nf = bounds.nfi * bounds.nfj * bounds.nfk * bounds.nfl; - double *dd_cache = dd_pool + blockIdx.x * nf * blockDim.x + sq_id; +#ifdef USE_SYCL + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &cell_j = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int cell_j, ish_cell0, jsh_cell0, i0, j0; @@ -55,7 +78,13 @@ void rys_ejk_ip1_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, __shared__ double aij_cache[2]; __shared__ int expi; __shared__ int expj; - +#endif + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + int threads = blockDim_x * blockDim_y; + int thread_id = threadIdx_x + blockDim_x * threadIdx_y; int *bas = envs.bas; double *env = envs.env; int li = bounds.li; @@ -88,15 +117,19 @@ void rys_ejk_ip1_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bounds, double *gx = shared_memory + nsq_per_block * 6 + sq_id; double *rw = shared_memory + nsq_per_block * (g_size*3+6) + sq_id; double *cicj_cache = shared_memory + reserved_shm_size; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lj); - int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lk); - int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.ll); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lj); + const int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lk); + const int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.ll); int do_j = jk.j_factor != 0.; int do_k = jk.k_factor != 0.; int *ao_loc = envs.ao_loc; double *dm = jk.dm; + + int64_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; + int nf = bounds.nfi * bounds.nfj * bounds.nfk * bounds.nfl; + double *dd_cache = dd_pool + blockIdx_x * nf * blockDim_x + sq_id; while (1) { __syncthreads(); if (thread_id == 0) { @@ -459,17 +492,41 @@ void rys_ejk_strain_deriv_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bou float *q_cond_ij, float *q_cond_kl, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, int64_t *pool, double *dd_pool, int *head, - int reserved_shm_size) + int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, double *shared_memory + #endif + ) { - int sq_id = threadIdx.x; - int nsq_per_block = blockDim.x; - int gout_id = threadIdx.y; - int gout_stride = blockDim.y; - int threads = blockDim.x * blockDim.y; - int thread_id = threadIdx.x + blockDim.x * threadIdx.y; - int64_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; - int nf = bounds.nfi * bounds.nfj * bounds.nfk * bounds.nfl; - double *dd_cache = dd_pool + blockIdx.x * nf * blockDim.x + sq_id; +#ifdef USE_SYCL + int blockDim_x = item.get_local_range(1); + int blockDim_y = item.get_local_range(0); + int threadIdx_x = item.get_local_id(1); + int threadIdx_y = item.get_local_id(0); + int blockIdx_x = item.get_group(1); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &cell_j = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rj)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + int blockIdx_x = blockIdx.x; + extern __shared__ double shared_memory[]; __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int cell_j, ish_cell0, jsh_cell0, i0, j0; @@ -479,7 +536,13 @@ void rys_ejk_strain_deriv_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bou __shared__ double aij_cache[2]; __shared__ int expi; __shared__ int expj; - +#endif + int sq_id = threadIdx_x; + int nsq_per_block = blockDim_x; + int gout_id = threadIdx_y; + int gout_stride = blockDim_y; + int threads = blockDim_x * blockDim_y; + int thread_id = threadIdx_x + blockDim_x * threadIdx_y; int *bas = envs.bas; double *env = envs.env; int li = bounds.li; @@ -512,16 +575,19 @@ void rys_ejk_strain_deriv_kernel(RysIntEnvVars envs, JKEnergy jk, BoundsInfo bou double *gx = shared_memory + nsq_per_block * 6 + sq_id; double *rw = shared_memory + nsq_per_block * (g_size*3+6) + sq_id; double *cicj_cache = shared_memory + reserved_shm_size; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lj); - int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lk); - int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.ll); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lj); + const int *idx_k = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.lk); + const int *idx_l = _c_cartesian_lexical_xyz + lex_xyz_offset(bounds.ll); int do_j = jk.j_factor != 0.; int do_k = jk.k_factor != 0.; int *ao_loc = envs.ao_loc; double *dm = jk.dm; + int64_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; + int nf = bounds.nfi * bounds.nfj * bounds.nfk * bounds.nfl; + double *dd_cache = dd_pool + blockIdx_x * nf * blockDim_x + sq_id; double sigma_xx = 0; double sigma_xy = 0; double sigma_xz = 0; @@ -962,9 +1028,6 @@ while (1) { atomicAdd(sigma+8, sigma_zz); } -//extern int rys_ejk_ip1_unrolled(RysIntEnvVars *envs, JKEnergy *jk, BoundsInfo *bounds, -// int *pool, double *dd_pool); - extern "C" { int PBC_per_atom_jk_ip1(double *ejk, double j_factor, double k_factor, double *dm, int n_dm, int nao, @@ -1013,6 +1076,9 @@ int PBC_per_atom_jk_ip1(double *ejk, double j_factor, double k_factor, // Additional factor 1/2 from the two-electron Coulomb operator JKEnergy jk = {ejk, dm, 2.*j_factor, -k_factor, n_dm, omega, 0, 1}; + #ifdef USE_SYCL + sycl::queue& stream = *sycl_get_queue(); + #endif cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; @@ -1023,24 +1089,42 @@ int PBC_per_atom_jk_ip1(double *ejk, double j_factor, double k_factor, int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block; int reserved_shm_size = max(buflen, 6*gout_stride*quartets_per_block); - buflen = (reserved_shm_size + ij_prims)*sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_ejk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + buflen = reserved_shm_size + ij_prims; + #ifdef USE_SYCL + sycl::range<2> threads(gout_stride, quartets_per_block); + sycl::range<2> blocks(1, workers); + auto dev_envs = *envs; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_ejk_ip1_kernel(dev_envs, jk, bounds, pair_ij_mapping, pair_kl_mapping, + bas_mask_idx, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, + q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, + pool, dd_pool, head, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + int buflen_bytes = buflen * sizeof(double); + if (buflen_bytes > 48000) { + cudaFuncSetAttribute(rys_ejk_ip1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen_bytes); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen_bytes, cudaGetErrorString(err)); return 1; } } - rys_ejk_ip1_kernel<<>>( + rys_ejk_ip1_kernel<<>>( *envs, jk, bounds, pair_ij_mapping, pair_kl_mapping, bas_mask_idx, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, pool, dd_pool, head, reserved_shm_size); + #endif } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -1100,6 +1184,9 @@ int PBC_jk_strain_deriv(double *ejk, double j_factor, double k_factor, jk.lr_factor = 0; jk.sr_factor = 1; + #ifdef USE_SYCL + sycl::queue& stream = *sycl_get_queue(); + #endif cudaDeviceProp prop; cudaGetDeviceProperties(&prop, 0); int workers = prop.multiProcessorCount; @@ -1110,24 +1197,41 @@ int PBC_jk_strain_deriv(double *ejk, double j_factor, double k_factor, int quartets_per_block = scheme[0]; int gout_stride = scheme[1]; int ij_prims = iprim * jprim; - dim3 threads(quartets_per_block, gout_stride); int buflen = (nroots*2 + g_size*3 + 6) * quartets_per_block; int reserved_shm_size = max(buflen, 6*gout_stride*quartets_per_block); - buflen = (reserved_shm_size + ij_prims)*sizeof(double); - if (buflen > 48000) { - cudaFuncSetAttribute(rys_ejk_strain_deriv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + buflen = reserved_shm_size + ij_prims; + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(gout_stride, quartets_per_block); + auto dev_envs = *envs; + stream.submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_ejk_strain_deriv_kernel(dev_envs, jk, bounds, sigma, pair_ij_mapping, pair_kl_mapping, + bas_mask_idx, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, + q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, + pool, dd_pool, head, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + dim3 threads(quartets_per_block, gout_stride); + int buflen_bytes = buflen * sizeof(double); + if (buflen_bytes > 48000) { + cudaFuncSetAttribute(rys_ejk_strain_deriv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen_bytes); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { - fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, + fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen_bytes, cudaGetErrorString(err)); return 1; } } - rys_ejk_strain_deriv_kernel<<>>( + rys_ejk_strain_deriv_kernel<<>>( *envs, jk, bounds, sigma, pair_ij_mapping, pair_kl_mapping, bas_mask_idx, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, pool, dd_pool, head, reserved_shm_size); + #endif } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { diff --git a/gpu4pyscf/lib/pbc/rys_contract_k.cu b/gpu4pyscf/lib/pbc/rys_contract_k.cu index 0f058b9be..d9671f39d 100644 --- a/gpu4pyscf/lib/pbc/rys_contract_k.cu +++ b/gpu4pyscf/lib/pbc/rys_contract_k.cu @@ -28,9 +28,24 @@ #include "gvhf-rys/build_rys_gxyz.cuh" #include "pbc/create_tasks.cu" +#ifdef USE_SYCL +// NOTE: this device_global and PBC_make_gxyz_offset() below are deliberately +// named differently from their libgvhf_rys counterparts. Both libraries are +// loaded into the SAME process (verified via /proc/self/maps), neither lists +// the other in DT_NEEDED, and the symbols are GLOBAL DEFAULT visibility. When +// both exported `s_pbc_gxyz_offset` / `PBC_make_gxyz_offset`, the dynamic linker +// bound every caller in the process to whichever library happened to load +// first -- so one library's device_global was written by the other library's +// host memcpy, leaving its own copy uninitialised. Reading it yields garbage +// int8_t offsets that flow into load_dm() as an arbitrary pointer +// displacement. Keep these names library-unique. +SYCL_EXTERNAL sycl_device_global s_pbc_gxyz_offset; +#endif + #define GOUT_WIDTH1 81 // gout_pattern = ((li == 0) << 3) | ((lj == 0) << 2) | ((lk == 0) << 1) | (ll == 0); +template __global__ static void rys_k_kernel(RysIntEnvVars envs, JKMatrix kmat, BoundsInfo bounds, int64_t *pair_ij_mapping, int64_t *pair_kl_mapping, @@ -38,16 +53,53 @@ void rys_k_kernel(RysIntEnvVars envs, JKMatrix kmat, BoundsInfo bounds, int nimgs, int nimgs_uniq_pair, int nbas_cell0, int nao, float *q_cond_ij, float *q_cond_kl, float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, - float dm_penalty, int64_t *pool, int *head, - int gout_pattern, int reserved_shm_size) + float dm_penalty, + int64_t *pool, int *head_base, const GXYZOffset *p_gxyz_offsets, + int gout_pattern, int reserved_shm_size + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { // sq is short for shl_quartet + #ifdef USE_SYCL + int sq_id = item.get_local_id(1); + int nsq_per_block = item.get_local_range(1); + int gout_id = item.get_local_id(0); + int gout_stride = item.get_local_range(0); + int thread_id = item.get_local_id(1) + item.get_local_range(1) * item.get_local_id(0); + + int t_id = item.get_local_id(0) * item.get_local_range(1) + item.get_local_id(1); + int64_t *bas_kl_idx = pool + item.get_group(1) * QUEUE_DEPTH; + + double *shared_memory = reinterpret_cast(shm_mem); + + auto thread_block = item.get_group(); + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &cell_j = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &ish_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &jsh_cell0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&aij_cache)[2] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + auto gxyz_offsets = s_pbc_gxyz_offset.get() + OFFSET; + #else int sq_id = threadIdx.x; int nsq_per_block = blockDim.x; int gout_id = threadIdx.y; int gout_stride = blockDim.y; + int thread_id = threadIdx.x + blockDim.x * threadIdx.y; + int t_id = threadIdx.y * blockDim.x + threadIdx.x; int64_t *bas_kl_idx = pool + blockIdx.x * QUEUE_DEPTH; + extern __shared__ double shared_memory[]; __shared__ int ntasks, pair_ij, pair_kl0; __shared__ int cell_j, ish_cell0, jsh_cell0, i0, j0; @@ -57,6 +109,12 @@ void rys_k_kernel(RysIntEnvVars envs, JKMatrix kmat, BoundsInfo bounds, __shared__ int expi; __shared__ int expj; + // c_gxyz_offset is a 256-entry __constant__; the launcher copies the + // chunk for this OFFSET into it before each launch, so no offset here. + const GXYZOffset *gxyz_offsets = p_gxyz_offsets; + #endif + int *head = head_base + OFFSET/256; + int li = bounds.li; int lj = bounds.lj; int lk = bounds.lk; @@ -252,7 +310,7 @@ while (1) { if (task_id >= ntasks) { continue; } - GXYZOffset goff = c_gxyz_offset[gout_id]; + GXYZOffset goff = gxyz_offsets[gout_id]; int *addr_i = idx_i + goff.ioff*3; int *addr_j = idx_j + goff.joff*3; int *addr_k = idx_k + goff.koff*3; @@ -281,7 +339,7 @@ while (1) { __syncthreads(); if (task_id < ntasks) { - GXYZOffset goff = c_gxyz_offset[gout_id]; + GXYZOffset goff = gxyz_offsets[gout_id]; int ioff = goff.ioff; int joff = goff.joff; int koff = goff.koff; @@ -329,7 +387,9 @@ while (1) { // Requires room for 256*3 entries: the tile count prod((nf+2)/3) reaches 625 // at (4,4,4,4), and the launcher copies it out in chunks of 256. -void RYS_make_gxyz_offset(GXYZOffset *gxyz_offset, BoundsInfo &bounds) +// NOTE: kept as PBC_make_gxyz_offset (not renamed to RYS_make_gxyz_offset to +// match upstream) -- see the symbol-collision comment above this function. +GXYZOffset *PBC_make_gxyz_offset(GXYZOffset *goff, BoundsInfo &bounds) { /* nfi = (li + 1) * (li + 2) // 2 @@ -353,19 +413,29 @@ void RYS_make_gxyz_offset(GXYZOffset *gxyz_offset, BoundsInfo &bounds) for (int j = 0; j < nfj; j += 3) { for (int k = 0; k < nfk; k += 3) { for (int l = 0; l < nfl; l += 3) { - gxyz_offset[nf].ioff = i; - gxyz_offset[nf].joff = j; - gxyz_offset[nf].koff = k; - gxyz_offset[nf].loff = l; + goff[nf].ioff = i; + goff[nf].joff = j; + goff[nf].koff = k; + goff[nf].loff = l; ++nf; } } } } // n+m must be clamped too: nf need not divide 256, so the last round // would otherwise write past entry 255. for (int n = nf; n < 256; n += nf) { for (int m = 0; m < nf && n+m < 256; ++m) { - gxyz_offset[n+m] = gxyz_offset[m]; + goff[n+m] = goff[m]; } } + #ifdef USE_SYCL + sycl_get_queue()->memcpy(s_pbc_gxyz_offset, goff, max(nf, 256)*sizeof(GXYZOffset)).wait(); + return nullptr; + #else + // c_gxyz_offset holds only 256 entries; the launchers copy each 256-tile + // chunk (goff+OFFSET) into it right before the corresponding launch. + GXYZOffset *p_gxyz_offset; + cudaGetSymbolAddress((void**)&p_gxyz_offset, c_gxyz_offset); + return p_gxyz_offset; + #endif } void threads_scheme_for_k(int *scheme, BoundsInfo &bounds, @@ -496,8 +566,8 @@ int PBC_build_k(double *vk, double *dm, int n_dm, int nao, supcell_shl, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, dm_penalty, pool, head, workers)) { - GXYZOffset gxyz_offset[256*3]; - RYS_make_gxyz_offset(gxyz_offset, bounds); + GXYZOffset gxyz_offset[625]; + GXYZOffset* p_gxyz_offset = PBC_make_gxyz_offset(gxyz_offset, bounds); int n_tiles = ntiles_i * ntiles_j * ntiles_k * ntiles_l; int gout_pattern = (((li == 0) << 3) | ((lj == 0) << 2) | @@ -505,15 +575,34 @@ int PBC_build_k(double *vk, double *dm, int n_dm, int nao, ( ll == 0)); auto launch = [&](auto offset, int tile_chunk) { - checkCudaErrors( - cudaMemcpyToSymbol(c_gxyz_offset, gxyz_offset+offset, - tile_chunk*sizeof(GXYZOffset), - 0, cudaMemcpyHostToDevice)); + constexpr int OFFSET = decltype(offset)::value; int scheme[4]; threads_scheme_for_k(scheme, bounds, shm_size, tile_chunk); int buflen = scheme[2]; + int reserved_shm_size = scheme[3]; + + #ifdef USE_SYCL + sycl::range<2> blocks(1, workers); + sycl::range<2> threads(scheme[1], scheme[0]); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + rys_k_kernel(dev_envs, kmat, bounds, pair_ij_mapping, pair_kl_mapping, + supcell_shl, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, + q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, + dm_penalty, pool, head, p_gxyz_offset, + gout_pattern, reserved_shm_size, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else + checkCudaErrors( + cudaMemcpyToSymbol(c_gxyz_offset, gxyz_offset+OFFSET, + tile_chunk*sizeof(GXYZOffset), + 0, cudaMemcpyHostToDevice)); if (buflen > 48000) { - cudaFuncSetAttribute(rys_k_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); + cudaFuncSetAttribute(rys_k_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, buflen); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { fprintf(stderr, "Failed to set CUDA shm size %d: %s\n", buflen, @@ -522,19 +611,20 @@ int PBC_build_k(double *vk, double *dm, int n_dm, int nao, } } dim3 threads(scheme[0], scheme[1]); - int reserved_shm_size = scheme[3]; - rys_k_kernel<<>>( + rys_k_kernel<<>>( *envs, kmat, bounds, pair_ij_mapping, pair_kl_mapping, supcell_shl, Ts_ij_lookup, nimgs, nimgs_uniq_pair, nbas_cell0, nao, q_cond_ij, q_cond_kl, s_cond_ij, s_cond_kl, diffuse_exps, - dm_penalty, pool, head + offset/256, + dm_penalty, pool, head, p_gxyz_offset, gout_pattern, reserved_shm_size); + #endif }; - launch(0, 256); - if (n_tiles > 256) launch(256, min(256, n_tiles-256)); - if (n_tiles > 512) launch(512, min(256, n_tiles-512)); + launch(std::integral_constant{}, 256); + if (n_tiles > 256) launch(std::integral_constant{}, min(256, n_tiles-256)); + if (n_tiles > 512) launch(std::integral_constant{}, min(256, n_tiles-512)); } + cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { int device_id = -1; diff --git a/gpu4pyscf/lib/pbc/rys_roots_dat.cu b/gpu4pyscf/lib/pbc/rys_roots_dat.cu index 580b374b7..a0e933752 100644 --- a/gpu4pyscf/lib/pbc/rys_roots_dat.cu +++ b/gpu4pyscf/lib/pbc/rys_roots_dat.cu @@ -1,2 +1,4 @@ +#ifndef USE_SYCL // header only relevant in CUDA! #include "gvhf-rys/rys_constant.cu" +#endif #include "gvhf-rys/rys_roots_dat.cu" diff --git a/gpu4pyscf/lib/pbc/sorting.c b/gpu4pyscf/lib/pbc/sorting.c index e7f79fa7c..1df1243c1 100644 --- a/gpu4pyscf/lib/pbc/sorting.c +++ b/gpu4pyscf/lib/pbc/sorting.c @@ -2,6 +2,10 @@ #include #include +#ifdef __cplusplus +extern "C" { +#endif + void condense_primitive_ovlp_mask(int8_t *c_ovlp_mask, int8_t *p_ovlp_mask, int *p2c_mapping, int c_nbas, int p_nbas) { @@ -33,3 +37,7 @@ void take2d_add(double *out, double *inp, int *idx, int nrow, int ncol, int idxl } } } } + +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/gpu4pyscf/lib/pbc/supmol_sr_estimator.cu b/gpu4pyscf/lib/pbc/supmol_sr_estimator.cu index 69528cd76..40be6bcde 100644 --- a/gpu4pyscf/lib/pbc/supmol_sr_estimator.cu +++ b/gpu4pyscf/lib/pbc/supmol_sr_estimator.cu @@ -42,10 +42,21 @@ void fill_s_estimator(float *s_estimator, RysIntEnvVars envs, int64_t *bas_ij_idx, int *bas_mask_idx, float *atom_diffuse_exps, float *diffuse_exps, float *diffuse_ctr_coef, float log_cutoff, int nbas_cell0, int natm_cell0, uint32_t npairs, - double omega, int tril_symmetry, int8_t *Ecut_mask) + double omega, int tril_symmetry, int8_t *Ecut_mask + #ifdef USE_SYCL + , sycl::nd_item<1> &item, std::byte *shm_mem + #endif + ) { + #ifdef USE_SYCL + uint32_t sp_block_id = item.get_group(0); + int t_id = item.get_local_id(0); + float *shared_memory = reinterpret_cast(shm_mem); + #else uint32_t sp_block_id = blockIdx.x; int t_id = threadIdx.x; + extern __shared__ float shared_memory[]; + #endif int *atm = envs.atm; int *bas = envs.bas; double *env = envs.env; @@ -56,7 +67,6 @@ void fill_s_estimator(float *s_estimator, RysIntEnvVars envs, int jsh0 = bas_ij0 % NBAS_MAX; int li = bas[ish0*BAS_SLOTS+ANG_OF]; int lj = bas[jsh0*BAS_SLOTS+ANG_OF]; - extern __shared__ float shared_memory[]; float *xyz_cache = shared_memory; for (int k = t_id; k < natm_cell0; k += THREADS) { double *rk = env + atm[k*ATM_SLOTS+PTR_COORD]; @@ -168,11 +178,23 @@ void fill_s_estimator(float *s_estimator, RysIntEnvVars envs, __global__ static void q_cond_kernel(float *q_cond, RysIntEnvVars envs, int64_t *bas_ij_idx, int *gout_stride_lookup, - uint32_t npairs, double omega) + uint32_t npairs, double omega + #ifdef USE_SYCL + , sycl::nd_item<2> &item, std::byte *shm_mem + #endif + ) { + #ifdef USE_SYCL + uint32_t sp_block_id = item.get_group(1); + int threads = item.get_local_range(1); + int t_id = item.get_local_id(1); + float *shared_memory = reinterpret_cast(shm_mem); + #else uint32_t sp_block_id = blockIdx.x; int threads = blockDim.x; int t_id = threadIdx.x; + extern __shared__ float shared_memory[]; + #endif int *bas = envs.bas; double *env = envs.env; uint32_t shl_pair0 = sp_block_id * SP_BLOCK_SIZE; @@ -195,7 +217,13 @@ void q_cond_kernel(float *q_cond, RysIntEnvVars envs, int stride_k = stride_j * (lj + 1); int nfij = nfi * nfj; + #ifdef USE_SYCL + auto thread_block = item.get_group(); + int &gout_stride = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + int &nsp_per_block = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + #else __shared__ int gout_stride, nsp_per_block; + #endif if (t_id == 0) { gout_stride = gout_stride_lookup[li*LMAX1+lj]; nsp_per_block = THREADS / gout_stride; @@ -205,7 +233,6 @@ void q_cond_kernel(float *q_cond, RysIntEnvVars envs, int gout_id = t_id / nsp_per_block; int g_size = stride_k; - extern __shared__ float shared_memory[]; float *rjri = shared_memory + sp_id; float *Rpq = shared_memory + nsp_per_block * 3 + sp_id; float *rw = shared_memory + nsp_per_block * 6 + sp_id; @@ -214,8 +241,8 @@ void q_cond_kernel(float *q_cond, RysIntEnvVars envs, float *gx = shared_memory + nsp_per_block * (nroots * 2 + 6) + sp_id; // gz can be reused for gbuf; gbuf size = (li+1)*(lj+1)*(lij+1) float *gbuf = gx + g_size * nsp_per_block * 2; - int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); + const int *idx_i = _c_cartesian_lexical_xyz + lex_xyz_offset(li); + const int *idx_j = _c_cartesian_lexical_xyz + lex_xyz_offset(lj); for (uint32_t task_id = shl_pair0+sp_id; task_id < shl_pair1+sp_id; task_id += nsp_per_block) { float gout[GOUT_WIDTH]; @@ -461,9 +488,16 @@ __global__ static void sort_pair_ij_kernel(int64_t *pair_ij, int *ish, int *jsh, int nish, int njsh, int nbas, int tile) { + #ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int t_id = item.get_local_id(0); + int threads = item.get_local_range(0); + int i_tile = item.get_group(0); + #else int t_id = threadIdx.x; int threads = blockDim.x; int i_tile = blockIdx.x; + #endif size_t off = i_tile * tile * (size_t)njsh; // when nish not divisible by tile int nish_rem = min(tile, nish - i_tile * tile); @@ -494,6 +528,18 @@ int PBCfill_s_estimator(float *s_estimator, RysIntEnvVars *envs, { int sp_blocks = (npairs + SP_BLOCK_SIZE - 1) / SP_BLOCK_SIZE; int buflen = max(512, natm_cell0 * 3) * sizeof(float); + #ifdef USE_SYCL + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); + cgh.parallel_for(sycl::nd_range<1>(sp_blocks * THREADS, THREADS), [=](auto item) { + fill_s_estimator(s_estimator, dev_envs, bas_ij_idx, bas_mask_idx, atom_diffuse_exps, + diffuse_exps, diffuse_ctr_coef, log_cutoff, nbas_cell0, natm_cell0, + npairs, omega, tril_symmetry, Ecut_mask, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else fill_s_estimator<<>>( s_estimator, *envs, bas_ij_idx, bas_mask_idx, atom_diffuse_exps, diffuse_exps, diffuse_ctr_coef, log_cutoff, nbas_cell0, natm_cell0, @@ -505,6 +551,7 @@ int PBCfill_s_estimator(float *s_estimator, RysIntEnvVars *envs, cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -513,6 +560,20 @@ int PBCfill_qcond(float *q_cond, RysIntEnvVars *envs, int shm_size, uint32_t npairs, double omega) { int sp_blocks = (npairs + SP_BLOCK_SIZE - 1) / SP_BLOCK_SIZE; + #ifdef USE_SYCL + // Though the kernel is 1D launch in CUDA, SYCL must do 2D because of the + // free-functions used in rys_roots_for_k() method + sycl::range<2> threads(1, THREADS); + sycl::range<2> blocks(1, sp_blocks); + auto dev_envs = *envs; + sycl_get_queue()->submit([&](sycl::handler &cgh) { + sycl::local_accessor local_acc(sycl::range<1>(shm_size), cgh); + cgh.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + q_cond_kernel(q_cond, dev_envs, bas_ij_idx, gout_stride_lookup, npairs, omega, + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); + }); + }); + #else q_cond_kernel<<>>( q_cond, *envs, bas_ij_idx, gout_stride_lookup, npairs, omega); @@ -522,6 +583,7 @@ int PBCfill_qcond(float *q_cond, RysIntEnvVars *envs, int shm_size, cudaGetErrorString(err)); return 1; } + #endif return 0; } @@ -529,6 +591,11 @@ int PBCsort_pair_ij(int64_t *pair_ij, int *ish, int *jsh, int nish, int njsh, int nbas, int tile) { int ntile = (nish + tile - 1) / tile; + #ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(ntile * THREADS, THREADS), [=](auto item) { + sort_pair_ij_kernel(pair_ij, ish, jsh, nish, njsh, nbas, tile); + }); + #else sort_pair_ij_kernel<<>>(pair_ij, ish, jsh, nish, njsh, nbas, tile); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -536,6 +603,7 @@ int PBCsort_pair_ij(int64_t *pair_ij, int *ish, int *jsh, int nish, int njsh, cudaGetErrorString(err)); return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/pbc/unrolled_kernels.cuh b/gpu4pyscf/lib/pbc/unrolled_kernels.cuh index b74e88585..2638e7b45 100644 --- a/gpu4pyscf/lib/pbc/unrolled_kernels.cuh +++ b/gpu4pyscf/lib/pbc/unrolled_kernels.cuh @@ -16,6 +16,59 @@ #include "gvhf-rys/vhf.cuh" +#ifdef USE_SYCL + +#define dim3 sycl::range<2> + +#define JKMATRIX_KERNEL_ARGS \ +RysIntEnvVars envs, JKMatrix kmat, BoundsInfo bounds, \ + int64_t *pair_ij_mapping, int64_t *pair_kl_mapping, \ + int *supcell_shl, int *Ts_ij_lookup, \ + int nimgs, int nimgs_uniq_pair, int nbas_cell0, int nao, \ + float *q_cond_ij, float *q_cond_kl, \ + float *s_cond_ij, float *s_cond_kl, float *diffuse_exps, \ + float dm_penalty, int64_t *pool, int *head \ + , sycl::nd_item<2> &item, double *shared_memory + +#define JKMATRIX_KERNEL_SETUP() \ + int sq_id = item.get_local_id(1); \ + int gout_id = item.get_local_id(0); \ + int _nsq_per_block = item.get_local_range(1); \ + int blockIdx_x = item.get_group(1); \ + int64_t *bas_kl_idx = pool + blockIdx_x * QUEUE_DEPTH; \ + auto thread_block = item.get_group(); \ + int &ntasks = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &pair_ij = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &pair_kl0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &cell_j = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &ish_cell0= *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &jsh_cell0= *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &i0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &j0 = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + double (&ri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + double (&rjri)[3] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &expi = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); \ + int &expj = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + +#define LAUNCH_JKMATRIX_KERNEL(KERNEL) \ + { \ + auto dev_envs = *envs; auto dev_kmat = *kmat; auto dev_bounds = *bounds; \ + sycl::range<2> blocks(1, workers); \ + sycl::range<2> cuda_threads(gout_stride, nsq_per_block); \ + sycl_get_queue()->submit([&](sycl::handler &cgh) { \ + sycl::local_accessor local_acc(sycl::range<1>(buflen), cgh); \ + cgh.parallel_for(sycl::nd_range<2>(blocks * cuda_threads, cuda_threads), [=](auto item) { \ + KERNEL(dev_envs, dev_kmat, dev_bounds, \ + pair_ij_mapping, pair_kl_mapping, supcell_shl, Ts_ij_lookup, \ + nimgs, nimgs_uniq_pair, nbas_cell0, nao, q_cond_ij, q_cond_kl, \ + s_cond_ij, s_cond_kl, diffuse_exps, dm_penalty, pool, head, \ + item, GPU4PYSCF_IMPL_SYCL_GET_MULTI_PTR(local_acc)); \ + }); \ + }); \ + } + +#else // USE_SYCL + #define JKMATRIX_KERNEL_ARGS \ RysIntEnvVars envs, JKMatrix kmat, BoundsInfo bounds, \ int64_t *pair_ij_mapping, int64_t *pair_kl_mapping, \ @@ -44,3 +97,5 @@ pair_ij_mapping, pair_kl_mapping, supcell_shl, Ts_ij_lookup, \ nimgs, nimgs_uniq_pair, nbas_cell0, nao, q_cond_ij, q_cond_kl, \ s_cond_ij, s_cond_kl, diffuse_exps, dm_penalty, pool, head) + +#endif // USE_SYCL diff --git a/gpu4pyscf/lib/sem/CMakeLists.txt b/gpu4pyscf/lib/sem/CMakeLists.txt index cae010b13..4835e3f96 100644 --- a/gpu4pyscf/lib/sem/CMakeLists.txt +++ b/gpu4pyscf/lib/sem/CMakeLists.txt @@ -14,7 +14,7 @@ set(CMAKE_VERBOSE_MAKEFILE ON) -add_library(sem SHARED +set(GPU_SRCS ss_kernel.cu fock.cu eri_1c2e_kernel.cu @@ -22,7 +22,18 @@ add_library(sem SHARED grad_kernel.cu ) -set_target_properties(sem PROPERTIES - LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} - CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}" -) \ No newline at end of file +add_library(sem SHARED ${GPU_SRCS}) + +if (USE_SYCL) + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + + set_target_properties(sem PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) + target_compile_options(sem PRIVATE -x c++ -nocudainc -nocudalib) + target_link_libraries(sem PRIVATE sycl_compat gsycl) +else() + set_target_properties(sem PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} + CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}" + ) +endif() diff --git a/gpu4pyscf/lib/sem/eri_1c2e_kernel.cu b/gpu4pyscf/lib/sem/eri_1c2e_kernel.cu index fd975ed7a..fed47b91e 100644 --- a/gpu4pyscf/lib/sem/eri_1c2e_kernel.cu +++ b/gpu4pyscf/lib/sem/eri_1c2e_kernel.cu @@ -36,7 +36,12 @@ __global__ void rsc_kernel( const double* __restrict__ b_table, // Size 30*30 flattened double* __restrict__ out_val ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= n_tasks) return; int k = k_vec[idx]; @@ -125,6 +130,13 @@ int launch_rsc_kernel_c( int threads = 128; int blocks = (n_tasks + threads - 1) / threads; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + rsc_kernel(n_tasks, hartree2ev, k_vec, + na, ea, nb, eb, nc, ec, nd, ed, + fx_table, b_table, out_val); + }); +#else rsc_kernel<<>>( n_tasks, hartree2ev, k_vec, na, ea, nb, eb, nc, ec, nd, ed, @@ -135,7 +147,8 @@ int launch_rsc_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } -} // extern "C" \ No newline at end of file +} // extern "C" diff --git a/gpu4pyscf/lib/sem/eri_2c2e_kernel.cu b/gpu4pyscf/lib/sem/eri_2c2e_kernel.cu index 4f40cfada..1081b685f 100644 --- a/gpu4pyscf/lib/sem/eri_2c2e_kernel.cu +++ b/gpu4pyscf/lib/sem/eri_2c2e_kernel.cu @@ -775,8 +775,8 @@ __device__ double ccrep_pm6_device( enuclr += scale_vdw; // Short distance repulsion - double zi = pow(ele_i, 0.3333); //follow mopac the 1/3 is set to 0.3333 - double zj = pow(ele_j, 0.3333); //follow mopac the 1/3 is set to 0.3333 + double zi = pow((double)ele_i, 0.3333); //follow mopac the 1/3 is set to 0.3333 + double zj = pow((double)ele_j, 0.3333); //follow mopac the 1/3 is set to 0.3333 // double zi = cbrt((double)ele_i); // double zj = cbrt((double)ele_j); double ax = r_angstrom / (zi + zj); @@ -804,7 +804,12 @@ __global__ void multipole_eval_kernel( const double* __restrict__ add_vec, // (n_pairs,) double* __restrict__ out_vec // (n_pairs,) ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= n_pairs) return; out_vec[idx] = charg_kernel_device( @@ -826,7 +831,12 @@ __global__ void solve_poij_kernel( double* __restrict__ rho_vec, // (N,) Output const double hartree2ev // Constant passed from Python ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= n_atoms) return; int l = l_vec[idx]; @@ -913,7 +923,12 @@ __global__ void test_rijkl_kernel( const double* __restrict__ ch, double* __restrict__ out_val ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= n_tasks) return; out_val[idx] = rijkl_device( @@ -965,7 +980,16 @@ __global__ void calc_local_rep_core_kernel( double* __restrict__ core_out, // (n_pairs, 10, 2) double* __restrict__ gab_out // (n_pairs) ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int p_idx = item.get_group(0); + int threadIdx_x = item.get_local_id(0); + int blockDim_x = item.get_local_range(0); +#else int p_idx = blockIdx.x; + int threadIdx_x = threadIdx.x; + int blockDim_x = blockDim.x; +#endif if (p_idx >= n_pairs) return; int ni = pair_i_vec[p_idx]; // Atom index @@ -975,15 +999,23 @@ __global__ void calc_local_rep_core_kernel( int e_i = ele_id[ni]; // Element index int e_j = ele_id[nj]; // Element index +#ifdef USE_SYCL + auto thread_block = item.get_group(); + double (&s_ri)[22] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_rep)[491] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_core)[20] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double &s_gab = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else __shared__ double s_ri[22]; // sp parts __shared__ double s_rep[491]; // spd parts __shared__ double s_core[20]; // 10 rows, 2 cols __shared__ double s_gab; +#endif - int tid = threadIdx.x; + int tid = threadIdx_x; if (tid < 20) s_core[tid] = 0.0; - for (int t = tid; t < 491; t += blockDim.x) s_rep[t] = 0.0; + for (int t = tid; t < 491; t += blockDim_x) s_rep[t] = 0.0; __syncthreads(); // Thread 0 handles serial computation of prerequisite physical quantities @@ -1012,7 +1044,7 @@ __global__ void calc_local_rep_core_kernel( __syncthreads(); // parallel evaluation of 491 terms (Stage 1: Direct computation) - for (int t = tid; t < 491; t += blockDim.x) { + for (int t = tid; t < 491; t += blockDim_x) { int action = task_action[t]; bool valid_i = dorbs[ni] ? true : (task_li[t] == 0 ? true : (task_li[t] <= 1 && e_i >= 3)); @@ -1032,7 +1064,7 @@ __global__ void calc_local_rep_core_kernel( __syncthreads(); // parallel evaluation of 491 terms (Stage 2: Symmetry copying) - for (int t = tid; t < 491; t += blockDim.x) { + for (int t = tid; t < 491; t += blockDim_x) { int action = task_action[t]; if (action == 2) { s_rep[t] = s_rep[task_target[t]]; @@ -1043,7 +1075,7 @@ __global__ void calc_local_rep_core_kernel( __syncthreads(); // Flush the computed results into Global Memory at once - for (int t = tid; t < 491; t += blockDim.x) { + for (int t = tid; t < 491; t += blockDim_x) { rep_out[p_idx * 491 + t] = s_rep[t]; } if (tid < 20) { @@ -1056,7 +1088,13 @@ __global__ void calc_local_rep_core_kernel( // HARDCODED MAPPINGS // Dense 1D index (0..44) to 2D orbital index (i) +#ifdef USE_SYCL +// SYCL has no cross-TU __device__ linkage for read-only tables; +// `static constexpr` gives the same device-side constant data. +static constexpr int DENSE_TO_I[45] = { +#else __device__ const int DENSE_TO_I[45] = { +#endif 0, 1, 1, 2, 2, 2, @@ -1069,7 +1107,11 @@ __device__ const int DENSE_TO_I[45] = { }; // Dense 1D index (0..44) to 2D orbital index (j) +#ifdef USE_SYCL +static constexpr int DENSE_TO_J[45] = { +#else __device__ const int DENSE_TO_J[45] = { +#endif 0, 0, 1, 0, 1, 2, @@ -1082,7 +1124,11 @@ __device__ const int DENSE_TO_J[45] = { }; // It is indexd in the mopac +#ifdef USE_SYCL +static constexpr int MOPAC_INDEXD[9][9] = { +#else __device__ const int MOPAC_INDEXD[9][9] = { +#endif { 0, 1, 2, 3, 4, 5, 6, 7, 8}, { 1, 9, 10, 11, 12, 13, 14, 15, 16}, { 2, 10, 17, 18, 19, 20, 21, 22, 23}, @@ -1116,10 +1162,18 @@ __global__ void global_transform_kernel( double* __restrict__ e2a_out, double* __restrict__ enuc_out ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int p_idx = item.get_group(0); + int tid = item.get_local_id(0); + int blockDim_x = item.get_local_range(0); +#else int p_idx = blockIdx.x; + int tid = threadIdx.x; + int blockDim_x = blockDim.x; +#endif if (p_idx >= n_pairs) return; - int tid = threadIdx.x; int ni = pair_i_vec[p_idx]; int nj = pair_j_vec[p_idx]; int ele_i = ele_id[ni]; @@ -1130,10 +1184,18 @@ __global__ void global_transform_kernel( int limij = ii * (ii + 1) / 2; int limkl = kk * (kk + 1) / 2; +#ifdef USE_SYCL + auto thread_block = item.get_group(); + double (&s_R)[45][45] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_V)[45][45] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_L_A)[45] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_L_B)[45] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); +#else __shared__ double s_R[45][45]; __shared__ double s_V[45][45]; __shared__ double s_L_A[45]; __shared__ double s_L_B[45]; +#endif if (tid == 0) { double xi = coords[ni * 3 + 0], yi = coords[ni * 3 + 1], zi = coords[ni * 3 + 2]; @@ -1190,7 +1252,7 @@ __global__ void global_transform_kernel( // __syncthreads(); // Tensor Contraction 1 - for (int idx = tid; idx < limij * limkl; idx += blockDim.x) { + for (int idx = tid; idx < limij * limkl; idx += blockDim_x) { int ij = idx / limkl; int KL = idx % limkl; @@ -1220,7 +1282,7 @@ __global__ void global_transform_kernel( // Tensor Contraction 2 int kr = kr_offsets[p_idx]; - for (int idx = tid; idx < limij * limkl; idx += blockDim.x) { + for (int idx = tid; idx < limij * limkl; idx += blockDim_x) { int IJ = idx / limkl; int KL = idx % limkl; double w_val = 0.0; @@ -1232,13 +1294,13 @@ __global__ void global_transform_kernel( } // Transform Elenuc Integrals - for (int IJ = tid; IJ < limij; IJ += blockDim.x) { + for (int IJ = tid; IJ < limij; IJ += blockDim_x) { double h_val = 0.0; for (int ij = 0; ij < limij; ++ij) h_val += s_R[IJ][ij] * s_L_A[ij]; e1b_out[p_idx * 45 + IJ] = h_val; } - for (int KL = tid; KL < limkl; KL += blockDim.x) { + for (int KL = tid; KL < limkl; KL += blockDim_x) { double h_val = 0.0; for (int kl = 0; kl < limkl; ++kl) h_val += s_R[KL][kl] * s_L_B[kl]; e2a_out[p_idx * 45 + KL] = h_val; @@ -1261,6 +1323,11 @@ int launch_multipole_eval_kernel_c( ) { int threads_per_block = 128; int blocks_per_grid = (n_pairs + threads_per_block - 1) / threads_per_block; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks_per_grid * threads_per_block, threads_per_block), [=](auto item) [[intel::kernel_args_restrict]] { + multipole_eval_kernel(n_pairs, r_vec, l1_vec, l2_vec, m_vec, da_vec, db_vec, add_vec, out_vec); + }); +#else multipole_eval_kernel<<>>( n_pairs, r_vec, l1_vec, l2_vec, m_vec, da_vec, db_vec, add_vec, out_vec ); @@ -1269,6 +1336,7 @@ int launch_multipole_eval_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -1283,6 +1351,11 @@ int launch_solve_poij_kernel_c( int threads_per_block = 128; int blocks_per_grid = (n_atoms + threads_per_block - 1) / threads_per_block; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks_per_grid * threads_per_block, threads_per_block), [=](auto item) [[intel::kernel_args_restrict]] { + solve_poij_kernel(n_atoms, l_vec, d_vec, fg_vec, rho_vec, hartree2ev); + }); +#else solve_poij_kernel<<>>( n_atoms, l_vec, d_vec, fg_vec, rho_vec, hartree2ev ); @@ -1291,6 +1364,7 @@ int launch_solve_poij_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -1308,6 +1382,16 @@ int launch_test_rijkl_kernel_c( int threads = 128; int blocks = (n_tasks + threads - 1) / threads; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + test_rijkl_kernel(n_tasks, n_atom, + ni_vec, nj_vec, ij_vec, kl_vec, + li_vec, lj_vec, lk_vec, ll_vec, + ic_vec, r_vec, + po_tensor, ddp_tensor, core_rho, ch, + out_val); + }); +#else test_rijkl_kernel<<>>( n_tasks, n_atom, ni_vec, nj_vec, ij_vec, kl_vec, @@ -1321,6 +1405,7 @@ int launch_test_rijkl_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -1342,6 +1427,18 @@ int launch_calc_local_rep_core_kernel_c( int threads = 128; int blocks = n_pairs; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + calc_local_rep_core_kernel(n_pairs, pair_i_vec, pair_j_vec, ele_id, r_vec, n_atom, + am, ad, aq, dd, qq, + po_tensor, ddp_tensor, core_rho, ch, + tore, natorb, dorbs, + task_action, task_target, task_ij, task_kl, + task_li, task_lj, task_lk, task_ll, + HATREE2EV, + rep_out, core_out, gab_out); + }); +#else calc_local_rep_core_kernel<<>>( n_pairs, pair_i_vec, pair_j_vec, ele_id, r_vec, n_atom, am, ad, aq, dd, qq, @@ -1357,6 +1454,7 @@ int launch_calc_local_rep_core_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -1373,6 +1471,14 @@ int launch_global_transform_kernel_c( int threads = 128; int blocks = n_pairs; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + global_transform_kernel(n_pairs, pair_i_vec, pair_j_vec, ele_id, coords, + rep_in, core_in, gab_in, ind2_arr, natorb, kr_offsets, + tore, xfac, alpb, guess1, guess2, guess3, v_par6, BOHR, + w_out, e1b_out, e2a_out, enuc_out); + }); +#else global_transform_kernel<<>>( n_pairs, pair_i_vec, pair_j_vec, ele_id, coords, rep_in, core_in, gab_in, ind2_arr, natorb, kr_offsets, @@ -1384,7 +1490,8 @@ int launch_global_transform_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } -} // extern "C" \ No newline at end of file +} // extern "C" diff --git a/gpu4pyscf/lib/sem/fock.cu b/gpu4pyscf/lib/sem/fock.cu index 31c88ff39..c1b58de6d 100644 --- a/gpu4pyscf/lib/sem/fock.cu +++ b/gpu4pyscf/lib/sem/fock.cu @@ -38,7 +38,12 @@ void build_jk_2c2e_kernel( int nao) { // Each block processes one pair of interacting atoms (Atom A and Atom B) +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int p = item.get_group(0); +#else int p = blockIdx.x; +#endif if (p >= npairs) return; int A = pair_i[p]; @@ -59,6 +64,21 @@ void build_jk_2c2e_kernel( // Allocate shared memory. // In PM6, the maximum number of orbitals per atom is 9 (s, p, d). +#ifdef USE_SYCL + auto thread_block = item.get_group(); + double (&s_PAA)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_PBB)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_PAB)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_PBA)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + double (&s_JAA)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_JBB)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_KAB)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + double (&s_KBA)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(thread_block); + + int tid = item.get_local_id(0); + int bdim = item.get_local_range(0); +#else __shared__ double s_PAA[81]; __shared__ double s_PBB[81]; __shared__ double s_PAB[81]; @@ -71,6 +91,7 @@ void build_jk_2c2e_kernel( int tid = threadIdx.x; int bdim = blockDim.x; +#endif // Initialize shared memory to zero for (int i = tid; i < 81; i += bdim) { @@ -235,17 +256,32 @@ void build_jk_1c2e_kernel( int num_d_pairs) { // Grid handles 1 atom per block +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int A = item.get_group(0); + int threadIdx_x = item.get_local_id(0); + int blockDim_x = item.get_local_range(0); +#else int A = blockIdx.x; + int threadIdx_x = threadIdx.x; + int blockDim_x = blockDim.x; +#endif if (A >= natm) return; int offset = aoslice[A * 2]; int nao_A = natorb[A]; +#ifdef USE_SYCL + double (&s_P)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&s_J)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); + double (&s_K)[81] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else __shared__ double s_P[81]; __shared__ double s_J[81]; __shared__ double s_K[81]; +#endif - for (int i = threadIdx.x; i < 81; i += blockDim.x) { + for (int i = threadIdx_x; i < 81; i += blockDim_x) { s_J[i] = 0.0; s_K[i] = 0.0; int row = i / 9; @@ -259,7 +295,7 @@ void build_jk_1c2e_kernel( __syncthreads(); // Thread 0 handles the small number of s and p orbital integrals - if (threadIdx.x == 0) { + if (threadIdx_x == 0) { // s-orbital apply_eri_1c2e(0, 0, 0, 0, gss[A], s_P, s_J, s_K); @@ -286,7 +322,7 @@ void build_jk_1c2e_kernel( // All threads cooperatively handle d-orbital combinations if (nao_A == 9 && num_d_pairs > 0) { - for (int idx = threadIdx.x; idx < num_d_pairs; idx += blockDim.x) { + for (int idx = threadIdx_x; idx < num_d_pairs; idx += blockDim_x) { int IJ = intij[idx]; int KL = intkl[idx]; int rp = intrep[idx]; @@ -309,7 +345,7 @@ void build_jk_1c2e_kernel( __syncthreads(); - for (int i = threadIdx.x; i < 81; i += blockDim.x) { + for (int i = threadIdx_x; i < 81; i += blockDim_x) { int row = i / 9; int col = i % 9; if (row < nao_A && col < nao_A) { @@ -347,6 +383,15 @@ extern "C" { int blocks = npairs; int threads = 256; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + build_jk_2c2e_kernel(w_1d, P, J, K, + pair_i, pair_j, kr_offsets, + aoslice, natorb, loc_row, loc_col, + npairs, nao); + }); + sycl_get_queue()->wait(); +#else build_jk_2c2e_kernel<<>>( w_1d, P, J, K, pair_i, pair_j, kr_offsets, @@ -359,6 +404,7 @@ extern "C" { return 1; } cudaDeviceSynchronize(); +#endif return 0; } @@ -390,6 +436,16 @@ extern "C" { // 64 threads per block is sufficient since max d-orbital combinations is 243 int threads = 64; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + build_jk_1c2e_kernel(P, J, K, + gss, gsp, hsp, gpp, gp2, repd, + intij, intkl, intrep, + aoslice, natorb, loc_row, loc_col, + natm, nao, num_d_pairs); + }); + sycl_get_queue()->wait(); +#else build_jk_1c2e_kernel<<>>( P, J, K, gss, gsp, hsp, gpp, gp2, repd, @@ -403,6 +459,7 @@ extern "C" { return 1; } cudaDeviceSynchronize(); +#endif return 0; } -} \ No newline at end of file +} diff --git a/gpu4pyscf/lib/sem/grad_kernel.cu b/gpu4pyscf/lib/sem/grad_kernel.cu index 55fecfedd..e74fb6a8f 100644 --- a/gpu4pyscf/lib/sem/grad_kernel.cu +++ b/gpu4pyscf/lib/sem/grad_kernel.cu @@ -31,7 +31,12 @@ __global__ void calc_pair_e2e_kernel( double* __restrict__ E_2e_out, // (n_pairs,) int n_pairs ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int p = item.get_global_id(0); +#else int p = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (p >= n_pairs) return; int A = pair_i[p]; @@ -89,13 +94,20 @@ int launch_calc_pair_e2e_c( ) { int threads = 256; int blocks = (n_pairs + threads - 1) / threads; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) [[intel::kernel_args_restrict]] { + calc_pair_e2e_kernel(w_1d, P_AA, P_BB, P_AB, + pair_i, pair_j, natorb, kr_offsets, E_2e_out, n_pairs); + }); +#else calc_pair_e2e_kernel<<>>( w_1d, P_AA, P_BB, P_AB, pair_i, pair_j, natorb, kr_offsets, E_2e_out, n_pairs ); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) return 1; +#endif return 0; } -} \ No newline at end of file +} diff --git a/gpu4pyscf/lib/sem/ss_kernel.cu b/gpu4pyscf/lib/sem/ss_kernel.cu index 3118e0744..ed0379df2 100644 --- a/gpu4pyscf/lib/sem/ss_kernel.cu +++ b/gpu4pyscf/lib/sem/ss_kernel.cu @@ -29,7 +29,12 @@ __global__ void afn_kernel( const double* __restrict__ p_vec, double* __restrict__ af_out // Shape: (n_data, 20) ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int tid = item.get_global_id(0); +#else int tid = blockIdx.x * blockDim.x + threadIdx.x; +#endif int out_id = tid * 20; if (tid >= n_data) return; @@ -53,7 +58,12 @@ __global__ void bfn_kernel( const double* __restrict__ taylor_coeffs, // Flattened (13 * 16) transposed taylor coeffs double* __restrict__ bf_out // Shape: (n_data, 13) ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= n_data) return; double x_val = x[idx]; @@ -123,7 +133,12 @@ __global__ void rotation_transform_kernel( const double* __restrict__ C_tensor, // Input: (N, 3, 5, 5) double* __restrict__ di_out // Output: (N, 9, 9) ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int idx = item.get_global_id(0); +#else int idx = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (idx >= n_pairs) return; // ival[shell_idx][local_k_index] @@ -205,7 +220,12 @@ __global__ void ss_summation_kernel( const double* __restrict__ binom, double* __restrict__ out ) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + int tid = item.get_global_id(0); +#else int tid = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (tid >= n_pairs) return; @@ -276,6 +296,11 @@ int launch_ss_kernel_c( int threads_per_block = 128; int blocks_per_grid = (n_pairs + threads_per_block - 1) / threads_per_block; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks_per_grid * threads_per_block, threads_per_block), [=](auto item) [[intel::kernel_args_restrict]] { + ss_summation_kernel(n_pairs, ia, ib, ic, id, m, iab, af, bf, binom, out); + }); +#else ss_summation_kernel<<>>( n_pairs, ia, ib, ic, id, m, iab, af, bf, binom, out ); @@ -284,6 +309,7 @@ int launch_ss_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -294,6 +320,11 @@ int launch_afn_kernel_c( ) { int threads_per_block = 128; int blocks_per_grid = (n_pairs + threads_per_block - 1) / threads_per_block; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks_per_grid * threads_per_block, threads_per_block), [=](auto item) [[intel::kernel_args_restrict]] { + afn_kernel(n_pairs, p_vec, af_out); + }); +#else afn_kernel<<>>( n_pairs, p_vec, af_out ); @@ -302,6 +333,7 @@ int launch_afn_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -313,6 +345,11 @@ int launch_bfn_kernel_c( ) { int threads_per_block = 128; int blocks_per_grid = (n_pairs + threads_per_block - 1) / threads_per_block; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks_per_grid * threads_per_block, threads_per_block), [=](auto item) [[intel::kernel_args_restrict]] { + bfn_kernel(n_pairs, x, taylor_coeffs, bf_out); + }); +#else bfn_kernel<<>>( n_pairs, x, taylor_coeffs, bf_out ); @@ -321,6 +358,7 @@ int launch_bfn_kernel_c( if (err != cudaSuccess) { return 1; } +#endif return 0; } @@ -333,6 +371,11 @@ int launch_rotation_transform_kernel( { int threads_per_block = 128; int blocks_per_grid = (n_pairs + threads_per_block - 1) / threads_per_block; +#ifdef USE_SYCL + sycl_get_queue()->parallel_for(sycl::nd_range<1>(blocks_per_grid * threads_per_block, threads_per_block), [=](auto item) [[intel::kernel_args_restrict]] { + rotation_transform_kernel(n_pairs, S_local, C_tensor, di_out); + }); +#else rotation_transform_kernel<<>>( n_pairs, S_local, C_tensor, di_out ); @@ -341,8 +384,9 @@ int launch_rotation_transform_kernel( if (err != cudaSuccess) { return 1; } +#endif return 0; } -} // extern "C" \ No newline at end of file +} // extern "C" diff --git a/gpu4pyscf/lib/solvent/CMakeLists.txt b/gpu4pyscf/lib/solvent/CMakeLists.txt index 86a169d36..23544d932 100644 --- a/gpu4pyscf/lib/solvent/CMakeLists.txt +++ b/gpu4pyscf/lib/solvent/CMakeLists.txt @@ -19,13 +19,19 @@ set(CMAKE_Fortran_FLAGS "${CMAKE_Fortran_FLAGS} -std=legacy -fPIC") set(CMAKE_VERBOSE_MAKEFILE ON) -add_library(solvent SHARED -mnsol_interface.f90 -mnsol_mem.F -mnsol.F -pcm.cu + +set(FTN_SRCS + mnsol_interface.f90 + mnsol_mem.F + mnsol.F +) + +set(GPU_SRCS + pcm.cu ) +add_library(solvent SHARED ${FTN_SRCS} ${GPU_SRCS}) + #option(BUILD_SHARED_LIBS "build shared libraries" 1) #option(ENABLE_STATIC "Enforce static library build" 0) #if(ENABLE_STATIC) @@ -34,6 +40,17 @@ pcm.cu set_target_properties(solvent PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}) -set_target_properties(solvent PROPERTIES - CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") +if (USE_SYCL) + set_source_files_properties(${GPU_SRCS} PROPERTIES LANGUAGE CXX) + # Unlike the other GPU targets, `solvent` also carries Fortran sources, so + # the SYCL/C++-only flags must be guarded by a COMPILE_LANGUAGE genex + # instead of applied to the whole target. + target_compile_options(solvent PRIVATE + $<$:-x c++ -nocudainc -nocudalib> + ) + target_link_libraries(solvent PRIVATE sycl_compat gsycl) +else() + set_target_properties(solvent PROPERTIES + CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") +endif() diff --git a/gpu4pyscf/lib/solvent/pcm.cu b/gpu4pyscf/lib/solvent/pcm.cu index 2961cd6de..6af3f3e4b 100644 --- a/gpu4pyscf/lib/solvent/pcm.cu +++ b/gpu4pyscf/lib/solvent/pcm.cu @@ -28,8 +28,14 @@ static void _pcm_d_s(double* __restrict__ matrix_d, double* __restrict__ matrix_ const double* __restrict__ charge_exp, const double* __restrict__ switch_fun, const int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int i = static_cast( item.get_global_id(1) ); + int j = static_cast( item.get_global_id(0) ); +#else int i = blockIdx.x * blockDim.x + threadIdx.x; int j = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= n || j >= n){ return; } @@ -82,7 +88,16 @@ static void _pcm_left_multiply_S_offdiagonal(double* __restrict__ output, const const int n) { // Attention: The coords is assumed to be in x1,x2,...,xn,y1,y2,...,yn,z1,z2,...,zn, which is different from all other kernels! +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = static_cast( item.get_global_id(0) ); + const int threadIdx_x = item.get_local_id(1); + const int blockDim_x = item.get_local_range(1); +#else const int i = blockIdx.y * blockDim.y + threadIdx.y; + const int threadIdx_x = threadIdx.x; + const int blockDim_x = blockDim.x; +#endif if (i >= n) { return; } @@ -93,7 +108,7 @@ static void _pcm_left_multiply_S_offdiagonal(double* __restrict__ output, const const double ei = charge_exp[i]; double sum_i = 0; - for (int j = threadIdx.x; j < n; j += blockDim.x) { + for (int j = threadIdx_x; j < n; j += blockDim_x) { // calculate xi const double ej = charge_exp[j]; const double xi_ij = ei * ej * rsqrt(ei*ei + ej*ej); @@ -118,7 +133,7 @@ static void _pcm_left_multiply_S_offdiagonal(double* __restrict__ output, const sum_i += __shfl_down_sync(mask, sum_i, offset); } - if (threadIdx.x == 0) { + if (threadIdx_x == 0) { output[i] = sum_i; } } @@ -128,7 +143,12 @@ static void _pcm_left_multiply_S_diagonal(double* __restrict__ output, const dou const double* __restrict__ S_diag, const int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<1>(); + const int i = static_cast( item.get_global_id(0) ); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; +#endif if (i >= n) { return; } @@ -142,7 +162,22 @@ static void _pcm_left_multiply_D(double* __restrict__ output, const double* __re const double* __restrict__ coords, const double* __restrict__ norm_vec, const double* __restrict__ r_vdw, const double* __restrict__ charge_exp, const int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = static_cast( item.get_global_id(1) ); + const int threadIdx_x = item.get_local_id(1); + const int threadIdx_y = item.get_local_id(0); + const int blockDim_y = item.get_local_range(0); + + double (&sum_shared)[THREADS * THREADS] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int threadIdx_x = threadIdx.x; + const int threadIdx_y = threadIdx.y; + const int blockDim_y = blockDim.y; + + __shared__ double sum_shared[THREADS * THREADS]; +#endif if (i >= n) { return; } @@ -158,7 +193,7 @@ static void _pcm_left_multiply_D(double* __restrict__ output, const double* __re double nzi = 0.0; if constexpr (transpose) nzi = norm_vec[3*i+2]; double sum_i = 0.0; - for (int j = threadIdx.y; j < n; j += blockDim.y) { + for (int j = threadIdx_y; j < n; j += blockDim_y) { // calculate xi const double ej = charge_exp[j]; const double xi_ij = ei * ej * rsqrt(ei*ei + ej*ej); @@ -191,18 +226,16 @@ static void _pcm_left_multiply_D(double* __restrict__ output, const double* __re sum_i += d * right_vector[j]; } - __shared__ double sum_shared[THREADS * THREADS]; - - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_i; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_i; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[i] = sum_shared[threadIdx_x]; } } @@ -212,8 +245,14 @@ static void _pcm_dD_dS(double* __restrict__ matrix_dd, double* __restrict__ matr const double* __restrict__ charge_exp, const int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + int i = static_cast( item.get_global_id(1) ); + int j = static_cast( item.get_global_id(0) ); +#else int i = blockIdx.x * blockDim.x + threadIdx.x; int j = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= n || j >= n){ return; } @@ -267,7 +306,22 @@ static void _pcm_left_multiply_dS(double* __restrict__ output, const double* __r const double* __restrict__ coords, const double* __restrict__ charge_exp, const int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = static_cast( item.get_global_id(1) ); + const int threadIdx_x = item.get_local_id(1); + const int threadIdx_y = item.get_local_id(0); + const int blockDim_y = item.get_local_range(0); + + double (&sum_shared)[THREADS * THREADS] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int threadIdx_x = threadIdx.x; + const int threadIdx_y = threadIdx.y; + const int blockDim_y = blockDim.y; + + __shared__ double sum_shared[THREADS * THREADS]; +#endif if (i >= n) { return; } @@ -280,7 +334,7 @@ static void _pcm_left_multiply_dS(double* __restrict__ output, const double* __r double sum_x = 0.0; double sum_y = 0.0; double sum_z = 0.0; - for (int j = threadIdx.y; j < n; j += blockDim.y) { + for (int j = threadIdx_y; j < n; j += blockDim_y) { // calculate xi const double ej = charge_exp[j]; const double xi_ij = ei * ej * rsqrt(ei*ei + ej*ej); @@ -315,42 +369,40 @@ static void _pcm_left_multiply_dS(double* __restrict__ output, const double* __r sum_z += dSz * right_vector_j; } - __shared__ double sum_shared[THREADS * THREADS]; - - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_x; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_x; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[ i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[ i] = sum_shared[threadIdx_x]; } - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_y; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_y; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[n + i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[n + i] = sum_shared[threadIdx_x]; } - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_z; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_z; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[n * 2 + i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[n * 2 + i] = sum_shared[threadIdx_x]; } } @@ -359,7 +411,22 @@ static void _pcm_left_multiply_dS_one_atom(double* __restrict__ output, const do const double* __restrict__ coords, const double* __restrict__ charge_exp, const int n, const int g0, const int g1) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = static_cast( item.get_global_id(1) ); + const int threadIdx_x = item.get_local_id(1); + const int threadIdx_y = item.get_local_id(0); + const int blockDim_y = item.get_local_range(0); + + double (&sum_shared)[THREADS * THREADS] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int threadIdx_x = threadIdx.x; + const int threadIdx_y = threadIdx.y; + const int blockDim_y = blockDim.y; + + __shared__ double sum_shared[THREADS * THREADS]; +#endif if (i >= n) { return; } @@ -372,7 +439,7 @@ static void _pcm_left_multiply_dS_one_atom(double* __restrict__ output, const do double sum_x = 0.0; double sum_y = 0.0; double sum_z = 0.0; - for (int j = threadIdx.y + g0; j < g1; j += blockDim.y) { + for (int j = threadIdx_y + g0; j < g1; j += blockDim_y) { // calculate xi const double ej = charge_exp[j]; const double xi_ij = ei * ej * rsqrt(ei*ei + ej*ej); @@ -407,42 +474,40 @@ static void _pcm_left_multiply_dS_one_atom(double* __restrict__ output, const do sum_z += dSz * right_vector_j; } - __shared__ double sum_shared[THREADS * THREADS]; - - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_x; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_x; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[ i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[ i] = sum_shared[threadIdx_x]; } - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_y; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_y; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[n + i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[n + i] = sum_shared[threadIdx_x]; } - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_z; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_z; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[n * 2 + i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[n * 2 + i] = sum_shared[threadIdx_x]; } } @@ -452,7 +517,22 @@ static void _pcm_left_multiply_dD(double* __restrict__ output, const double* __r const double* __restrict__ coords, const double* __restrict__ charge_exp, const double* __restrict__ norm_vec, const int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = static_cast( item.get_global_id(1) ); + const int threadIdx_x = item.get_local_id(1); + const int threadIdx_y = item.get_local_id(0); + const int blockDim_y = item.get_local_range(0); + + double (&sum_shared)[THREADS * THREADS] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int threadIdx_x = threadIdx.x; + const int threadIdx_y = threadIdx.y; + const int blockDim_y = blockDim.y; + + __shared__ double sum_shared[THREADS * THREADS]; +#endif if (i >= n) { return; } @@ -469,7 +549,7 @@ static void _pcm_left_multiply_dD(double* __restrict__ output, const double* __r double sum_x = 0.0; double sum_y = 0.0; double sum_z = 0.0; - for (int j = threadIdx.y; j < n; j += blockDim.y) { + for (int j = threadIdx_y; j < n; j += blockDim_y) { // calculate xi const double ej = charge_exp[j]; const double xi_ij = ei * ej * rsqrt(ei*ei + ej*ej); @@ -512,42 +592,40 @@ static void _pcm_left_multiply_dD(double* __restrict__ output, const double* __r sum_z += dDz * right_vector_j; } - __shared__ double sum_shared[THREADS * THREADS]; - - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_x; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_x; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[ i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[ i] = sum_shared[threadIdx_x]; } - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_y; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_y; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[n + i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[n + i] = sum_shared[threadIdx_x]; } - sum_shared[threadIdx.y * THREADS + threadIdx.x] = sum_z; + sum_shared[threadIdx_y * THREADS + threadIdx_x] = sum_z; __syncthreads(); for (int stride = THREADS / 2; stride > 0; stride >>= 1) { - if (threadIdx.y < stride) { - sum_shared[threadIdx.y * THREADS + threadIdx.x] += sum_shared[(threadIdx.y + stride) * THREADS + threadIdx.x]; + if (threadIdx_y < stride) { + sum_shared[threadIdx_y * THREADS + threadIdx_x] += sum_shared[(threadIdx_y + stride) * THREADS + threadIdx_x]; } __syncthreads(); } - if (threadIdx.y == 0) { - output[n * 2 + i] = sum_shared[threadIdx.x]; + if (threadIdx_y == 0) { + output[n * 2 + i] = sum_shared[threadIdx_x]; } } @@ -557,8 +635,14 @@ static void _pcm_d2D_d2S(double* __restrict__ matrix_d2D, double* __restrict__ m const double* __restrict__ charge_exp, const int n) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i = static_cast( item.get_global_id(1) ); + const int j = static_cast( item.get_global_id(0) ); +#else const int i = blockIdx.x * blockDim.x + threadIdx.x; const int j = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i >= n || j >= n) { return; } @@ -633,15 +717,29 @@ static void _pcm_contract_d2S_offdiagonal(double* __restrict__ output, const double* __restrict__ coords, const double* __restrict__ charge_exp, const int ngrids, const int natm) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i_atom = static_cast( item.get_group(1) ); + const int j_atom = static_cast( item.get_group(0) ); + const int threadIdx_x = item.get_local_id(1); + const int threadIdx_y = item.get_local_id(0); + + double (&sum_shared)[n_thread_per_block * n_thread_per_block] = *sycl::ext::oneapi::group_local_memory_for_overwrite(item.get_group()); +#else const int i_atom = blockIdx.x; const int j_atom = blockIdx.y; + const int threadIdx_x = threadIdx.x; + const int threadIdx_y = threadIdx.y; + + __shared__ double sum_shared[n_thread_per_block * n_thread_per_block]; +#endif const int i_grid_start = gridslice[i_atom * 2 + 0]; const int i_grid_end = gridslice[i_atom * 2 + 1]; const int j_grid_start = gridslice[j_atom * 2 + 0]; const int j_grid_end = gridslice[j_atom * 2 + 1]; double sandwiched_d2S[9] { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 }; - for (int i_grid = i_grid_start + threadIdx.x; i_grid < i_grid_end; i_grid += n_thread_per_block) { + for (int i_grid = i_grid_start + threadIdx_x; i_grid < i_grid_end; i_grid += n_thread_per_block) { const double ei = charge_exp[i_grid]; const double rix = coords[3*i_grid]; @@ -650,7 +748,7 @@ static void _pcm_contract_d2S_offdiagonal(double* __restrict__ output, const double left_i = left_vector[i_grid]; - for (int j_grid = j_grid_start + threadIdx.y; j_grid < j_grid_end; j_grid += n_thread_per_block) { + for (int j_grid = j_grid_start + threadIdx_y; j_grid < j_grid_end; j_grid += n_thread_per_block) { const double ej = charge_exp[j_grid]; const double eij = ei * ej * rsqrt(ei*ei + ej*ej); @@ -694,8 +792,7 @@ static void _pcm_contract_d2S_offdiagonal(double* __restrict__ output, } } - __shared__ double sum_shared[n_thread_per_block * n_thread_per_block]; - const int tid = threadIdx.y * n_thread_per_block + threadIdx.x; + const int tid = threadIdx_y * n_thread_per_block + threadIdx_x; for (int i_xyz = 0; i_xyz < 9; i_xyz++) { __syncthreads(); @@ -718,8 +815,14 @@ __global__ static void _pcm_d2F_to_d2Sii(const double* __restrict__ F, const double* __restrict__ dF, const double* __restrict__ d2F, const double* __restrict__ charge_exp, double* __restrict__ d2Sii, const int n_atom, const int n_grid) { +#ifdef USE_SYCL + auto item = syclex::this_work_item::get_nd_item<2>(); + const int i_grid = static_cast( item.get_global_id(1) ); + const int ij_atom = static_cast( item.get_global_id(0) ); +#else const int i_grid = blockIdx.x * blockDim.x + threadIdx.x; const int ij_atom = blockIdx.y * blockDim.y + threadIdx.y; +#endif if (i_grid >= n_grid || ij_atom >= n_atom * n_atom) { return; } @@ -769,6 +872,13 @@ int pcm_d_s(cudaStream_t stream, double *matrix_d, double *matrix_s, { int ntilex = (n + THREADS - 1) / THREADS; int ntiley = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(ntiley, ntilex); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_d_s(matrix_d, matrix_s, coords, norm_vec, r_vdw, charge_exp, switch_fun, n); + }); + #else dim3 threads(THREADS, THREADS); dim3 blocks(ntilex, ntiley); _pcm_d_s<<>>(matrix_d, matrix_s, coords, norm_vec, r_vdw, charge_exp, switch_fun, n); @@ -776,6 +886,7 @@ int pcm_d_s(cudaStream_t stream, double *matrix_d, double *matrix_s, if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -785,6 +896,13 @@ int pcm_left_multiply_s(const cudaStream_t stream, double *output, const double { { const int ntiley = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(ntiley, 1); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_S_offdiagonal(output, right_vector, coords, charge_exp, n); + }); + #else const dim3 threads(THREADS, THREADS); const dim3 blocks(1, ntiley); _pcm_left_multiply_S_offdiagonal<<>>(output, right_vector, coords, charge_exp, n); @@ -792,9 +910,17 @@ int pcm_left_multiply_s(const cudaStream_t stream, double *output, const double if (err != cudaSuccess) { return 1; } + #endif } { const int ntilex = (n + THREADS * THREADS - 1) / (THREADS * THREADS); + #ifdef USE_SYCL + sycl::range<1> threads(THREADS * THREADS); + sycl::range<1> blocks(ntilex); + stream.parallel_for(sycl::nd_range<1>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_S_diagonal(output, right_vector, S_diag, n); + }); + #else const dim3 threads(THREADS * THREADS); const dim3 blocks(ntilex); _pcm_left_multiply_S_diagonal<<>>(output, right_vector, S_diag, n); @@ -802,6 +928,7 @@ int pcm_left_multiply_s(const cudaStream_t stream, double *output, const double if (err != cudaSuccess) { return 1; } + #endif } return 0; } @@ -811,6 +938,20 @@ int pcm_left_multiply_d(const cudaStream_t stream, double *output, const double const int n, const bool transpose) { const int ntilex = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(1, ntilex); + if (transpose) { + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_D< true> (output, right_vector, coords, norm_vec, r_vdw, charge_exp, n); + }); + } + else { + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_D (output, right_vector, coords, norm_vec, r_vdw, charge_exp, n); + }); + } + #else const dim3 threads(THREADS, THREADS); const dim3 blocks(ntilex, 1); if (transpose) @@ -821,6 +962,7 @@ int pcm_left_multiply_d(const cudaStream_t stream, double *output, const double if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -831,6 +973,13 @@ int pcm_dd_ds(cudaStream_t stream, double *matrix_dD, double *matrix_dS, { int ntilex = (n + THREADS - 1) / THREADS; int ntiley = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(ntiley, ntilex); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_dD_dS(matrix_dD, matrix_dS, coords, norm_vec, charge_exp, n); + }); + #else dim3 threads(THREADS, THREADS); dim3 blocks(ntilex, ntiley); _pcm_dD_dS<<>>(matrix_dD, matrix_dS, coords, norm_vec, charge_exp, n); @@ -838,6 +987,7 @@ int pcm_dd_ds(cudaStream_t stream, double *matrix_dD, double *matrix_dS, if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -846,6 +996,13 @@ int pcm_left_multiply_ds(const cudaStream_t stream, double *output, const double const int n, const bool transpose) { const int ntilex = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(1, ntilex); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_dS(output, right_vector, coords, charge_exp, n); + }); + #else const dim3 threads(THREADS, THREADS); const dim3 blocks(ntilex, 1); _pcm_left_multiply_dS<<>>(output, right_vector, coords, charge_exp, n); @@ -853,6 +1010,7 @@ int pcm_left_multiply_ds(const cudaStream_t stream, double *output, const double if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -864,6 +1022,13 @@ int pcm_left_multiply_ds_one_atom(const cudaStream_t stream, double *output, con return 1; } const int ntilex = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + const sycl::range<2> threads(THREADS, THREADS); + const sycl::range<2> blocks(1, ntilex); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_dS_one_atom(output, right_vector, coords, charge_exp, n, g0, g1); + }); + #else const dim3 threads(THREADS, THREADS); const dim3 blocks(ntilex, 1); _pcm_left_multiply_dS_one_atom<<>>(output, right_vector, coords, charge_exp, n, g0, g1); @@ -871,6 +1036,7 @@ int pcm_left_multiply_ds_one_atom(const cudaStream_t stream, double *output, con if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -879,6 +1045,20 @@ int pcm_left_multiply_dd(const cudaStream_t stream, double *output, const double const int n, const bool transpose) { const int ntilex = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + sycl::range<2> threads(THREADS, THREADS); + sycl::range<2> blocks(1, ntilex); + if (transpose) { + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_dD< true> (output, right_vector, coords, charge_exp, norm_vec, n); + }); + } + else { + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_left_multiply_dD (output, right_vector, coords, charge_exp, norm_vec, n); + }); + } + #else const dim3 threads(THREADS, THREADS); const dim3 blocks(ntilex, 1); if (transpose) @@ -889,6 +1069,7 @@ int pcm_left_multiply_dd(const cudaStream_t stream, double *output, const double if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -899,6 +1080,13 @@ int pcm_d2d_d2s(cudaStream_t stream, double *matrix_d2D, double *matrix_d2S, { const int ntilex = (n + THREADS - 1) / THREADS; const int ntiley = (n + THREADS - 1) / THREADS; + #ifdef USE_SYCL + const sycl::range<2> threads(THREADS, THREADS); + const sycl::range<2> blocks(ntiley, ntilex); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_d2D_d2S(matrix_d2D, matrix_d2S, coords, norm_vec, charge_exp, n); + }); + #else const dim3 threads(THREADS, THREADS); const dim3 blocks(ntilex, ntiley); _pcm_d2D_d2S<<>>(matrix_d2D, matrix_d2S, coords, norm_vec, charge_exp, n); @@ -906,6 +1094,7 @@ int pcm_d2d_d2s(cudaStream_t stream, double *matrix_d2D, double *matrix_d2S, if (err != cudaSuccess) { return 1; } + #endif return 0; } @@ -915,6 +1104,13 @@ int pcm_contract_d2s_offdiagonal(const cudaStream_t stream, double *output, const int ngrids, const int natm) { constexpr int n_thread_per_block = 16; // 32 will cause "too many resources requested for launch", out of register + #ifdef USE_SYCL + const sycl::range<2> threads(n_thread_per_block, n_thread_per_block); + const sycl::range<2> blocks(natm, natm); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_contract_d2S_offdiagonal(output, left_vector, right_vector, gridslice, coords, charge_exp, ngrids, natm); + }); + #else const dim3 threads(n_thread_per_block, n_thread_per_block); const dim3 blocks(natm, natm); _pcm_contract_d2S_offdiagonal <<>> @@ -925,6 +1121,7 @@ int pcm_contract_d2s_offdiagonal(const cudaStream_t stream, double *output, printf("pcm_contract_d2s_offdiagonal failed with error %d, error message: %s, ngrids = %d, natm = %d, n_thread_per_block = %d\n", err, cudaGetErrorString(err), ngrids, natm, n_thread_per_block); return 1; } + #endif return 0; } @@ -933,6 +1130,13 @@ int pcm_d2f_to_d2sii(cudaStream_t stream, const double* F, const double* dF, con { const int ntilex = (n_grid + THREADS - 1) / THREADS; const int ntiley = (n_atom * n_atom + THREADS - 1) / THREADS; + #ifdef USE_SYCL + const sycl::range<2> threads(THREADS, THREADS); + const sycl::range<2> blocks(ntiley, ntilex); + stream.parallel_for(sycl::nd_range<2>(blocks * threads, threads), [=](auto item) { + _pcm_d2F_to_d2Sii(F, dF, d2F, charge_exp, d2Sii, n_atom, n_grid); + }); + #else const dim3 threads(THREADS, THREADS); const dim3 blocks(ntilex, ntiley); _pcm_d2F_to_d2Sii<<>>(F, dF, d2F, charge_exp, d2Sii, n_atom, n_grid); @@ -940,6 +1144,7 @@ int pcm_d2f_to_d2sii(cudaStream_t stream, const double* F, const double* dF, con if (err != cudaSuccess) { return 1; } + #endif return 0; } } diff --git a/gpu4pyscf/lib/utils.py b/gpu4pyscf/lib/utils.py index c61708bc8..2675ecc80 100644 --- a/gpu4pyscf/lib/utils.py +++ b/gpu4pyscf/lib/utils.py @@ -110,41 +110,80 @@ def device(obj): def format_sys_info(): '''Format a list of system information for printing.''' import gpu4pyscf - from cupyx._runtime import get_runtime_info from gpu4pyscf.__config__ import num_devices, mem_fraction, props as device_props pyscf_info = lib.repo_info(pyscf.__file__) gpu4pyscf_info = lib.repo_info(os.path.join(__file__, '..', '..')) - cuda_version = cupy.cuda.runtime.runtimeGetVersion() - cuda_version = f"{cuda_version // 1000}.{(cuda_version % 1000) // 10}" - - runtime_info = get_runtime_info() - result = [ - f'System: {platform.uname()} Threads {lib.num_threads()}', - f'Python {sys.version}', - f'numpy {numpy.__version__} scipy {scipy.__version__} ' - f'h5py {h5py.__version__}', - f'Date: {time.ctime()}', - f'PySCF version {pyscf.__version__}', - f'PySCF path {pyscf_info["path"]}', - 'CUDA Environment', - f' CuPy {runtime_info.cupy_version}', - f' CUDA Path {runtime_info.cuda_path}', - f' CUDA Build Version {runtime_info.cuda_build_version}', - f' CUDA Driver Version {runtime_info.cuda_driver_version}', - f' CUDA Runtime Version {runtime_info.cuda_runtime_version}', - 'CUDA toolkit', - f' cuSolver {runtime_info.cusolver_version}', - f' cuBLAS {runtime_info.cublas_version}', - f' cuTENSOR {runtime_info.cutensor_version}', - 'Device info', - f' Device name {device_props["name"]}', - f' Device global memory {device_props["totalGlobalMem"] / 1024**3:.2f} GB', - f' CuPy memory fraction {mem_fraction}', - f' Num. Devices {num_devices}', - f'GPU4PySCF {gpu4pyscf.__version__}', - f'GPU4PySCF path {gpu4pyscf_info["path"]}' - ] + + from importlib.util import find_spec + has_dpctl = find_spec("dpctl") + if not has_dpctl: + from cupyx._runtime import get_runtime_info + cuda_version = cupy.cuda.runtime.runtimeGetVersion() + cuda_version = f"{cuda_version // 1000}.{(cuda_version % 1000) // 10}" + + runtime_info = get_runtime_info() + result = [ + f'System: {platform.uname()} Threads {lib.num_threads()}', + f'Python {sys.version}', + f'numpy {numpy.__version__} scipy {scipy.__version__} ' + f'h5py {h5py.__version__}', + f'Date: {time.ctime()}', + f'PySCF version {pyscf.__version__}', + f'PySCF path {pyscf_info["path"]}', + 'CUDA Environment', + f' CuPy {runtime_info.cupy_version}', + f' CUDA Path {runtime_info.cuda_path}', + f' CUDA Build Version {runtime_info.cuda_build_version}', + f' CUDA Driver Version {runtime_info.cuda_driver_version}', + f' CUDA Runtime Version {runtime_info.cuda_runtime_version}', + 'CUDA toolkit', + f' cuSolver {runtime_info.cusolver_version}', + f' cuBLAS {runtime_info.cublas_version}', + f' cuTENSOR {runtime_info.cutensor_version}', + 'Device info', + f' Device name {device_props["name"]}', + f' Device global memory {device_props["totalGlobalMem"] / 1024**3:.2f} GB', + f' CuPy memory fraction {mem_fraction}', + f' Num. Devices {num_devices}', + f'GPU4PySCF {gpu4pyscf.__version__}', + f'GPU4PySCF path {gpu4pyscf_info["path"]}' + ] + else: + import dpnp, dpctl + # DPCTL device info: pick default device or first GPU device + try: + device = dpctl.get_devices(device_type='gpu')[0] + except IndexError: + device = dpctl.get_default_device() # fallback to any device + + # Get device properties + dev_name = device.name + dev_driver_version = device.driver_version if hasattr(device, 'driver_version') else 'Unknown' +# dev_platform_version = device.platform.version if hasattr(device.platform, 'version') else 'Unknown' + dev_global_mem_bytes = device.global_mem_size + + result = [ + f'System: {platform.uname()} Threads {lib.num_threads()}', + f'Python {sys.version}', + f'numpy {numpy.__version__} scipy {scipy.__version__} ' + f'h5py {h5py.__version__}', + f'Date: {time.ctime()}', + f'PySCF version {pyscf.__version__}', + f'PySCF path {pyscf_info["path"]}', + 'SYCL / DPNP / DPCTL Environment', + f' Device name {dev_name}', + # f' Device platform version {dev_platform_version}', + f' Device driver version {dev_driver_version}', + f' Device max alloc size {dev_global_mem_bytes / 1024**3:.2f} GB', + f' DPNP version {dpnp.__version__}', + f' DPCTL version {dpctl.__version__}', + f' DPNP memory fraction {mem_fraction}', + f' Num. Devices {num_devices}', + f'GPU4PySCF {gpu4pyscf.__version__}', + f'GPU4PySCF path {gpu4pyscf_info["path"]}' + ] + if 'git' in pyscf_info: result.append(pyscf_info['git']) return result diff --git a/gpu4pyscf/nac/finite_diff.py b/gpu4pyscf/nac/finite_diff.py index 1802eed3f..a000e17c6 100644 --- a/gpu4pyscf/nac/finite_diff.py +++ b/gpu4pyscf/nac/finite_diff.py @@ -18,8 +18,85 @@ from gpu4pyscf.lib import logger from gpu4pyscf.tdscf import ris from scipy.optimize import linear_sum_assignment -from gpu4pyscf.grad.tests.test_tdrhf_grad import diagonalize_tda +def diagonalize(a, b, nroots=5): + nocc, nvir = a.shape[:2] + nov = nocc * nvir + a = a.reshape(nov, nov) + b = b.reshape(nov, nov) + h = np.block([[a, b], + [-b.conj(), -a.conj()]]) + e, xy = np.linalg.eig(np.asarray(h)) + assert np.max(np.abs(e.imag)) < 1e-14 + assert np.max(np.abs(xy.imag)) < 1e-14 + e = e.real + xy = xy.real + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + +def diagonalize_tda(a, nroots=5): + nocc, nvir = a.shape[:2] + nov = nocc * nvir + a = a.reshape(nov, nov) + e, xy = np.linalg.eigh(np.asarray(a)) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + +def diagonalize_u(a, b, nroots=5): + a_aa, a_ab, a_bb = a + b_aa, b_ab, b_bb = b + nocc_a, nvir_a, nocc_b, nvir_b = a_ab.shape + a_aa = a_aa.reshape((nocc_a * nvir_a, nocc_a * nvir_a)) + a_ab = a_ab.reshape((nocc_a * nvir_a, nocc_b * nvir_b)) + a_bb = a_bb.reshape((nocc_b * nvir_b, nocc_b * nvir_b)) + b_aa = b_aa.reshape((nocc_a * nvir_a, nocc_a * nvir_a)) + b_ab = b_ab.reshape((nocc_a * nvir_a, nocc_b * nvir_b)) + b_bb = b_bb.reshape((nocc_b * nvir_b, nocc_b * nvir_b)) + a = np.block([[a_aa, a_ab], [a_ab.T, a_bb]]) + b = np.block([[b_aa, b_ab], [b_ab.T, b_bb]]) + abba = np.asarray(np.block([[a, b], [-b.conj(), -a.conj()]])) + e, xy = np.linalg.eig(abba) + assert np.max(np.abs(e.imag)) < 1e-14 + assert np.max(np.abs(xy.imag)) < 1e-14 + e = e.real + xy = xy.real + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + +def diagonalize_tda_u(a, nroots=5): + a_aa, a_ab, a_bb = a + nocc_a, nvir_a, nocc_b, nvir_b = a_ab.shape + a_aa = a_aa.reshape((nocc_a * nvir_a, nocc_a * nvir_a)) + a_ab = a_ab.reshape((nocc_a * nvir_a, nocc_b * nvir_b)) + a_bb = a_bb.reshape((nocc_b * nvir_b, nocc_b * nvir_b)) + a = np.block([[a_aa, a_ab], [a_ab.T, a_bb]]) + e, xy = np.linalg.eigh(a) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] def change_sign(s12_ao, mo_coeff_b ,mo_coeff): mo_coeff_new = mo_coeff*1.0 @@ -50,16 +127,16 @@ def match_and_reorder_mos(s12_ao, mo_coeff_b, mo_coeff, threshold=0.4): below_threshold_mask = abs_mo_overlap < threshold infinity_cost = mo_coeff_b.shape[1] + 1 cost_matrix[below_threshold_mask] = infinity_cost - + row_ind, col_ind = linear_sum_assignment(cost_matrix.get()) matching_indices = col_ind - + mo2_reordered = mo_coeff[:, matching_indices] final_chosen_overlaps = abs_mo_overlap[row_ind, col_ind] invalid_matches_mask = final_chosen_overlaps < threshold - + if cp.any(invalid_matches_mask): num_invalid = cp.sum(invalid_matches_mask) print( @@ -137,7 +214,7 @@ def get_nacv_ge(td_nac, x_yI, delta=0.001, with_ris=False, singlet=True, atmlst= yI = cp.zeros_like(xI) yI = yI.reshape(nocc, nvir) - gamma = np.block([[np.zeros((nocc, nocc)), xI.get()], + gamma = np.block([[np.zeros((nocc, nocc)), xI.get()], [(xI.T*0.0).get(), np.zeros((nvir, nvir))]]) gamma = cp.asarray(gamma)*2 gamma_ao = mo_coeff @ gamma @ mo_coeff.T @@ -160,7 +237,7 @@ def get_nacv_ge(td_nac, x_yI, delta=0.001, with_ris=False, singlet=True, atmlst= offsetdic = mol.offset_nr_by_atom() s12_deriv = mol.intor('int1e_ipovlp') s12_deriv = cp.asarray(s12_deriv) - for k, ia in enumerate(atmlst): + for k, ia in enumerate(atmlst): shl0, shl1, p0, p1 = offsetdic[ia] s12_deriv_tmp = s12_deriv*1.0 ds1_tmp = s12_deriv_tmp.transpose(0,2,1) @@ -192,7 +269,7 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, nJ, delta=0.001, with_ris=False, singlet=Tru if not isinstance(yJ, np.ndarray) and not isinstance(yJ, cp.ndarray): yJ = cp.zeros_like(xJ) yJ = cp.asarray(yJ).reshape(nocc, nvir) - gamma = np.block([[(-xJ@xI.T).get(), np.zeros((nocc, nvir))], + gamma = np.block([[(-xJ@xI.T).get(), np.zeros((nocc, nvir))], [np.zeros((nvir, nocc)), (xI.T@xJ).get()]]) * 2 gamma = cp.asarray(gamma) gamma_ao = mo_coeff @ gamma @ mo_coeff.T @@ -213,20 +290,20 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, nJ, delta=0.001, with_ris=False, singlet=Tru sign1 = -1.0 if (xJ*xJ_minus).sum() < 0.0: sign2 = -1.0 - + mo_diff = (mf_add.mo_coeff - mf_minus.mo_coeff)/(delta*2.0)*0.52917721092 dpq = mo_coeff.T @ s @ mo_diff nac[iatm, icart] = (gamma*dpq).sum() t_diff = (xJ_add*sign1 - xJ_minus*sign2)/(delta*2.0)*0.52917721092 nac3[iatm, icart] = (xI*t_diff).sum()*2 # for double occupancy - + nac2 = np.zeros((natm, 3)) atmlst = range(mol.natm) offsetdic = mol.offset_nr_by_atom() s12_deriv = mol.intor('int1e_ipovlp') s12_deriv = cp.asarray(s12_deriv) - for k, ia in enumerate(atmlst): + for k, ia in enumerate(atmlst): shl0, shl1, p0, p1 = offsetdic[ia] s12_deriv_tmp = s12_deriv*1.0 ds1_tmp = s12_deriv_tmp.transpose(0,2,1) @@ -234,4 +311,3 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, nJ, delta=0.001, with_ris=False, singlet=Tru ds1_tmp[:,:,p1:] = 0 nac2[k] = cp.einsum('xij,ij->x', ds1_tmp, gamma_ao).get() return nac - nac2 + nac3 - \ No newline at end of file diff --git a/gpu4pyscf/pbc/df/ft_ao.py b/gpu4pyscf/pbc/df/ft_ao.py index 6c23e5a97..2f49f27c5 100644 --- a/gpu4pyscf/pbc/df/ft_ao.py +++ b/gpu4pyscf/pbc/df/ft_ao.py @@ -362,7 +362,7 @@ def ft_evaluator(self, batch_size=None, compressing=True, cart=None, dims, tmp = np.empty_like(dims), dims dims[cell.sorted_idx] = tmp ao_loc = cp.asarray(np.append(0, np.cumsum(dims.ravel()))) - ao_loc = np.append(ao_loc[cell.sorted_idx], nao) + ao_loc = cp.append(ao_loc[cell.sorted_idx], nao) ao_loc = cp.asarray(ao_loc, dtype=np.int32) if batch_size is None: diff --git a/gpu4pyscf/pbc/df/int3c2e.py b/gpu4pyscf/pbc/df/int3c2e.py index c1cc7d11d..e84df4645 100644 --- a/gpu4pyscf/pbc/df/int3c2e.py +++ b/gpu4pyscf/pbc/df/int3c2e.py @@ -128,7 +128,7 @@ def sr_aux_e2(cell, auxcell, omega, kpts=None, bvk_kmesh=None, j_only=False): out = cp.empty((nkpts,nkpts,naux,nao,nao), dtype=np.complex128) kk_conserv = double_translation_indices(int3c2e_opt.bvk_kmesh) for k in range(nkpts): - ki_idx, kj_idx = np.where(kk_conserv == k) + ki_idx, kj_idx = cp.where(kk_conserv == k) out[k] = _unpack_cderi_v2(j3c[k], pair_address, kj_idx, conj_mapping, expLk, nao, axis) j3c = None diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py b/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py index ffb4c26e9..05624b902 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py @@ -324,7 +324,7 @@ def test_contract_dm_kpts(): np.random.seed(9) auxvec = np.random.rand(auxcell.nao) vj = opt.contract_auxvec(opt.auxcell.apply_C_dot(auxvec), kpts=kpts) - ref = cp.einsum('kpqr,r->kpq', j3c, auxvec) + ref = cp.einsum('kpqr,r->kpq', j3c, cp.asarray(auxvec)) # auxvec is host data assert abs(vj - ref).max() < 1e-10 def test_int3c2e_batch_evaluation(): diff --git a/gpu4pyscf/pbc/dft/multigrid.py b/gpu4pyscf/pbc/dft/multigrid.py index cfc2f2fab..ce0d51e01 100644 --- a/gpu4pyscf/pbc/dft/multigrid.py +++ b/gpu4pyscf/pbc/dft/multigrid.py @@ -1211,7 +1211,7 @@ def sub_tasks_for_l(mesh, n_radius, mask): Ecut_threshold /= 2 continue - n_radius = int(np.ceil(r_active.max() / dh)) + n_radius = int(cp.ceil(r_active.max() / dh)) n_radius = max(n_radius, 4) sub_tasks = sub_tasks_for_l(mesh, n_radius, mask) tasks.append(sub_tasks) @@ -1226,7 +1226,7 @@ def sub_tasks_for_l(mesh, n_radius, mask): if cp.any(remaining_mask): # TODO: Using a regular FFTDF task than the MG algorithm? dh = (cell_len / mesh).min() - n_radius = int(np.ceil(radius[remaining_mask].max() / dh)) + n_radius = int(cp.ceil(radius[remaining_mask].max() / dh)) n_radius = max(n_radius, 4) sub_tasks = sub_tasks_for_l(mesh, n_radius, remaining_mask) tasks.append(sub_tasks) diff --git a/gpu4pyscf/pbc/dft/multigrid_v2.py b/gpu4pyscf/pbc/dft/multigrid_v2.py index b9b5d2df1..18b89d892 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/multigrid_v2.py @@ -1162,12 +1162,15 @@ def convert_xc_on_g_mesh_to_fock( xc_on_g_mesh = xc_on_g_mesh.reshape(n_channels, density_slices, *mydf.mesh) if kpts is None: - n_k_points = 1 - at_gamma_point = True - else: - assert kpts.ndim == 2 - n_k_points = len(kpts) - at_gamma_point = multigrid.gamma_point(kpts) + kpts = np.zeros((1,3)) + elif isinstance(kpts, KPoints): + kpts = kpts.kpts + is_single_kpt = kpts.ndim == 1 + if is_single_kpt: + kpts = kpts.reshape(1, 3) + assert kpts.ndim == 2 + n_k_points = len(kpts) + at_gamma_point = multigrid.gamma_point(kpts) if hermi != 1: raise NotImplementedError @@ -2030,6 +2033,9 @@ def nr_rks_fxc(self, cell, grids, xc_code, dm0, dms, hermi=0, fxc=None, elif isinstance(kpts, KPoints): kpts = kpts.kpts_ibz + is_single_kpt = kpts.ndim == 1 + if is_single_kpt: + kpts = kpts.reshape(1, 3) assert kpts.ndim == 2 assert dms.ndim == 4 nset, nkpts, nao = dms.shape[:3] @@ -2094,6 +2100,9 @@ def nr_uks_fxc(self, cell, grids, xc_code, dm0, dms, hermi=0, fxc=None, elif isinstance(kpts, KPoints): kpts = kpts.kpts_ibz + is_single_kpt = kpts.ndim == 1 + if is_single_kpt: + kpts = kpts.reshape(1, 3) assert kpts.ndim == 2 assert dms.ndim == 5 nset, nkpts, nao = dms.shape[1:4] diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py index 64b505f30..accb57fd5 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py @@ -197,6 +197,8 @@ def eval_nucG_SI_gradient(cell, mesh, rho_g): class KnownValues(unittest.TestCase): def test_get_pp(self): ref = MultiGridNumInt_cpu(cell_orth).get_pp() + if ref.ndim == 2: # In pyscf==2.8.0 + ref = ref[None,:,:] out = multigrid.MultiGridNumInt(cell_orth).get_pp().get() self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) @@ -204,12 +206,16 @@ def test_get_pp(self): def test_get_nuc(self): ref = MultiGridNumInt_cpu(cell_orth).get_nuc() out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get() + if ref.ndim == 2: # In pyscf==2.8.0 + ref = ref[None,:,:] self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc_nonorth(self): ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc() out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get() + if ref.ndim == 2: # In pyscf==2.8.0 + ref = ref[None,:,:] self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 7) diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py index 02f4bc0e4..e32a8e9dc 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py @@ -366,7 +366,7 @@ def test_uniform_grid_division_mode(self): test_coords = [] for frag_grids in grids.loop_grids(): test_coords.append(frag_grids.coords) - assert np.max(np.abs(frag_grids.weights - ref_weight)) < 1e-14 + assert cp.max(cp.abs(frag_grids.weights - ref_weight)) < 1e-14 test_coords = cp.vstack(test_coords).get() ref_coords = ref_coords[np.lexsort((ref_coords[:, 2], ref_coords[:, 1], ref_coords[:, 0])), :] diff --git a/gpu4pyscf/pbc/hessian/gamma_hessian.py b/gpu4pyscf/pbc/hessian/gamma_hessian.py index 8fc7824ad..6825f81e9 100644 --- a/gpu4pyscf/pbc/hessian/gamma_hessian.py +++ b/gpu4pyscf/pbc/hessian/gamma_hessian.py @@ -36,6 +36,7 @@ except ImportError: HAS_PHONOPY = False from gpu4pyscf.lib import logger +from gpu4pyscf.pbc.tools.discretization import freeze_mesh GRAD_TO_FORCE = -(HARTREE2EV / BOHR_TO_ANGSTROM) EV_A2_TO_HA_BOHR2 = BOHR_TO_ANGSTROM**2 / HARTREE2EV @@ -48,7 +49,9 @@ class GammaHessian(lib.StreamObject): ``primitive_matrix`` maps the input cell to the primitive cell and must be supplied explicitly. ``kernel`` stores the corresponding phonopy - object in ``phonon``. + object in ``phonon``. The current ``cell.mesh`` is fixed when ``kernel`` + starts. If geometry-optimization mesh metadata is available, that mesh is + reused. """ def __init__( @@ -93,6 +96,12 @@ def kernel(self): mf = self.mf cell = mf.cell original_coords = cell.atom_coords() # Bohr + geomopt_mesh = getattr( + mf, + "_geomopt_mesh", + getattr(cell, "_geomopt_mesh", None), + ) + reference_mesh = freeze_mesh(mf, cell, geomopt_mesh) # default calculator is vasp, default unit in phonopy for vasp: # | Distance Atomic mass Force Force constants @@ -133,6 +142,7 @@ def kernel(self): for index, displaced in enumerate(displaced_cells, 1): disp_cell = cell.set_geom_(displaced.positions, unit="Angstrom", inplace=False) mf_disp = mf.copy().reset(disp_cell) + freeze_mesh(mf_disp, disp_cell, reference_mesh) logger.info( mf, "Running displaced SCF and gradient %d/%d", diff --git a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py index cada001a1..c7c49b57b 100644 --- a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py +++ b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py @@ -40,7 +40,7 @@ def test_krhf_smearing(self): mf = cell.KRHF(kpts=cell.make_kpts([2,1,1])).to_gpu() mf = mf.smearing(0.1, 'fermi') nkpts = len(mf.kpts) - mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)]) + mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)]) mf.get_occ(mo_energy_kpts) self.assertAlmostEqual(mf.entropy, 6.1656394960533021/2, 9) @@ -56,7 +56,7 @@ def test_kuhf_smearing(self): mf = cell.KUHF(kpts=cell.make_kpts([2,1,1])).to_gpu() mf = mf.smearing(0.1, 'fermi') nkpts = len(mf.kpts) - mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)]) + mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)]) mo_energy_kpts = cp.array([mo_energy_kpts, mo_energy_kpts+cp.cos(mo_energy_kpts)*.02]) mf.get_occ(mo_energy_kpts) self.assertAlmostEqual(mf.entropy, 6.1803390081500869/2, 9) diff --git a/gpu4pyscf/pbc/tools/discretization.py b/gpu4pyscf/pbc/tools/discretization.py new file mode 100644 index 000000000..e76e8cdf2 --- /dev/null +++ b/gpu4pyscf/pbc/tools/discretization.py @@ -0,0 +1,37 @@ +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np + + +def freeze_mesh(method, cell=None, mesh=None): + """Make the current PBC mesh explicit and share it with method objects.""" + if cell is None: + cell = method.cell + if mesh is None: + mesh = cell.mesh + if mesh is None: + raise RuntimeError("PBC mesh is not initialized") + + mesh = np.asarray(mesh, dtype=np.int32) + cell.mesh = mesh + objects = ( + getattr(method, "_numint", None), + getattr(method, "with_df", None), + getattr(method, "grids", None), + ) + for obj in objects: + if obj is not None and hasattr(obj, "mesh"): + obj.mesh = mesh + return mesh diff --git a/gpu4pyscf/pbc/tools/k2gamma.py b/gpu4pyscf/pbc/tools/k2gamma.py index f27685bfa..5d0215b34 100644 --- a/gpu4pyscf/pbc/tools/k2gamma.py +++ b/gpu4pyscf/pbc/tools/k2gamma.py @@ -82,7 +82,8 @@ def double_translation_indices(kmesh): tz = cp.array(translation_map(kmesh[2]), dtype=np.int32) idx = cp.ravel_multi_index([tx[:,None,None,:,None,None], ty[None,:,None,None,:,None], - tz[None,None,:,None,None,:]], kmesh) + tz[None,None,:,None,None,:]], + tuple(int(n) for n in kmesh)) nk = np.prod(kmesh) return idx.reshape(nk, nk) diff --git a/gpu4pyscf/pbc/tools/tests/test_discretization.py b/gpu4pyscf/pbc/tools/tests/test_discretization.py new file mode 100644 index 000000000..9e86babd5 --- /dev/null +++ b/gpu4pyscf/pbc/tools/tests/test_discretization.py @@ -0,0 +1,65 @@ +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +import numpy as np +from gpu4pyscf.pbc.tools.discretization import freeze_mesh + + +class _Cell: + def __init__(self, mesh, mesh_from_build=True): + self._mesh = np.asarray(mesh) + self._mesh_from_build = mesh_from_build + + @property + def mesh(self): + return self._mesh + + @mesh.setter + def mesh(self, value): + self._mesh = value + self._mesh_from_build = False + + +def _method(cell_mesh=(20, 22, 24), object_mesh=(10, 10, 10)): + cell = _Cell(cell_mesh) + return SimpleNamespace( + cell=cell, + _numint=SimpleNamespace(mesh=object_mesh), + with_df=SimpleNamespace(mesh=object_mesh), + grids=SimpleNamespace(mesh=object_mesh), + ) + + +def test_freezes_automatic_cell_mesh_and_synchronizes_objects(): + method = _method() + + mesh = freeze_mesh(method) + + np.testing.assert_array_equal(mesh, (20, 22, 24)) + assert method.cell._mesh_from_build is False + for obj in (method.cell, method._numint, method.with_df, method.grids): + assert obj.mesh is mesh + np.testing.assert_array_equal(obj.mesh, mesh) + + +def test_reuses_reference_mesh(): + method = _method(cell_mesh=(22, 24, 26)) + + mesh = freeze_mesh(method, mesh=(20, 22, 24)) + + np.testing.assert_array_equal(mesh, (20, 22, 24)) + for obj in (method.cell, method._numint, method.with_df, method.grids): + assert obj.mesh is mesh + np.testing.assert_array_equal(obj.mesh, mesh) diff --git a/gpu4pyscf/tools/ase_interface.py b/gpu4pyscf/tools/ase_interface.py index 15dfa646b..7266e3d79 100644 --- a/gpu4pyscf/tools/ase_interface.py +++ b/gpu4pyscf/tools/ase_interface.py @@ -28,6 +28,7 @@ from pyscf.gto.mole import charge from pyscf.pbc.gto.cell import Cell from pyscf.pbc.tools.pyscf_ase import ase_atoms_to_pyscf +from gpu4pyscf.pbc.tools.discretization import freeze_mesh # These functions are copied from the development branch of PySCF and will be # provided by the pyscf.pbc.tools.pyscf_ase module in PySCF 2.11. @@ -138,8 +139,10 @@ def __init__(self, restart=None, label='PySCF', atoms=None, directory='.', self.method = method self.pbc = hasattr(method, 'cell') + self.mesh = None if self.pbc: mol = method.cell + self.mesh = freeze_mesh(method) else: mol = method.mol self.mol = mol @@ -169,6 +172,11 @@ def calculate(self, atoms=None, properties=['energy'], self.mol.set_geom_(_atoms, a=np.asarray(atoms.cell), unit='Angstrom') else: self.mol.set_geom_(_atoms, unit='Angstrom') + if self.pbc: + base_method = self.method + if self.method_scan is not None: + base_method = self.method_scan + freeze_mesh(base_method, self.mol, self.mesh) with_grad = 'forces' in properties or 'stress' in properties with_energy = with_grad or 'energy' in properties or 'dipole' in properties diff --git a/setup.py b/setup.py index a794fce72..b16b487a4 100755 --- a/setup.py +++ b/setup.py @@ -33,6 +33,42 @@ CLASSIFIERS = None PLATFORMS = None +# Selects the compute backend for this build. CUDA is the default. To build +# SYCL: +# +# python setup.py --sycl build +# CMAKE_CONFIGURE_ARGS=-DUSE_SYCL=ON pip install . +# +# The second form covers pip and other PEP 517 frontends, which do not forward +# unknown flags to setup.py; CMAKE_CONFIGURE_ARGS has to be set for a SYCL +# build regardless, to point cmake at icpx. +def build_backend(): + # --sycl / --cuda are ours, not setuptools', so strip them from argv before + # setuptools parses it -- an unrecognised global option is a hard error. + # Precedence: command line, then -DUSE_SYCL=ON in CMAKE_CONFIGURE_ARGS, + # else CUDA. + backend = None + for flag in ('--sycl', '--cuda'): + while flag in sys.argv: + sys.argv.remove(flag) + backend = flag[2:] + if backend is not None: + return backend + + if 'USE_SYCL=ON' in os.getenv('CMAKE_CONFIGURE_ARGS', ''): + return 'sycl' + return 'cuda' + + +BACKEND = build_backend() + + +def get_sycl_version(): + icpx_out = subprocess.check_output(["icpx", "--version"]).decode('utf-8') + m = re.search(r"[0-9]+\.[0-9]+\.[0-9]+", icpx_out) + return m.group(0) + + def get_cuda_version(): nvcc_out = subprocess.check_output(["nvcc", "--version"]).decode('utf-8') m = re.search(r"V[0-9]+.[0-9]+", nvcc_out) @@ -70,7 +106,20 @@ def run(self): self.announce('Configuring extensions', level=3) src_dir = os.path.abspath(os.path.join(__file__, '..', 'gpu4pyscf', 'lib')) dest_dir = os.path.join(self.build_temp, 'gpu4pyscf') - cmd = ['cmake', f'-S{src_dir}', f'-B{dest_dir}', '-DBUILD_LIBXC=OFF'] + if BACKEND == 'sycl': + # USE_SYCL defaults OFF in gpu4pyscf/lib/CMakeLists.txt so that a + # plain CUDA build needs no flags at all. SYCL uses ExchCXX where + # CUDA uses libxc, and there is no wheel for ExchCXX, so it has to + # be built from source, hence BUILD_LIBXC=ON rather than the OFF + # that CUDA passes. + libxc_arg = '-DBUILD_LIBXC=ON' + backend_args = ['-DUSE_SYCL=ON'] + else: + # CUDA takes libxc from the gpu4pyscf-libxc-cuda* wheel listed in + # install_requires rather than building it (upstream, since #110). + libxc_arg = '-DBUILD_LIBXC=OFF' + backend_args = [] + cmd = ['cmake', f'-S{src_dir}', f'-B{dest_dir}', libxc_arg] + backend_args configure_args = os.getenv('CMAKE_CONFIGURE_ARGS') if configure_args: cmd.extend(configure_args.split(' ')) @@ -115,10 +164,32 @@ def initialize_with_default_plat_name(self): # The sdist release package_name = NAME CUDA_VERSION = '12x' +elif BACKEND == 'sycl': + package_name = NAME + '-sycl' + get_sycl_version() else: CUDA_VERSION = get_cuda_version() package_name = NAME + '-cuda' + CUDA_VERSION +if BACKEND == 'sycl': + # dpnp replaces cupy, and the SYCL build supplies libxc through ExchCXX. + INSTALL_REQUIRES = [ + 'pyscf>=2.8.0', + 'pyscf-dispersion', + 'dpnp', + 'geometric', + 'packaging', + ] +else: + INSTALL_REQUIRES = [ + 'pyscf>=2.8.0', + 'pyscf-dispersion', + # Due to expm in cupyx.scipy.linalg and cutensor 2.0 + f'cupy-cuda{CUDA_VERSION}>=13.0,!=13.4.0', + 'geometric', + f'gpu4pyscf-libxc-cuda{CUDA_VERSION}==0.8.1', + 'packaging', + ] + setup( name=package_name, version=VERSION, @@ -138,12 +209,5 @@ def initialize_with_default_plat_name(self): "pytest-coverage==0.0", ], cmdclass={'build_py': CMakeBuildPy}, - install_requires=[ - 'pyscf>=2.8.0', - 'pyscf-dispersion', - f'cupy-cuda{CUDA_VERSION}>=13.0,!=13.4.0', # Due to expm in cupyx.scipy.linalg and cutensor 2.0 - 'geometric', - f'gpu4pyscf-libxc-cuda{CUDA_VERSION}==0.8.1', - 'packaging', - ] + install_requires=INSTALL_REQUIRES, )