diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 21eacea..f4f2995 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -1,36 +1,39 @@ -name: gpucpp-ci +name: gpu.cpp -on: +on: push: - branches: - - main + branches: [main] pull_request: - types: [opened, reopened, labeled, unlabeled, synchronize] - branches: - - main workflow_dispatch: jobs: - build: - runs-on: ubuntu-latest + test: + strategy: + matrix: + os: [macos-latest, ubuntu-latest] + runs-on: ${{ matrix.os }} steps: - - name: Checkout repository - uses: actions/checkout@v2 - - - name: No-op Step - run: echo "This is a no-op action" - - - name: Install CMake - run: sudo apt-get install cmake - - - name: Install dependencies - run: | - sudo apt-get update - sudo apt-get install -y libvulkan1 mesa-vulkan-drivers vulkan-tools - sudo apt-get install -y libxrandr-dev + - uses: actions/checkout@v4 + - if: runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y libvulkan1 mesa-vulkan-drivers ninja-build + - run: python3 -m pip install numpy + - run: tools/build_dawn.sh + - run: ./test - - name: Build - run: make all - - - name: Run hello world - run: make + browser-build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/checkout@v4 + with: + repository: emscripten-core/emsdk + path: emsdk + fetch-depth: 1 + - run: | + sudo apt-get update + sudo apt-get install -y ninja-build + - run: ./test --build-web + env: + GPUCPP_EMSDK_ROOT: ${{ github.workspace }}/emsdk diff --git a/.gitignore b/.gitignore index 1a8b5bc..a6f10b7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ build/* +build-web/ +build-dawn/ # any build subdirectory in the tree **/build/ examples/hello_gpu/build/* @@ -8,6 +10,8 @@ source .DS_Store third_party/lib/* third_party/local/* +third_party/dawn/* +third_party/emdawnwebgpu/* # formatter files .cmake-format.py @@ -19,4 +23,3 @@ build # clangd files .cache compile_commands.json - diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..2740924 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,35 @@ +# Changelog + +## 0.2.0 - 2026-08-21 + +This release modernizes gpu.cpp around current Dawn and adds a maintained +browser target. + +### Added + +- A browser backend using Emdawnwebgpu, Emscripten, JSPI, and an ES-module + Embind API for persistent contexts and JavaScript typed arrays. +- Explicit read/read-write bindings, reusable kernels, uniform parameters, and + asynchronous dispatch and readback through `gpu::Future`. +- Native SPIR-V input with an executable WebGPU-profile contract test. +- C++ stories for WGSL, native f16, and SPIR-V; a NumPy Python story; and a + Chrome story covering browser errors, context reuse, dispatch, and readback. +- Reproducible scripts pinning Dawn, Emdawn, emsdk, SPIRV-Tools, and + SPIRV-Headers to exact revisions. + +### Changed + +- WebGPU objects now use Dawn's generated C++ RAII facade and normal C++ value + semantics. +- The Python binding now follows the C++ API and preserves NumPy shape and + `float16`, `float32`, and `int32` dtypes. +- Host-side f16 uses native `_Float16` where supported; the portable core + treats f16 as two-byte IEEE 754 storage. +- Examples and the build are consolidated under the root CMake project and the + single `./test` entrypoint. + +### Removed + +- The legacy raw WebGPU C API implementation, manual resource pools, custom + half implementation, Haskell binding, obsolete build files, and abandoned + experimental targets. diff --git a/CMakeLists.txt b/CMakeLists.txt index db89df7..89643b7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,50 +1,154 @@ cmake_minimum_required(VERSION 3.28) -project(gpu) +project(gpu.cpp LANGUAGES CXX) -include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/webgpu.cmake") - -set(CMAKE_EXPORT_COMPILE_COMMANDS ON) # export compile_commands.json to use with - # LSP -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +if(EMSCRIPTEN) + set(GPUCPP_BUILD_PYTHON_DEFAULT OFF) + set(GPUCPP_BUILD_EXAMPLES_DEFAULT OFF) +else() + set(GPUCPP_BUILD_PYTHON_DEFAULT ${PROJECT_IS_TOP_LEVEL}) + set(GPUCPP_BUILD_EXAMPLES_DEFAULT ON) +endif() +option(GPUCPP_BUILD_PYTHON "Build the Python binding" ${GPUCPP_BUILD_PYTHON_DEFAULT}) +option(GPUCPP_BUILD_EXAMPLES "Build gpu.cpp examples" ${GPUCPP_BUILD_EXAMPLES_DEFAULT}) + +add_library(gpucpp INTERFACE) +target_include_directories(gpucpp INTERFACE "${CMAKE_CURRENT_SOURCE_DIR}") -option(USE_LOCAL_LIBS - "Use local libraries instead of fetching from the internet" OFF) +if(EMSCRIPTEN) + set(GPUCPP_EMDAWN_ROOT + "${CMAKE_CURRENT_SOURCE_DIR}/third_party/emdawnwebgpu" CACHE PATH + "Staged Emdawnwebgpu package") + set(EMDAWN_PORT + "--use-port=${GPUCPP_EMDAWN_ROOT}/emdawnwebgpu.port.py") + if(NOT EXISTS "${GPUCPP_EMDAWN_ROOT}/emdawnwebgpu.port.py") + message(FATAL_ERROR "Emdawnwebgpu is not built; run tools/build_emdawn.sh") + endif() + target_compile_options(gpucpp INTERFACE "${EMDAWN_PORT}" "-fwasm-exceptions") + target_link_options(gpucpp INTERFACE "${EMDAWN_PORT}" "-fwasm-exceptions" + "-sJSPI=1" "--closure=1") + if(CMAKE_HOST_APPLE) + target_link_options(gpucpp INTERFACE "--closure-args=--platform=java") + endif() +else() -# Ensure the build type is set -if(NOT CMAKE_BUILD_TYPE) - set(CMAKE_BUILD_TYPE - Release - CACHE STRING "Choose the type of build: Debug or Release" FORCE) +set(GPUCPP_DAWN_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/third_party/dawn" CACHE PATH + "Staged Dawn distribution") +set(GPUCPP_LOCAL_DAWN "${CMAKE_CURRENT_SOURCE_DIR}/third_party/local/dawn" CACHE PATH + "Developer Dawn source/build tree") + +if(EXISTS "${GPUCPP_DAWN_ROOT}/include/dawn/webgpu_cpp.h") + set(DAWN_INCLUDE_DIRS "${GPUCPP_DAWN_ROOT}/include") + if(APPLE) + set(DAWN_LIBRARY "${GPUCPP_DAWN_ROOT}/lib/libwebgpu_dawn.dylib") + elseif(UNIX) + set(DAWN_LIBRARY "${GPUCPP_DAWN_ROOT}/lib/libwebgpu_dawn.so") + endif() + set(SPIRV_AS "${GPUCPP_DAWN_ROOT}/bin/spirv-as") +elseif(EXISTS "${GPUCPP_LOCAL_DAWN}/build-latest/gen/include/dawn/webgpu_cpp.h") + set(DAWN_INCLUDE_DIRS + "${GPUCPP_LOCAL_DAWN}/build-latest/gen/include" + "${GPUCPP_LOCAL_DAWN}/source/include") + if(APPLE) + set(DAWN_LIBRARY + "${GPUCPP_LOCAL_DAWN}/build-latest/src/dawn/native/libwebgpu_dawn.dylib") + elseif(UNIX) + set(DAWN_LIBRARY + "${GPUCPP_LOCAL_DAWN}/build-latest/src/dawn/native/libwebgpu_dawn.so") + endif() + set(SPIRV_AS "${GPUCPP_LOCAL_DAWN}/build-spirv-tools/tools/spirv-as") +else() + message(FATAL_ERROR "Dawn is not built; run tools/build_dawn.sh") endif() -option(FASTBUILD "Option to enable fast builds" OFF) -if(FASTBUILD) - set(CMAKE_BUILD_TYPE None) # Avoid default flags of predefined build types - set(CMAKE_CXX_FLAGS "-O0") +if(NOT EXISTS "${DAWN_LIBRARY}") + message(FATAL_ERROR "Dawn library not found at ${DAWN_LIBRARY}") +endif() +if(NOT EXISTS "${SPIRV_AS}") + message(FATAL_ERROR "spirv-as not found at ${SPIRV_AS}") endif() -option(DEBUG "Option to enable debug flags" OFF) -if(DEBUG) - set(CMAKE_BUILD_TYPE Debug) - set(CMAKE_CXX_FLAGS "-O0 -g") +add_library(Dawn::webgpu SHARED IMPORTED) +set_target_properties(Dawn::webgpu PROPERTIES + IMPORTED_LOCATION "${DAWN_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${DAWN_INCLUDE_DIRS}") + +target_link_libraries(gpucpp INTERFACE Dawn::webgpu) endif() -if(WIN64) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DWEBGPU_BACKEND_DAWN") +if(GPUCPP_BUILD_PYTHON) + find_package(Python 3.9 COMPONENTS Interpreter Development.Module REQUIRED) + include(FetchContent) + FetchContent_Declare(pybind11 + GIT_REPOSITORY https://github.com/pybind/pybind11.git + GIT_TAG v3.0.4 + GIT_SHALLOW TRUE) + FetchContent_MakeAvailable(pybind11) + + pybind11_add_module(gpu_cpp bindings/python/gpu_cpp.cpp) + target_link_libraries(gpu_cpp PRIVATE gpucpp) endif() -include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/gpu.cmake") +if(GPUCPP_BUILD_EXAMPLES) + add_executable(hello_gpu examples/hello_world/run.cpp) + add_executable(float16_gpu examples/float16/run.cpp) + add_executable(gpu_puzzles examples/gpu_puzzles/run.cpp) + add_executable(gpu_puzzles_key examples/gpu_puzzles/key.cpp) + add_executable(matmul_gpu examples/matmul/run.cpp) + add_executable(physics_gpu examples/physics/run.cpp) + add_executable(render_gpu examples/render/run.cpp) + add_executable(shadertui_gpu examples/shadertui/run.cpp) + add_executable(transpose_gpu examples/transpose/run.cpp) -message(STATUS "CMAKE_CURRENT_SOURCE_DIR: ${CMAKE_CURRENT_SOURCE_DIR}") -message( - STATUS - "Include directories for wgpu: ${CMAKE_CURRENT_SOURCE_DIR}/third_party/headers" -) + foreach(target IN ITEMS hello_gpu float16_gpu gpu_puzzles gpu_puzzles_key + matmul_gpu physics_gpu render_gpu shadertui_gpu + transpose_gpu) + target_link_libraries(${target} PRIVATE gpucpp) + endforeach() + target_include_directories(matmul_gpu PRIVATE third_party/headers) + target_include_directories(transpose_gpu PRIVATE third_party/headers) -add_library(gpud SHARED gpu.hpp) -set_target_properties(gpud PROPERTIES LINKER_LANGUAGE CXX) -target_link_libraries(gpud PRIVATE wgpu) -target_link_libraries(gpud PRIVATE webgpu) -target_link_libraries(gpud PRIVATE gpu) -install(TARGETS gpud) +endif() + +enable_testing() +if(EMSCRIPTEN) + add_library(gpu_cpp_web_bindings OBJECT bindings/web/gpu_cpp.cpp) + target_link_libraries(gpu_cpp_web_bindings PRIVATE gpucpp) + + add_executable(gpu_cpp_web $) + set_target_properties(gpu_cpp_web PROPERTIES SUFFIX ".mjs") + target_link_libraries(gpu_cpp_web PRIVATE gpucpp) + target_link_options(gpu_cpp_web PRIVATE "--bind" "--no-entry" + "-sMODULARIZE=1" "-sEXPORT_ES6=1" "-sEXPORT_NAME=createGpuCpp" + "-sENVIRONMENT=web" "-sALLOW_MEMORY_GROWTH=1") + + add_executable(web_binding_test tests/web_binding_test.cpp + $) + set_target_properties(web_binding_test PROPERTIES SUFFIX ".html") + target_link_libraries(web_binding_test PRIVATE gpucpp) + target_link_options(web_binding_test PRIVATE "--bind" "--emrun" + "-sMODULARIZE=1" "-sEXPORT_ES6=1" + "-sALLOW_MEMORY_GROWTH=1") +else() + set(TEST_SPIRV "${CMAKE_CURRENT_BINARY_DIR}/write42.spv") + add_custom_command( + OUTPUT "${TEST_SPIRV}" + COMMAND "${SPIRV_AS}" --target-env spv1.3 + "${CMAKE_CURRENT_SOURCE_DIR}/tests/write42.spvasm" -o "${TEST_SPIRV}" + DEPENDS tests/write42.spvasm + VERBATIM) + + add_executable(gpu_test tests/gpu_test.cpp "${TEST_SPIRV}") + target_link_libraries(gpu_test PRIVATE gpucpp) + add_test(NAME gpu COMMAND gpu_test "${TEST_SPIRV}") +endif() +if(GPUCPP_BUILD_PYTHON) + add_test(NAME python + COMMAND "${Python_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/bindings/python/test_gpu_cpp.py") + set_tests_properties(python PROPERTIES + ENVIRONMENT "PYTHONPATH=$") +endif() diff --git a/DEV.md b/DEV.md new file mode 100644 index 0000000..f6b445c --- /dev/null +++ b/DEV.md @@ -0,0 +1,91 @@ +# Development + +## Dependency model + +Dawn and its native dependencies are pinned in `tools/dependencies.sh`. +`tools/build_dawn.sh` maintains ignored working +trees under `third_party/local/dawn/` and stages the small runtime distribution +under `third_party/dawn/`. Neither directory is committed. + +The same file pins the emsdk revision and Emscripten version recorded in +Dawn's DEPS. `tools/build_emdawn.sh` builds `emdawnwebgpu_pkg` from that exact +Dawn source and stages its local Emscripten port under +`third_party/emdawnwebgpu/`. The sibling `../emsdk` clone and staged package are +not committed. Closure is required by Emdawn release links and uses OpenJDK on +macOS ARM. + +The top-level build fetches the tagged pybind11 release declared in +`CMakeLists.txt`; consumers using only the C++ target do not fetch or build it. +The Python runtime test also requires NumPy. + +The current Dawn revision needs compatible SPIRV-Tools and SPIRV-Headers +revisions for `TINT_BUILD_SPV_READER`; those exact revisions are pinned beside +the Dawn revision. This is dependency selection, not a Dawn source patch. + +The build enables only the native platform backend (Metal or Vulkan), the WGSL +reader, and the SPIR-V reader. Dawn's Null backend is disabled because it +accepts work but intentionally does not execute it. + +The Emdawn package build disables protobuf, Tint's protobuf-based IR format, +and C++ modules because none are part of the browser port. This also avoids a +false-positive Emscripten C++-module probe at the current Dawn revision. + +To update Dawn and Emdawn: + +1. Change the Dawn, SPIRV-Tools, and SPIRV-Headers revisions in + `tools/dependencies.sh`. +2. Copy the emsdk revision and corresponding Emscripten version from the new + Dawn revision's `DEPS`. +3. Run `./test --rebuild-dawn`, then `./test` on a machine with a real GPU. +4. Run `./test --rebuild-web` in Chrome. +5. Update the SPIR-V version/profile documentation if Tint's reader changed. + +No gpu.cpp patches are applied to the Dawn source tree. + +## Tests + +`./test` is the only test entrypoint. It configures with CMake/Ninja, builds all +maintained examples, and runs two readable end-to-end stories through CTest. +The C++ story covers: + +1. WGSL upload, explicit read/read-write bindings, dispatch, and readback. +2. Native `_Float16` upload and execution through an f16 WGSL pipeline. +3. Assembly and execution of a minimal WebGPU-compatible SPIR-V module. + +The SPIR-V assembly is kept as text so the compiler/runtime contract is +reviewable. Avoid adding tests that merely restate Dawn validation; add a new +story only when gpu.cpp itself owns meaningful behavior. + +The Python story covers NumPy upload/readback, explicit binding access, +uniform parameters, and reuse of a compiled kernel with updated tensor data. + +`./test --web` cross-compiles the same public core and Embind API against +Emdawnwebgpu, opens the result with `emrun`, and checks rejected invalid WGSL, +context reuse, typed-array upload, dispatch, and readback in Chrome. The +release build uses Closure, Wasm exceptions, and JSPI. CI uses +`./test --build-web` to compile the same artifacts without requiring a browser +GPU; the complete story remains the local runtime contract. + +## Architecture + +`gpu.hpp` is the library. Public objects own Dawn's `wgpu` RAII handles, so +normal C++ lifetime rules replace the old tensor/kernel pools and manual C API +release bookkeeping. + +Pipeline creation uses a validation error scope and returns ordinary C++ +exceptions with Dawn's diagnostic. Permanent device callbacks never throw; +they write into `Context` error state, which foreground API calls surface. + +`gpu::Future` retains both the callback result and Dawn's `wgpu::Future`. +Native waits pump `ProcessEvents`; browser waits call `WaitAny`, allowing JSPI +to suspend Wasm while JavaScript and WebGPU make progress. + +The browser Embind layer owns one persistent `Context`. Each `run()` creates +short-lived tensors and a kernel from JavaScript typed arrays, rejects +overlapping calls, and copies `readWrite` results back in place. Wasm +exceptions turn C++ validation failures into rejected JavaScript Promises. + +Each dispatch creates a command encoder because WebGPU command buffers are +single-use. Pipelines and bind groups remain reusable. Readback uses a staging +buffer and returns a future whose callback owns that buffer until mapping is +complete. diff --git a/Makefile b/Makefile index 8e5d67b..5ee7d49 100644 --- a/Makefile +++ b/Makefile @@ -1,159 +1,16 @@ -NUM_JOBS=$(shell nproc) -CXX=clang++ +.PHONY: all build dawn run test -.PHONY: default examples/hello_world/build/hello_world tests libgpu debug build check-clang clean-build clean all watch-tests docs +all: build -GPUCPP ?= $(PWD) -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -INCLUDES ?= -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif +build: + cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release + cmake --build build -default: examples/hello_world/build/hello_world +dawn: + tools/build_dawn.sh -pch: - mkdir -p build && $(CXX) -std=c++17 $(INCLUDES) -x c++-header gpu.hpp -o build/gpu.hpp.pch +run: build + ./build/hello_gpu -# TODO(avh): change extension based on platform -# Get the current OS name -OS = $(shell uname | tr -d '\n') -# Set the specific variables for each platform -LIB_PATH ?= /usr/lib -HEADER_PATH ?= /usr/include -ifeq ($(OS), Linux) -OS_TYPE ?= Linux -GPU_CPP_LIB_NAME ?= libgpucpp.so -DAWN_LIB_NAME ?= libwebgpu_dawn.so -else ifeq ($(OS), Darwin) -OS_TYPE ?= macOS -GPU_CPP_LIB_NAME ?= libgpucpp.dylib -DAWN_LIB_NAME ?= libwebgpu_dawn.dylib -else -OS_TYPE ?= unknown -endif - -lib: check-clang dawnlib - mkdir -p build && $(CXX) -std=c++17 $(INCLUDES) -L$(LIBDIR) -lwebgpu_dawn -ldl -shared -fPIC gpu.cpp -o build/$(GPU_CPP_LIB_NAME) - python3 build.py - cp third_party/lib/$(DAWN_LIB_NAME) build/ - -install: - cp build/$(GPU_CPP_LIB_NAME) $(LIB_PATH) - cp build/$(DAWN_LIB_NAME) $(LIB_PATH) - cp build/gpu.hpp $(HEADER_PATH) - -uninstall: - rm $(LIB_PATH)/$(GPU_CPP_LIB_NAME) - rm $(LIB_PATH)/$(DAWN_LIB_NAME) - rm $(HEADER_PATH)/gpu.hpp - -examples/hello_world/build/hello_world: check-clang dawnlib examples/hello_world/run.cpp check-linux-vulkan - $(LIBSPEC) && cd examples/hello_world && make build/hello_world && ./build/hello_world - -dawnlib: $(if $(wildcard third_party/lib/libwebgpu_dawn.so third_party/lib/libwebgpu_dawn.dylib),,run_setup) - -run_setup: check-python - python3 setup.py - -all: dawnlib check-clang check-linux-vulkan lib pch - cd examples/float16 && make build/float16 - cd examples/gpu_puzzles && make build/gpu_puzzles - cd examples/hello_world && make build/hello_world - cd examples/matmul && make build/matmul - cd examples/physics && make build/physics - cd examples/render && make build/render - cd examples/shadertui && make build/shadertui - cd examples/transpose && make build/transpose - -# Test 16-bit floating point type -test-half: dawnlib check-clang - $(LIBSPEC) && clang++ -std=c++17 $(INCLUDES) numeric_types/half.cpp -L$(LIBDIR) -lwebgpu_dawn -ldl -o build/half && ./build/half - -docs: Doxyfile - doxygen Doxyfile - -################################################################################ -# cmake targets (optional - precompiled binaries is preferred) -################################################################################ - -CMAKE_CMD = mkdir -p build && cd build && cmake .. -# Add --trace to see the cmake commands -FLAGS = -DCMAKE_VERBOSE_MAKEFILE:BOOL=ON -DCMAKE_CXX_COMPILER=$(CXX) -DABSL_INTERNAL_AT_LEAST_CXX20=OFF -FASTBUILD_FLAGS = $(FLAGS) -DFASTBUILD:BOOL=ON -DEBUG_FLAGS = $(FLAGS) -DDEBUG:BOOL=ON -RELEASE_FLAGS = $(FLAGS) -DFASTBUILD:BOOL=OFF -TARGET_LIB=gpu - -libgpu-cmake: check-clang check-cmake - $(CMAKE_CMD) $(RELEASE_FLAGS) && make -j$(NUM_JOBS) gpu - -debug-cmake: check-clang check-cmake - $(CMAKE_CMD) $(DEBUG_FLAGS) && make -j$(NUM_JOBS) $(TARGET_ALL) - -all-cmake: check-clang check-cmake - $(CMAKE_CMD) $(RELEASE_FLAGS) && make -j$(NUM_JOBS) $(TARGET_ALL) - -################################################################################ -# Cleanup -################################################################################ - -clean-dawnlib: - rm -f third_party/lib/libwebgpu_dawn.so third_party/lib/libwebgpu_dawn.dylib - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* - rm -rf examples/float16/build/* - rm -rf examples/gpu_puzzles/build/* - rm -rf examples/hello_world/build/* - rm -rf examples/matmul/build/matmul - rm -rf examples/physics/build/* - rm -rf examples/render/build/* - rm -rf examples/shadertui/build/* - rm -rf examples/transpose/build/transpose - rm -f build/gpu.hpp.pch - rm -f build/libgpucpp.so - rm -f build/half - -clean-all: - read -r -p "This will delete the contents of build/* and third_party/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* third_party/fetchcontent/* third_party/gpu-build third_party/gpu-subbuild third_party/gpu-src third_party/lib/libwebgpu_dawn.so third_party/lib/libwebgpu_dawn.dylib - -################################################################################ -# Checks -################################################################################ - -# Check all -check-all: check-os check-clang check-cmake check-python - -# check the os -check-os: -ifeq ($(OS_TYPE), unknown) -$(error Unsupported operating system) -endif - -# check for the existence of clang++ and cmake -check-clang: - @command -v clang++ >/dev/null 2>&1 || { echo -e >&2 "Clang++ is not installed. Please install clang++ to continue.\nOn Debian / Ubuntu: 'sudo apt-get install clang' or 'brew install llvm'\nOn Centos: 'sudo yum install clang'"; exit 1; } - -check-cmake: - @command -v cmake >/dev/null 2>&1 || { echo -e >&2 "Cmake is not installed. Please install cmake to continue.\nOn Debian / Ubuntu: 'sudo apt-get install cmake' or 'brew install cmake'\nOn Centos: 'sudo yum install cmake'"; exit 1; } - -check-python: - @command -v python3 >/dev/null 2>&1 || { echo -e >&2 "Python is not installed. Please install python to continue.\nOn Debian / Ubuntu: 'sudo apt-get install python'\nOn Centos: 'sudo yum install python'"; exit 1; } - -check-linux-vulkan: - @echo "Checking system type and Vulkan availability..." - @if [ "$$(uname)" = "Linux" ]; then \ - if command -v vulkaninfo >/dev/null 2>&1; then \ - echo "Vulkan is installed."; \ - vulkaninfo; \ - else \ - echo -e "Vulkan is not installed. Please install Vulkan drivers to continue.\nOn Debian / Ubuntu: 'sudo apt install libvulkan1 mesa-vulkan-drivers vulkan-tools'.\nOn Centos: 'sudo yum install vulkan vulkan-tools.'"; \ - exit 1; \ - fi \ - else \ - echo "Non-Linux system detected. Skipping Vulkan check."; \ - fi +test: + ./test diff --git a/README.md b/README.md index 46340b7..02c5e0d 100644 --- a/README.md +++ b/README.md @@ -1,333 +1,201 @@ # gpu.cpp -gpu.cpp is a lightweight library that makes portable GPU compute with C++ simple. +gpu.cpp is a small, header-only C++20 interface for GPU compute through +[Dawn](https://dawn.googlesource.com/dawn), Google's WebGPU implementation. It +uses Metal on macOS, Vulkan on Linux, and Emdawnwebgpu in browsers. -It focuses on general purpose native GPU computation, leveraging the WebGPU -specification as a portable low-level GPU interface. This means we can drop in -GPU code in C++ projects and have it run on Nvidia, Intel, AMD, and other GPUs. -The same C++ code can work on a wide variety of laptops, workstations, mobile -devices or virtually any hardware with Vulkan, Metal, or DirectX support. +The library deliberately stays close to WebGPU: it owns the native resources, +handles asynchronous work and errors, and removes repetitive descriptor code, +but shaders still declare their resources and execution model explicitly. -## Objectives: Lightweight, Fast Iteration, and Low Boilerplate +## Build and test -With gpu.cpp we want to enable a high-leverage library for individual developers and researchers to incorporate GPU computation into programs relying on nothing more than a standard C++ compiler as tooling. Our goals are: +Install CMake, Ninja, Python 3, and a C++20 compiler. Linux also needs a Vulkan +driver; Mesa's Vulkan driver is sufficient for development and CI. -- High power-to-weight ratio API: Provide the smallest API surface area that can cover the full range of GPU compute needs. -- Fast compile/run cycles: Ensure projects can build nearly instantaneously, compile/run cycles should be <5 seconds on a modern laptop. -- Minimal dependencies and tooling overhead: A standard clang C++ compiler should be enough, no external library dependencies beyond the WebGPU native implementation. - -The implementation aims for a small API surface area with minimum boilerplate. There are a small number of library operations to carry out an broad range of low-level GPU operations. We avoid abstractions that add layers of indirection, making the mapping between the gpu.cpp library to raw WebGPU API clear when it's needed. - -In this spirit of fast experimentation, we also want near-instantaneous C++ builds taking no more than a second or two even on modestly capable personal computing devices. With this in mind, we not only keep the API surface area small, but also keep the implementation small and we also provide a prebuilt binary of the Dawn native WebGPU implementation. - -The core library implementation in the header-only `gpu.hpp` source code is around 1000 lines of code. In addition to enabling instantaneous, semi-interactive compilation cycles, the small implementation surface area keeps maintenance burden low and the velocity of improvements high. -We also pre-build Google's Dawn WebGPU implementation as a shared library binary. This allows builds to link the shared library with each build and incorporate Google's powerful native WebGPU implementation without paying the cost of re-compiling Dawn during development cycles. - -For more advanced users and release deployments, we include `cmake` examples for building both Dawn with gpu.cpp end-to-end, but this is not required nor recommended for most users to get started. - -## Quick Start: Building and Running - -To build a gpu.cpp project, you will need to have installed on your system: - -- `clang++` compiler installed with support for C++17. -- `python3` and above, to run the script which downloads the Dawn shared library. -- `make` to build the project. -- Only on Linux systems - Vulkan drivers. If Vulkan is not installed, you can run `sudo apt install libvulkan1 mesa-vulkan-drivers vulkan-tools` to install them. - -The only library dependency of gpu.cpp is a WebGPU implementation. Currently we support the Dawn native backend, but we plan to support other targets and WebGPU implementations (web browsers or other native implementations such as wgpu). Currently we support MacOS, Linux, and Windows (via WSL). - -Optionally, Dawn can be built from scratch with gpu.cpp using the cmake build scripts provided - see the -cmake targets in the Makefile. However, this is recommended for advanced users only. Building Dawn dependencies with cmake takes much longer than using the precompiled Dawn shared library. - -After cloning the repo, from the top-level gpu.cpp, you should be able to build and run the hello world GELU example by typing: - -``` -make +```bash +./test --rebuild-dawn # first setup, or after changing the pinned Dawn revision +./test # configure, build, and run the native GPU stories +./test --rebuild-web # first browser setup +./test --web # build and run the browser story in Chrome +make run # run the hello-world example ``` -The first time you build and run the project this way, it will download a prebuilt shared library for the Dawn native WebGPU implementation automatically (using the setup.py script). This places the Dawn shared library in the `third_party/lib` directory. Afterwards you should see `libdawn.dylib` on MacOS or `libdawn.so` on Linux. This download only occurs once. +`tools/build_dawn.sh` checks out exact revisions, builds a monolithic shared +Dawn library, and stages its headers, library, and `spirv-as` under +`third_party/dawn/`. During Dawn development, CMake can instead use the source +and build trees under `third_party/local/dawn/` directly. -The build process itself should take a few seconds. If the build and executions is successful, you should see the output of the GELU computation: +Browser builds require Chrome with WebGPU and JSPI, plus a sibling `../emsdk` +clone. On macOS they also require OpenJDK for Closure: +```bash +git clone --depth 1 https://github.com/emscripten-core/emsdk.git ../emsdk +brew install openjdk +./test --rebuild-web ``` -Hello gpu.cpp! --------------- - - gelu(0.00) = 0.00 - gelu(0.10) = 0.05 - gelu(0.20) = 0.12 - gelu(0.30) = 0.19 - gelu(0.40) = 0.26 - gelu(0.50) = 0.35 - gelu(0.60) = 0.44 - gelu(0.70) = 0.53 - gelu(0.80) = 0.63 - gelu(0.90) = 0.73 - gelu(1.00) = 0.84 - gelu(1.10) = 0.95 - ... - -Computed 10000 values of GELU(x) -``` - -If you need to clean up the build artifacts, you can run: - -``` -make clean -``` - -## Hello World Tutorial: A GELU Kernel -As a real-world example for how to use gpu.cpp, let's start with a practical-but-simple example of a GPU kernel from neural networks. +`tools/build_emdawn.sh` moves that clone to the exact emsdk revision pinned by +Dawn, installs its Emscripten SDK, and stages an Emdawn local port under +`third_party/emdawnwebgpu/`. Nothing needs to be added to the shell profile. -GELU is a non-linear embarassingly parallel operation often used in modern large language model transformer-based architectures. +Native tests must have access to a real Metal or Vulkan adapter. gpu.cpp +disables Dawn's Null backend so an inaccessible GPU fails clearly rather than +producing plausible-looking no-op executions. -It takes as input a vector of floats and applies the GELU function to each element of the vector. The function is nonlinear, attenuating values below zero to near zero, approximating the y = x identity function for large positive values. For values close to zero, GELU smoothly interpolates between the identity function and the zero function. - -The GELU code below will illustrate the three main aspects of setting up a GPU computation with gpu.cpp: - -1. The code that runs on the GPU (in WebGPU Shading Language, or WGSL), implementing the compute operation. - -2. The code that runs on the CPU (in C++) that sets up the GPU computation by allocating and preparing resources. For high performance, this code should be run ahead-of-time from the hot paths of the application. - -3. The code that runs on the CPU (in C++) that dispatches the GPU computation and retrieves the results. The key concern of hot-path dispatch code is to eliminate or minimize any unnecessary resource allocation or data movement (offloading such concerns to step 2). A secondary consideration is that GPU dispatches are asynchronous. We work with standard C++ asynchronous primitives to manage the asynchronous aspect of kernel dispatch. - -Here's a GELU kernel implemented (based on the CUDA implementation in [llm.c](https://github.com/karpathy/llm.c)) as on-device WebGPU WGSL code and invoked from the host using gpu.cpp library functions and types. It can be compiled using a standard C++ compiler (we recommend Clang): +## Example ```cpp -#include -#include -#include - #include "gpu.hpp" -using namespace gpu; // createContext, createTensor, createKernel, - // dispatchKernel, wait, toCPU Bindings, - // Tensor, Kernel, Context, Shape, kf32 +#include +#include + +using namespace gpu; + +static constexpr auto twice = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; -static const char *kGelu = R"( -const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI) -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; @compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&inp)) { - let x: f32 = inp[i]; - out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR - * (x + .044715 * x * x * x))), x, x > 10.0); - } +fn main(@builtin(global_invocation_id) id: vec3) { + if (id.x < arrayLength(&input)) { + output[id.x] = input[id.x] * 2; + } } )"; -int main(int argc, char **argv) { - Context ctx = createContext(); - static constexpr size_t N = 10000; - std::array inputArr, outputArr; - for (int i = 0; i < N; ++i) { - inputArr[i] = static_cast(i) / 10.0; // dummy input data - } - Tensor input = createTensor(ctx, Shape{N}, kf32, inputArr.data()); - Tensor output = createTensor(ctx, Shape{N}, kf32); - std::promise promise; - std::future future = promise.get_future(); - Kernel op = createKernel(ctx, {kGelu, /* 1-D workgroup size */ 256, kf32}, - Bindings{input, output}, - /* number of workgroups */ {cdiv(N, 256), 1, 1}); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, outputArr.data(), sizeof(outputArr)); - for (int i = 0; i < 16; ++i) { - printf(" gelu(%.2f) = %.2f\n", inputArr[i], outputArr[i]); - } - return 0; +int main() { + auto context = createContext(); + std::vector input{1, 2, 3, 4}; + std::vector output(input.size()); + auto gpuInput = createTensor(context, {input.size()}, ki32, input); + auto gpuOutput = createTensor(context, {output.size()}, ki32); + auto kernel = createKernel(context, WGSL{twice, 4}, + Bindings{read(gpuInput), readWrite(gpuOutput)}); + + auto dispatched = dispatchKernel(context, kernel); + wait(context, dispatched); + auto downloaded = toCPU(context, gpuOutput, output); + wait(context, downloaded); } ``` -Here we see the GPU code is quoted in a domain specific language called WGSL (WebGPU Shading Language). In a larger project, you might store this code in a separate file to be loaded at runtime (see [examples/shadertui](https://github.com/AnswerDotAI/gpu.cpp/tree/main/examples/shadertui) for a demonstration of live WGSL code re-loading). - -The CPU code in main() sets up the host coordination for the GPU computation. -We can think of the use of gpu.cpp library as a collection of GPU nouns and -verbs. - -The "nouns" are GPU resources modeled by the type definitions of the library -and the "verbs" actions on GPU resources, modeled by the functions of the -library. The ahead-of-time resource acquisition functions are prefaced with -`create*`, such as: - -- `createContext()` - constructs a reference to the GPU device context (`Context`). -- `createTensor()` - acquires a contiguous buffer on the GPU (`Tensor`). -- `createKernel()` - constructs a handle to resources for the GPU computation (`Kernel`), taking the shader code as input and the tensor resources to bind. - -These resource acquisition functions are tied to resource types for interacting with the GPU: - -- `Context` - a handle to the state of resources for interacting with the GPU device. -- `Tensor` - a buffer of data on the GPU. -- `KernelCode` - the code for a WGSL program that can be dispatched to the - GPU. This is a thin wrapper around a WGSL string and also includes the - workgroup size the code is designed to run with. -- `Kernel` - a GPU program that can be dispatched to the GPU. This accepts a - `KernelCode` and a list of `Tensor` resources to bind for the dispatch - computation. This takes an argument `Bindings` that is a list of `Tensor` instances and should map the bindings declared at the top of the WGSL code. In this example there's two bindings corresponding to the `input` buffer on the GPU and the `ouptut` buffer on the GPU. - -In this example, the GELU computation is performed only once and the program immediately exits so preparing resources and dispatch are side-by-side. Other examples in the [examples/](https://github.com/AnswerDotAI/gpu.cpp/blob/main/examples/) directory illustrate how resource acquisition is prepared ahead of time and dispatch occurs in the hot path like a render, model inference, or simulation loop. - -Besides the `create*` resource acquisition functions, there are a few more "verbs" in the gpu.cpp library for handling dispatching execution to the GPU and data movement: +`read()` and `readWrite()` are explicit because WebGPU pipeline layouts must +agree with the shader declarations. Dawn validates the agreement when +`createKernel()` builds the pipeline. -- `dispatchKernel()` - dispatches a `Kernel` to the GPU for computation. This is an asynchronous operation that returns immediately. -- `wait()` - blocks until the GPU computation is complete. This is a standard C++ future/promise pattern. -- `toCPU()` - moves data from the GPU to the CPU. This is a synchronous operation that blocks until the data is copied. -- `toGPU()` - moves data from the CPU to the GPU. This is a synchronous operation that blocks until the data is copied. In this particular example, `toGPU()` is not used because there's only one data movement from CPU to GPU in the program and that happens when the `createTensor()` function is called. +## API shape -This example is available in [examples/hello_world/run.cpp](https://github.com/AnswerDotAI/gpu.cpp/blob/main/examples/hello_world/run.cpp). +- `Context` owns the instance, native adapter, device, queue, and error state. +- `Tensor` owns a WebGPU buffer plus its shape and numeric type. +- `WGSL` and `SPIRV` own shader source and entry-point metadata. +- `Kernel` owns a reusable compute pipeline, bind group, and optional uniform + parameter buffer. +- `dispatchKernel()` and `toCPU()` return `gpu::Future`; `wait()` pumps Dawn + events natively or suspends through JSPI in a browser, and propagates + asynchronous failures. +- `toGPU()` updates an existing tensor or a kernel's parameter buffer. -## Other Examples: Matrix Multiplication, Physics Sim, and SDF Rendering +The portable core treats f16 host data as two-byte IEEE 754 storage. Native +code may include `numeric_types/half.hpp` for the compiler's `_Float16` type; +browser bindings can use `Float16Array` or raw `Uint16Array` storage. -You can explore the example projects in -[examples/](https://github.com/AnswerDotAI/gpu.cpp/blob/main/examples/) which -illustrate how to use gpu.cpp as a library. +All WebGPU handles use Dawn's generated `wgpu` C++ RAII facade. gpu.cpp does +not maintain parallel resource pools or manually release C handles. -After you have run `make` in the top-level directory which retrieves the prebuilt Dawn shared library, you can run each example by navigating to its directory and running `make` from the example's directory. +## SPIR-V input -An example of tiled matrix multiplication is in [examples/matmul](https://github.com/AnswerDotAI/gpu.cpp/blob/main/examples/matmul/). This implements a WebGPU version of the first few kernels of Simon Boehm's [How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog](https://siboehm.com/articles/22/CUDA-MMM) post. It currently runs at ~ 3.5+ TFLOPs on a Macbook Pro M1 Max laptop. Contributions to optimize this further are welcome. +SPIR-V is an opt-in Dawn instance feature: -A parallel physics simulation of an ensemble of double pendulums simulated in parallel with different initial conditions on the GPU is shown in [examples/physics](https://github.com/AnswerDotAI/gpu.cpp/tree/main/examples/physics). - -
-matmul example output -physics example animated gif -
- -We also show some examples of signed distance function computations, rendered in the terminal as ascii. A 3D SDF of spheres is shown in [examples/render](https://github.com/AnswerDotAI/gpu.cpp/tree/main/examples/render) and a shadertoy-like live-reloading example is in [examples/shadertui](https://github.com/AnswerDotAI/gpu.cpp/tree/main/examples/shadertui). - -
- shadertui example animated gif -
- -## Who is gpu.cpp for? - -gpu.cpp is aimed at enabling projects requiring portable on-device GPU computation with minimal implementation complexity and friction. Some example use cases are: - -- Development of GPU algorithms to be run on personal computing devices -- Direct standalone implementations of neural network models -- Physics simulations and simulation environments -- Multimodal applications - audio and video processing -- Offline graphics rendering -- ML inference engines and runtimes -- Parallel compute intensive data processing applications - -Although gpu.cpp is meant for any general purpose GPU computation and not strictly AI, one area we're interested in is pushing the limits exploring the intersection of new algorithms for post-training and on-device compute. - -To date, AI research has primarily been built with CUDA as the privileged first-class target. CUDA has been dominant at large scale training and inference but at the other end of the the spectrum in the world of GPU compute on personal devices, there exists far more heterogeneity in the hardware and software stack. - -GPU compute in this personal device ecosystem has been largely limited to a small group of experts such as game engine developers and engineers working directly on ML compilers or inference runtimes. Along with that, implementing against the Vulkan or even WebGPU API directly tends to be targeted mostly towards infrastructure scale efforts - game engines, production ML inference engines, large software packages. - -We want to make it easier for a broader range of projects to harness the power of GPUs on personal devices. With a small amount of code, we can access the GPU at a low-level, focusing on directly implementing algorithms rather than the scaffolding and tech stack around the GPU. For example, in our AI research there's much to explore with the various forms of dynamic/conditional post-training computation - dynamic use of adapters, sparsity, model compression, realtime multimodal integrations etc. - -gpu.cpp lets us implement and drop-in any algorithm with fine-grained control of data movement and GPU code, and explore outside boundaries of what is supported by existing production-oriented inference runtimes. At the same time we can write code that is portable and immediately usable on a wide variety of and GPU vendors and compute form factors - workstations, laptops, mobile, or even emerging hardware platforms such as AR/VR and robotics. - -## What gpu.cpp is not - -gpu.cpp is meant for developers with some familiarity with C++ and GPU programming. It is not a high-level numerical computing or machine learning framework or inference engine, though it can be used in support of such implementations. - -Second, in spite of the name, WebGPU has native implementations decoupled from the web and the browser. If you find it counterintuitive, watch Elie Michel's excellent talk ["WebGPU is Not Just About the Web"](https://www.youtube.com/watch?v=qHrx41aOTUQ). - -Finally, the focus of gpu.cpp is general-purpose GPU computation rather than rendering/graphics on the GPU, although it can be useful for offline rendering or video processing use cases. We may explore directions with graphics in the future, but for now our focus is GPU compute. - -## Limitations and Upcoming Features - -_Browser Targets_ - In spite of using WebGPU we haven't tested builds targeting the browser yet though this is a short-term priority. - -_Reusable Kernel Library_ - Currently the core library is strictly the operations and types for interfacing with the WebGPU API, with some specific use case example WGSL implementations in `examples/`. Over time, as kernel implementations mature we may migrate some of the reusable operations from specific examples into a small reusable kernel library. - -## Troubleshooting - -If you run into issues building the project, please open an issue. - -## Acknowledgements - -gpu.cpp makes use of: - -- [Dawn](https://dawn.googlesource.com/dawn) as the WebGPU implementation -- [webgpu-dawn-binaries](https://github.com/jspanchu/webgpu-dawn-binaries) by - @jspanchu to build a binary artifact of Dawn. -- [webgpu-distribution](https://github.com/eliemichel/WebGPU-distribution) by - @eliemichel for cmake builds. - -Thanks also to fellow colleagues at Answer.AI team for their support, testing help, and feedback. - -## Discord Community and Contributing - -Join our community in the `#gpu-cpp` channel on the [AnswerDotAI Discord with this invite link](https://discord.gg/zmJVhXsC7f). Feel free to get in touch via X [@austinvhuang](https://twitter.com/austinvhuang) as well. - -Feedback, issues and pull requests are welcome. - -## Code Guidelines for Contributors - -For contributors, here are general rules of thumb regarding the design and -style of the gpu.cpp library: - -Aesthetics - Maximize Leverage and Account for Sources of Friction: +```cpp +auto context = createContext({.enableSPIRV = true}); +auto kernel = createKernel(context, SPIRV{words}, + Bindings{readWrite(output)}); +``` -- In addition to performance, time-to-grok the codebase, compilation time, - number of failure modes for builds are things worth optimizing for. -- Increase the implementation surface area only when there's a clear goal - behind doing so. This maximizes leverage per unit effort, increases - optionality in how the library can be used, and keeps compile times low. -- Taking inspiration from the time-tested horizontal extensibility - of neural network libraries like PyTorch, to a first approximation the library - architecture could be described as a bag of composable functions. -- Design choices general attempt to blend the composability of functional - programming with the performance awareness of data oriented design. +Input must satisfy Dawn/Tint's WebGPU SPIR-V profile. The current pinned Dawn +accepts SPIR-V 1.3; `tests/write42.spvasm` is the executable contract fixture. +This is intentionally a WebGPU-facing compiler contract, not arbitrary Vulkan +SPIR-V passthrough. -Overloads and Templates: +## Using gpu.cpp from CMake -- Prefer value-level types over type-level templates, especially for core - implementation code. It's easy to add a more typesafe templated wrapper - around a value type core implementation. Whereas moving templated core - implementations from comptime to runtime leads to a more significant - refactor. -- For comptime polymorphism, prefer trivial function overloads over templates. - Besides compile time benefits, this reasoning about which version of a - function is being called becomes explicit and scanable in the codebase. +After staging Dawn, add this repository as a subdirectory and link the +interface target: -Avoid Encapsulation and Methods: +```cmake +add_subdirectory(path/to/gpu.cpp) +target_link_libraries(my_program PRIVATE gpucpp) +``` -- To build systems effectively, we need to construct them out of subsystems for - which the behavior is known and thereby composable and predictable. - Therefore, we prefer transparency and avoid encapsulation. Don't use abstract - classes as interface specifications, the library and its function signatures - is the interface. -- Use struct as a default over class unless there's a clear reason otherwise. -- Instead of methods, pass the "owning object" object as a reference to a - function. In general this convention can perform any operation that a method - can, but with more flexibility and less coupling. Using mutating functions - generalizes more cleanly to operations that have side effects on more than - one parameter, whereas methods priveledge the the owning class, treating the - single variable case as a special case and making it harder to generalize to - multiple parameters. -- Methods are usually only used for constructor/destructor/operator priveledged - cases. -- For operations requesting GPU resources and more complex initialization, use - factory functions following the `create[X]` convention - createTensor, - createKernel, createContext etc. -- Use (as-trivial-as-possible) constructors for simple supporting types (mostly - providing metadata for a dispatch) Shape, KernelCode, etc. +See [DEV.md](DEV.md) for the dependency layout and update process, and +[CHANGELOG.md](CHANGELOG.md) for release notes. + +## Python + +The optional pybind11 module is built by default when gpu.cpp is the top-level +CMake project. It accepts C-contiguous NumPy arrays and preserves tensor shape +and `float16`, `float32`, or `int32` dtype information: + +```python +import numpy as np +import gpu_cpp as gpu + +context = gpu.Context() +values = np.arange(4, dtype=np.float32) +gpu_values = gpu.tensor(context, values) +gpu_output = gpu.create_tensor(context, values.shape, gpu.f32) +kernel = gpu.create_kernel( + context, gpu.WGSL(source, workgroup_size=[4, 1, 1]), + [gpu.read(gpu_values), gpu.read_write(gpu_output)]) +dispatched = gpu.dispatch_kernel(context, kernel) +gpu.wait(context, dispatched) +result = gpu.to_numpy(context, gpu_output) +``` -Ownership: +Set `GPUCPP_BUILD_PYTHON=OFF` when embedding gpu.cpp in a CMake project that +does not need the module. + +## Browser + +The same `gpucpp` CMake target uses Dawn's Emdawnwebgpu port under Emscripten. +Browser builds accept WGSL, use the browser-selected WebGPU adapter, and use +JSPI for gpu.cpp's synchronous-looking `wait()` calls. SPIR-V is intentionally +native-only because browser WebGPU does not accept it. + +The build also produces `gpu_cpp_web.mjs` and its Wasm file. Its persistent +context accepts ordinary JavaScript typed arrays; `readWrite` arrays are +updated in place: + +```js +import createGpuCpp from "./gpu_cpp_web.mjs"; + +const gpu = await createGpuCpp(); +const context = await gpu.createContext(false); // true enables shader-f16 +const input = new Float32Array([1, 2, 3, 4]); +const output = new Float32Array(4); + +await context.run({ + code: wgsl, + workgroupSize: [4, 1, 1], + workgroups: [1, 1, 1], + bindings: [ + {data: input, access: "read"}, + {data: output, access: "readWrite"}, + ], + // parameters: new Uint8Array(...), // optional uniform bytes +}); +context.delete(); +``` -- Prefer stack allocation for ownership, use unique_ptr for ownership when the - heap is needed. Use raw pointers only for non-owning views. Avoid shared_ptr - unless there's a clear rationale for shared ownership. -- Use pools as a single point of control to manage sets of resources. Consider - incorporating a pool in Context if the resource is universal enough to the - overall API. +Bindings accept `Float32Array`, `Int32Array`, and either `Float16Array` or raw +IEEE-754 half bits in a `Uint16Array`. Calls on one context must be awaited +sequentially; validation and runtime failures reject the returned Promise. -Separating Resource Acquisition from Hot Paths: +## Scope -- In general, resource acquisition should be done ahead of time from the hot - paths of the application. This is to ensure that the hot paths are as fast as - possible and don't have to deal with resource allocation or data movement. -- Operations in the API should be implemented with a use in mind - typically - either ahead-of-time resource preparation/acquisition, hot-paths, or - non-critical testing/observability code. +gpu.cpp targets general-purpose WebGPU compute. It is not a tensor +framework, graph compiler, or rendering engine. The goal is a concise layer for +projects that need direct control over shaders, bindings, dispatch, and data +movement without carrying raw WebGPU setup throughout their code. diff --git a/bindings/haskell/CHANGELOG.md b/bindings/haskell/CHANGELOG.md deleted file mode 100644 index d20679e..0000000 --- a/bindings/haskell/CHANGELOG.md +++ /dev/null @@ -1,5 +0,0 @@ -# Revision history for gpu-cpp - -## 0.1.0.0 -- 2024-12-28 - -* First version. diff --git a/bindings/haskell/Makefile b/bindings/haskell/Makefile deleted file mode 100644 index 7ca37a0..0000000 --- a/bindings/haskell/Makefile +++ /dev/null @@ -1,3 +0,0 @@ -all: - cabal configure --extra-include-dirs=$(PWD)/../.. --extra-include-dirs=$(PWD)/../../third_party/headers --extra-lib-dirs=$(PWD)/../../third_party/lib - cabal build . diff --git a/bindings/haskell/app/Main.hs b/bindings/haskell/app/Main.hs deleted file mode 100644 index ba1ae6d..0000000 --- a/bindings/haskell/app/Main.hs +++ /dev/null @@ -1,37 +0,0 @@ -module Main where - -import GpuCpp.Types -import GpuCpp -import qualified Data.Vector.Storable as V -import Foreign.C.Types - -main :: IO () -main = do - context <- createContext - input <- createTensor context [12] kf32 - output <- createTensor context [12] kf32 - kernelCode <- createKernelCode - ( - "const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI)\n" <> - "@group(0) @binding(0) var inp: array<{{precision}}>;\n" <> - "@group(0) @binding(1) var out: array<{{precision}}>;\n" <> - "@group(0) @binding(1) var dummy: array<{{precision}}>;\n" <> - "@compute @workgroup_size({{workgroupSize}})\n" <> - "fn main(\n" <> - " @builtin(global_invocation_id) GlobalInvocationID: vec3) {\n" <> - " let i: u32 = GlobalInvocationID.x;\n" <> - " if (i < arrayLength(&inp)) {\n" <> - " let x: f32 = inp[i];\n" <> - " out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR \n" <> - " * (x + .044715 * x * x * x))), x, x > 10.0);\n" <> - " }\n" <> - "}\n" - ) - 256 - kf32 - kernel <- createKernel context kernelCode [input, output] [0,0] [12,1,1] - toGpu context (V.fromList [1 :: CFloat,2,3,4,1,2,3,4,1,2,3,4]) input - async <- dispatchKernel context kernel - wait context async - vec <- toCpu context output :: IO (V.Vector CFloat) - print vec diff --git a/bindings/haskell/gpu-cpp.cabal b/bindings/haskell/gpu-cpp.cabal deleted file mode 100644 index 90cb4fa..0000000 --- a/bindings/haskell/gpu-cpp.cabal +++ /dev/null @@ -1,49 +0,0 @@ -cabal-version: 3.0 -name: gpu-cpp -version: 0.1.0.0 -license: BSD-3-Clause -author: Junji Hashimoto -maintainer: junji.hashimoto@gmail.com -category: Math -build-type: Simple - -extra-doc-files: CHANGELOG.md - -common warnings - ghc-options: -Wall - -library - import: warnings - exposed-modules: GpuCpp - , GpuCpp.Types - build-depends: base ^>=4.18.1.0 - , inline-c - , inline-c-cpp - , containers - , template-haskell - , safe-exceptions - , vector - hs-source-dirs: src - default-language: Haskell2010 - ghc-options: -optcxx-std=c++17 - extra-libraries: webgpu_dawn - -executable gpu-cpp - import: warnings - main-is: Main.hs - build-depends: base ^>=4.18.1.0 - , gpu-cpp - , vector - hs-source-dirs: app - default-language: Haskell2010 - -test-suite gpu-cpp-test - import: warnings - default-language: Haskell2010 - type: exitcode-stdio-1.0 - hs-source-dirs: test - main-is: Main.hs - build-depends: base ^>=4.18.1.0 - , gpu-cpp - , vector - , hspec diff --git a/bindings/haskell/src/GpuCpp.hs b/bindings/haskell/src/GpuCpp.hs deleted file mode 100644 index 2177ecf..0000000 --- a/bindings/haskell/src/GpuCpp.hs +++ /dev/null @@ -1,207 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE QuasiQuotes #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE FlexibleInstances #-} - -module GpuCpp where - -import qualified Language.C.Inline.Cpp as C -import qualified Language.C.Inline.Cpp.Unsafe as C -import qualified Language.C.Inline.Context as C -import Foreign.C.String -import Foreign.C.Types -import GHC.Int -import GHC.ForeignPtr(mallocPlainForeignPtrBytes) -import Foreign -import Control.Monad (forM_) -import GpuCpp.Types -import Control.Exception.Safe (bracket) -import qualified Data.Vector.Storable as V - -C.context $ C.cppCtx <> mempty { C.ctxTypesTable = typeTable } - -C.include "" -C.include "" -C.include "" - -[C.emitBlock| -struct GpuAsync { - std::promise promise; - std::future future; - GpuAsync(): future(promise.get_future()){ - } -}; - -gpu::Shape vector_to_shape(const std::vector &dims) { - switch(dims.size()){ - case 1: - return gpu::Shape{(unsigned long)dims[0]}; - break; - case 2: - return gpu::Shape{(unsigned long)dims[0],(unsigned long)dims[1]}; - break; - case 3: - return gpu::Shape{(unsigned long)dims[0],(unsigned long)dims[1],(unsigned long)dims[2]}; - break; - case 4: - return gpu::Shape{(unsigned long)dims[0],(unsigned long)dims[1],(unsigned long)dims[2],(unsigned long)dims[3]}; - break; - case 5: - return gpu::Shape{(unsigned long)dims[0],(unsigned long)dims[1],(unsigned long)dims[2],(unsigned long)dims[3],(unsigned long)dims[4]}; - break; - } - return gpu::Shape{0}; -} -|] - -kf32 :: CInt -kf32 = [C.pure| int { (int)gpu::kf32 } |] - -createContext :: IO (ForeignPtr Context) -createContext = - [C.throwBlock| gpu::Context* { return new gpu::Context(gpu::createContext()); }|] >>= - newForeignPtr - [C.funPtr| void deleteContext(gpu::Context* ptr) { delete ptr; }|] - - -createKernelCode :: String -> CInt -> CInt -> IO (ForeignPtr KernelCode) -createKernelCode kernelString workgroupSize precision = - withCString kernelString $ \pData -> - [C.throwBlock| gpu::KernelCode* { return new gpu::KernelCode($(char* pData), $(int workgroupSize), (gpu::NumType)$(int precision)); }|] >>= - newForeignPtr - [C.funPtr| void deleteKernelCode(gpu::KernelCode* ptr) { delete ptr; }|] - - -dispatchKernel :: ForeignPtr Context -> ForeignPtr Kernel -> IO (ForeignPtr GpuAsync) -dispatchKernel context kernel = - withForeignPtr context $ \c -> - withForeignPtr kernel $ \k -> - [C.throwBlock| GpuAsync* { - auto async = new GpuAsync(); - gpu::dispatchKernel(*$(gpu::Context* c), *$(gpu::Kernel* k), async->promise); - return async; }|] >>= - newForeignPtr - [C.funPtr| void deleteGpuAsync(GpuAsync* ptr) { delete ptr; }|] - -wait :: ForeignPtr Context -> ForeignPtr GpuAsync -> IO () -wait context async = - withForeignPtr context $ \c -> - withForeignPtr async $ \a -> - [C.throwBlock| void { - gpu::wait(*$(gpu::Context* c), $(GpuAsync* a)->future); - }|] - -instance WithVector CInt Int64 where - withVector shape func = - bracket - (do - let len = fromIntegral $ length shape - vec <- [C.throwBlock| std::vector* { - return new std::vector($(int len)); - }|] - ptr <- [C.throwBlock| int64_t* { - return $(std::vector* vec)->data(); - }|] - pokeArray ptr (map fromIntegral shape) - return vec - ) - (\vec -> [C.block| void { delete $(std::vector* vec); }|]) - (\vec -> func vec) - -instance WithVector CInt CSize where - withVector shape func = - bracket - (do - let len = fromIntegral $ length shape - vec <- [C.throwBlock| std::vector* { - return new std::vector($(int len)); - }|] - ptr <- [C.throwBlock| size_t* { - return $(std::vector* vec)->data(); - }|] - pokeArray ptr (map fromIntegral shape) - return vec - ) - (\vec -> [C.block| void { delete $(std::vector* vec); }|]) - (\vec -> func vec) - -instance WithVector (Ptr Tensor) Tensor where - withVector ptrs func = - bracket (do - vec <- [C.throwBlock| std::vector* { return new std::vector(); }|] - forM_ ptrs $ do - \ptr -> [C.throwBlock| void { $(std::vector* vec)->push_back(*$(gpu::Tensor* ptr)); }|] - return vec - ) - (\vec -> [C.block| void { delete $(std::vector* vec); }|]) - (\vec -> func vec) - -withForeignPtrs :: [ForeignPtr a] -> ([Ptr a] -> IO b) -> IO b -withForeignPtrs [] func = func [] -withForeignPtrs (x:xs) func = - withForeignPtr x $ \x' -> - withForeignPtrs xs $ \xs' -> - func (x':xs') - -createKernel :: ForeignPtr Context -> ForeignPtr KernelCode -> [ForeignPtr Tensor] -> [Int] -> [Int] -> IO (ForeignPtr Kernel) -createKernel context kernelCode dataBindings viewOffsets totalWorkgroups = - withForeignPtr context $ \c -> - withForeignPtr kernelCode $ \k -> - withForeignPtrs dataBindings $ \b -> - withVector b $ \b' -> - withVector @CInt (map fromIntegral viewOffsets) $ \v -> - withVector @CInt (map fromIntegral totalWorkgroups) $ \w -> - [C.throwBlock| gpu::Kernel* { - return new gpu::Kernel(gpu::createKernel( - *$(gpu::Context* c), - *$(gpu::KernelCode* k), - $(std::vector* b')->data(), - $(std::vector* b')->size(), - $(std::vector* v)->data(), - vector_to_shape(*$(std::vector* w)))); - }|] >>= - newForeignPtr - [C.funPtr| void deleteKernel(gpu::Kernel* ptr) { delete ptr; }|] - -createTensor :: ForeignPtr Context -> [CInt] -> CInt -> IO (ForeignPtr Tensor) -createTensor context shape dtype = - withVector shape $ \s -> - withForeignPtr context $ \c -> - [C.throwBlock| gpu::Tensor* { - return new gpu::Tensor(gpu::createTensor(*$(gpu::Context* c), vector_to_shape(*$(std::vector* s)), (gpu::NumType)$(int dtype))); - }|] >>= - newForeignPtr - [C.funPtr| void deleteTensor(gpu::Tensor* ptr) { delete ptr; }|] - -createVector :: forall a. Storable a => Int -> IO (V.Vector a) -createVector n = do - ptr <- mallocPlainForeignPtrBytes (n * sizeOf (undefined :: a)) - return $ V.unsafeFromForeignPtr ptr 0 n - -instance GpuStorable CFloat where - toGpu context array tensor = - withForeignPtr context $ \c -> - withForeignPtr tensor $ \t -> - V.unsafeWith array $ \ptr -> - [C.throwBlock| void { - gpu::toGPU(*$(gpu::Context* c), $(float* ptr), *$(gpu::Tensor* t)); - }|] - toCpu context tensor = - withForeignPtr context $ \c -> - withForeignPtr tensor $ \t -> do - (size :: CInt) <- [C.block| int { - size_t u = sizeof(float); - size_t len = $(gpu::Tensor* t)->data.size; - return len/u; - }|] - array <- createVector (fromIntegral size) - V.unsafeWith array $ \ptr -> - [C.throwBlock| void { - gpu::toCPU(*$(gpu::Context* c), *$(gpu::Tensor* t), $(float* ptr), $(int size) * sizeof(float)); - }|] - return array diff --git a/bindings/haskell/src/GpuCpp/Types.hs b/bindings/haskell/src/GpuCpp/Types.hs deleted file mode 100644 index 3905aa7..0000000 --- a/bindings/haskell/src/GpuCpp/Types.hs +++ /dev/null @@ -1,40 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE QuasiQuotes #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE MultiParamTypeClasses #-} - -module GpuCpp.Types where - -import qualified Language.C.Types as C -import qualified Language.Haskell.TH.Lib as TH -import qualified Data.Map as Map -import Foreign -import qualified Data.Vector.Storable as V - -data Context -data Tensor -data Kernel -data KernelCode -data GpuAsync -data StdVector a - -typeTable :: Map.Map C.TypeSpecifier TH.TypeQ -typeTable = Map.fromList [ - (C.TypeName "gpu::Context", [t|Context|]) - , (C.TypeName "gpu::Tensor", [t|Tensor|]) - , (C.TypeName "gpu::Kernel", [t|Kernel|]) - , (C.TypeName "gpu::KernelCode", [t|KernelCode|]) - , (C.TypeName "GpuAsync", [t|GpuAsync|]) - , (C.TypeName "std::vector", [t|StdVector|]) - ] - - -class WithVector a b where - withVector :: [a] -> (Ptr (StdVector b) -> IO c) -> IO c - -class GpuStorable a where - toGpu :: ForeignPtr Context -> V.Vector a -> ForeignPtr Tensor -> IO () - toCpu :: ForeignPtr Context -> ForeignPtr Tensor -> IO (V.Vector a) - diff --git a/bindings/haskell/test/Main.hs b/bindings/haskell/test/Main.hs deleted file mode 100644 index d66e5c1..0000000 --- a/bindings/haskell/test/Main.hs +++ /dev/null @@ -1,49 +0,0 @@ -module Main (main) where - -import Test.Hspec -import GpuCpp.Types -import GpuCpp -import qualified Data.Vector.Storable as V -import Foreign.C.Types - -gelu :: String -gelu= "const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI)\n" <> - "@group(0) @binding(0) var inp: array<{{precision}}>;\n" <> - "@group(0) @binding(1) var out: array<{{precision}}>;\n" <> - "@group(0) @binding(1) var dummy: array<{{precision}}>;\n" <> - "@compute @workgroup_size({{workgroupSize}})\n" <> - "fn main(\n" <> - " @builtin(global_invocation_id) GlobalInvocationID: vec3) {\n" <> - " let i: u32 = GlobalInvocationID.x;\n" <> - " if (i < arrayLength(&inp)) {\n" <> - " let x: f32 = inp[i];\n" <> - " out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR \n" <> - " * (x + .044715 * x * x * x))), x, x > 10.0);\n" <> - " }\n" <> - "}\n" - -main :: IO () -main = do - hspec $ do - describe "toCPU and toGPU" $ do - it "writes and reads back" $ do - context <- createContext - input <- createTensor context [12] kf32 - toGpu context (V.fromList [1 :: CFloat,2,3,4,1,2,3,4,1,2,3,4]) input - output <- toCpu context input :: IO (V.Vector CFloat) - V.toList output `shouldBe` [1,2,3,4,1,2,3,4,1,2,3,4] - describe "call kernel" $ do - it "gelu" $ do - context <- createContext - input <- createTensor context [12] kf32 - output <- createTensor context [12] kf32 - kernelCode <- createKernelCode gelu 256 kf32 - kernel <- createKernel context kernelCode [input, output] [0,0] [12,1,1] - toGpu context (V.fromList [1 :: CFloat,2,3,4,1,2,3,4,1,2,3,4]) input - async <- dispatchKernel context kernel - wait context async - vec <- toCpu context output :: IO (V.Vector CFloat) - V.toList (V.zipWith (\a b -> abs (a - b)) - vec - (V.fromList [0.841192,1.9545977,2.9963627,3.9999297,0.841192,1.9545977,2.9963627,3.9999297,0.841192,1.9545977,2.9963627,3.9999297])) - `shouldSatisfy` all (< 0.001) diff --git a/bindings/python/Makefile b/bindings/python/Makefile deleted file mode 100644 index 78e0b58..0000000 --- a/bindings/python/Makefile +++ /dev/null @@ -1,25 +0,0 @@ -CXX=clang++ -PYTHON=python3 -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source - -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif - -FLAGS=-shared -fPIC -std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib -lwebgpu_dawn \ - `python3 -m pybind11 --includes` \ - `python3-config --includes --ldflags` - -SUFFIX=$(shell $(PYTHON)-config --extension-suffix) - -gpu_cpp$(SUFFIX): gpu_cpp.cpp - $(CXX) $(FLAGS) -o $@ $< - -test: test_gpu_cpp.py gpu_cpp$(SUFFIX) - $(PYTHON) test_gpu_cpp.py - -.PHONY: test diff --git a/bindings/python/gpu_cpp.cpp b/bindings/python/gpu_cpp.cpp index 8bd762d..e292097 100644 --- a/bindings/python/gpu_cpp.cpp +++ b/bindings/python/gpu_cpp.cpp @@ -1,111 +1,215 @@ #include "gpu.hpp" -#include -#include -#include -using namespace gpu; - -#include #include +#include #include +#include +#include +#include +#include +#include + namespace py = pybind11; +using namespace gpu; + +namespace { -Shape vector_to_shape(const std::vector &dims) { - switch(dims.size()){ - case 1: - return Shape{(unsigned long)dims[0]}; - break; - case 2: - return Shape{(unsigned long)dims[0],(unsigned long)dims[1]}; - break; - case 3: - return Shape{(unsigned long)dims[0],(unsigned long)dims[1],(unsigned long)dims[2]}; - break; - case 4: - return Shape{(unsigned long)dims[0],(unsigned long)dims[1],(unsigned long)dims[2],(unsigned long)dims[3]}; - break; - case 5: - return Shape{(unsigned long)dims[0],(unsigned long)dims[1],(unsigned long)dims[2],(unsigned long)dims[3],(unsigned long)dims[4]}; - break; +Shape shape(const std::vector &dimensions) { + if (dimensions.size() > Shape::kMaxRank) + throw std::invalid_argument("tensor rank exceeds 8"); + Shape result; + result.rank = dimensions.size(); + for (size_t i = 0; i < dimensions.size(); ++i) { + if (dimensions[i] < 0) + throw std::invalid_argument("tensor dimensions cannot be negative"); + result.data[i] = dimensions[i]; } - return Shape{0}; + return result; } -Context* py_createContext() { - return new Context(createContext()); +std::vector dimensions(const Shape &shape) { + std::vector result(shape.rank); + for (size_t i = 0; i < shape.rank; ++i) result[i] = shape[i]; + return result; } -KernelCode* py_createKernelCode(const std::string &pData, size_t workgroupSize, int precision) { - return new KernelCode(pData, workgroupSize, (NumType)precision); +void requireContiguous(const py::array &array) { + if (!(array.flags() & py::array::c_style)) + throw std::invalid_argument("array must be C-contiguous"); } -Kernel py_createKernel(Context *ctx, const KernelCode *code, - // const Tensor *dataBindings, size_t numTensors, - const py::list& dataBindings_py, - // const size_t *viewOffsets, - const py::list& viewOffsets_py, - const std::vector &totalWorkgroups){ - std::vector bindings; - for (auto item : dataBindings_py) { - bindings.push_back(item.cast()); - } - std::vector viewOffsets; - for (auto item : viewOffsets_py) { - viewOffsets.push_back(item.cast()); +NumType numType(const py::array &array) { + if (array.dtype().is(py::dtype::of())) return kf32; + if (array.dtype().is(py::dtype::of())) return ki32; + if (array.dtype().is(py::dtype("float16"))) return kf16; + throw std::invalid_argument("array dtype must be float16, float32, or int32"); +} + +py::dtype dtype(NumType type) { + switch (type) { + case kf16: return py::dtype("float16"); + case kf32: return py::dtype::of(); + case ki32: return py::dtype::of(); } - return createKernel(*ctx, *code, bindings.data(), bindings.size(), viewOffsets.data(), vector_to_shape(totalWorkgroups)); + throw std::invalid_argument("unknown numeric type"); } -Tensor* py_createTensor(Context *ctx, const std::vector &dims, int dtype) { - return new Tensor(createTensor(*ctx, vector_to_shape(dims), (NumType)dtype)); +template +std::span input(const py::buffer_info &buffer) { + return {static_cast(buffer.ptr), + static_cast(buffer.size)}; } -py::array_t py_toCPU_float(Context *ctx, Tensor* tensor) { - auto result = py::array_t(tensor->data.size/sizeof(float)); - py::buffer_info buf = result.request(); - toCPU(*ctx, *tensor, static_cast(buf.ptr), tensor->data.size); - return result; +Tensor tensor(Context &context, const py::array &array) { + requireContiguous(array); + const auto type = numType(array); + const auto buffer = array.request(); + const auto tensorShape = shape(buffer.shape); + switch (type) { + case kf16: + return createTensor(context, tensorShape, type, input(buffer)); + case kf32: + return createTensor(context, tensorShape, type, input(buffer)); + case ki32: + return createTensor(context, tensorShape, type, input(buffer)); + } + throw std::invalid_argument("unknown numeric type"); } +void upload(Context &context, const Tensor &tensor, const py::array &array) { + requireContiguous(array); + if (numType(array) != tensor.type) + throw std::invalid_argument("array and tensor dtypes differ"); + const auto buffer = array.request(); + switch (tensor.type) { + case kf16: return toGPU(context, tensor, input(buffer)); + case kf32: return toGPU(context, tensor, input(buffer)); + case ki32: return toGPU(context, tensor, input(buffer)); + } + throw std::invalid_argument("unknown numeric type"); +} -void py_toGPU_float(Context *ctx, py::array_t array, Tensor *tensor) { - py::buffer_info buf = array.request(); - float *ptr = static_cast(buf.ptr); - toGPU(*ctx, ptr, *tensor); +template +void download(Context &context, const Tensor &tensor, const py::buffer_info &buffer) { + auto future = toCPU( + context, tensor, + std::span(static_cast(buffer.ptr), static_cast(buffer.size))); + py::gil_scoped_release release; + wait(context, future); } -struct GpuAsync { - std::promise promise; - std::future future ; - GpuAsync(): future(promise.get_future()){ +py::array numpy(Context &context, const Tensor &tensor) { + py::array result(dtype(tensor.type), dimensions(tensor.shape)); + const auto buffer = result.request(); + switch (tensor.type) { + case kf16: download(context, tensor, buffer); break; + case kf32: download(context, tensor, buffer); break; + case ki32: download(context, tensor, buffer); break; } + return result; +} + +struct PythonFuture { + explicit PythonFuture(gpu::Future value) : value(std::move(value)) {} + PythonFuture(const PythonFuture &) = delete; + PythonFuture &operator=(const PythonFuture &) = delete; + PythonFuture(PythonFuture &&) = default; + PythonFuture &operator=(PythonFuture &&) = default; + gpu::Future value; }; -GpuAsync* py_dispatchKernel(Context *ctx, Kernel kernel) { - auto async = new GpuAsync(); - dispatchKernel(*ctx, kernel, async->promise); - return async; -} +} // namespace -void py_wait(Context *ctx, GpuAsync* async) { - wait(*ctx, async->future); -} +PYBIND11_MODULE(gpu_cpp, module) { + module.doc() = "Native WebGPU compute with gpu.cpp"; + + py::enum_(module, "NumType") + .value("f16", kf16) + .value("f32", kf32) + .value("i32", ki32) + .export_values(); + + py::class_(module, "Context") + .def(py::init([](bool shaderF16, bool spirv) { + ContextOptions options{.enableSPIRV = spirv}; + if (shaderF16) + options.requiredFeatures = {wgpu::FeatureName::ShaderF16}; + py::gil_scoped_release release; + return std::make_unique(createContext(options)); + }), + py::kw_only(), py::arg("shader_f16") = false, + py::arg("spirv") = false); + + py::class_(module, "WGSL") + .def(py::init([](std::string code, + const std::vector &workgroupSize, + NumType precision) { + return WGSL(std::move(code), shape(workgroupSize), precision); + }), + py::arg("code"), py::arg("workgroup_size") = + std::vector{256, 1, 1}, + py::arg("precision") = kf32) + .def_readwrite("code", &WGSL::code) + .def_readwrite("label", &WGSL::label) + .def_readwrite("entry_point", &WGSL::entryPoint) + .def_property_readonly("workgroup_size", + [](const WGSL &shader) { + return dimensions(shader.workgroupSize); + }); + + py::class_(module, "SPIRV") + .def(py::init([](std::vector words) { + return SPIRV{std::move(words)}; + }), + py::arg("words")) + .def_readwrite("words", &SPIRV::code) + .def_readwrite("label", &SPIRV::label) + .def_readwrite("entry_point", &SPIRV::entryPoint); + + py::class_(module, "Tensor") + .def_property_readonly("shape", + [](const Tensor &tensor) { + return dimensions(tensor.shape); + }) + .def_readonly("dtype", &Tensor::type); + py::class_(module, "Binding"); + py::class_(module, "Kernel"); + py::class_(module, "Future"); -PYBIND11_MODULE(gpu_cpp, m) { - m.doc() = "gpu.cpp plugin"; - py::class_(m, "Context"); - py::class_(m, "Tensor"); - py::class_>(m, "Kernel"); - py::class_(m, "KernelCode"); - py::class_(m, "GpuAsync"); - m.def("create_context", &py_createContext, py::return_value_policy::take_ownership); - m.def("create_tensor", &py_createTensor, py::return_value_policy::take_ownership); - m.def("create_kernel", &py_createKernel); - m.def("create_kernel_code", &py_createKernelCode, py::return_value_policy::take_ownership); - m.def("dispatch_kernel", &py_dispatchKernel, py::return_value_policy::take_ownership); - m.def("wait", &py_wait, "Wait for GPU"); - m.def("to_cpu_float", &py_toCPU_float); - m.def("to_gpu_float", &py_toGPU_float); - m.attr("kf32") = (int)kf32; + module.def("create_tensor", + [](Context &context, + const std::vector &tensorShape, NumType type) { + return createTensor(context, shape(tensorShape), type); + }, + py::arg("context"), py::arg("shape"), py::arg("dtype")); + module.def("tensor", &tensor, py::arg("context"), py::arg("array")); + module.def("read", py::overload_cast(&read), + py::arg("tensor")); + module.def("read_write", py::overload_cast(&readWrite), + py::arg("tensor")); + module.def( + "create_kernel", + [](Context &context, const Shader &shader, + const std::vector &bindings, + const std::vector &workgroups, py::bytes parameters) { + const std::string bytes = parameters; + const auto *data = reinterpret_cast(bytes.data()); + return createKernel(context, shader, std::span(bindings), + shape(workgroups), {data, bytes.size()}); + }, + py::arg("context"), py::arg("shader"), py::arg("bindings"), + py::arg("workgroups") = std::vector{1, 1, 1}, + py::arg("parameters") = py::bytes()); + module.def("dispatch_kernel", + [](Context &context, const Kernel &kernel) { + return PythonFuture(dispatchKernel(context, kernel)); + }); + module.def("wait", [](Context &context, PythonFuture &future) { + py::gil_scoped_release release; + wait(context, future.value); + }); + module.def("to_gpu", &upload, py::arg("context"), py::arg("tensor"), + py::arg("array")); + module.def("to_numpy", &numpy, py::arg("context"), py::arg("tensor")); } diff --git a/bindings/python/test_gpu_cpp.py b/bindings/python/test_gpu_cpp.py index ad50c6a..431c10c 100644 --- a/bindings/python/test_gpu_cpp.py +++ b/bindings/python/test_gpu_cpp.py @@ -1,39 +1,45 @@ -import gpu_cpp as gpu +import struct + import numpy as np -ctx = gpu.create_context() - -N = 12 - -input = gpu.create_tensor(ctx, [N], gpu.kf32) -output = gpu.create_tensor(ctx, [N], gpu.kf32) -kernel_code = gpu.create_kernel_code( - """ - const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI) - @group(0) @binding(0) var inp: array<{{precision}}>; - @group(0) @binding(1) var out: array<{{precision}}>; - @group(0) @binding(1) var dummy: array<{{precision}}>; - @compute @workgroup_size({{workgroupSize}}) - fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&inp)) { - let x: f32 = inp[i]; - out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR - * (x + .044715 * x * x * x))), x, x > 10.0); - } - } - """, - 256, - gpu.kf32 - ) - -kernel = gpu.create_kernel(ctx, kernel_code, [input, output], [0,0], [12,1,1]) - -gpu.to_gpu_float(ctx, np.array([1,2,3,4,1,2,3,4,1,2,3,4],np.float32), input) - -gpu_async = gpu.dispatch_kernel(ctx, kernel); - -gpu.wait(ctx, gpu_async); - -print(gpu.to_cpu_float(ctx, output)) +import gpu_cpp as gpu + + +scale = r''' +struct Params { factor: f32 } +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size({{workgroupSize}}) +fn main(@builtin(global_invocation_id) id: vec3) { + if (id.x < arrayLength(&input)) { + output[id.x] = input[id.x] * params.factor; + } +} +''' + + +# Upload NumPy data, bind access explicitly, dispatch, and recover its shape and dtype. +context = gpu.Context() +values = np.arange(12, dtype=np.float32) +gpu_values = gpu.tensor(context, values) +gpu_output = gpu.create_tensor(context, values.shape, gpu.f32) +shader = gpu.WGSL(scale, workgroup_size=[4, 1, 1]) +bindings = [gpu.read(gpu_values), gpu.read_write(gpu_output)] +kernel = gpu.create_kernel(context, shader, bindings, workgroups=[3, 1, 1], parameters=struct.pack(' + +#include +#include +#include +#include +#include +#include + +namespace { + +using emscripten::val; +using namespace gpu; + +Shape shape(const val &value, const char *name) { + if (value["length"].as() != 3) + throw std::invalid_argument(std::string(name) + " must have three dimensions"); + return {value[0].as(), value[1].as(), value[2].as()}; +} + +NumType numType(const val &array) { + const auto name = array["constructor"]["name"].as(); + if (name == "Float32Array") return kf32; + if (name == "Int32Array") return ki32; + if (name == "Float16Array" || name == "Uint16Array") return kf16; + throw std::invalid_argument( + "binding data must be Float16Array, Uint16Array, Float32Array, or Int32Array"); +} + +std::vector bytes(const val &array) { + const auto length = array["byteLength"].as(); + std::vector result(length); + auto source = val::global("Uint8Array") + .new_(array["buffer"], array["byteOffset"], length); + auto destination = val(emscripten::typed_memory_view( + length, reinterpret_cast(result.data()))); + destination.call("set", source); + return result; +} + +void copy(const void *source, size_t length, const val &array) { + auto sourceView = val(emscripten::typed_memory_view( + length, static_cast(source))); + auto destination = val::global("Uint8Array") + .new_(array["buffer"], array["byteOffset"], length); + destination.call("set", sourceView); +} + +class BrowserContext { + Context context; + bool running = false; + + struct RunGuard { + bool &running; + explicit RunGuard(bool &running) : running(running) { + if (std::exchange(running, true)) + throw std::runtime_error("a gpu.cpp run is already in progress"); + } + ~RunGuard() { running = false; } + }; + + template + void download(const Tensor &tensor, const val &array) { + std::vector values(size(tensor.shape)); + auto future = toCPU(context, tensor, values); + wait(context, future); + copy(values.data(), values.size() * sizeof(T), array); + } + +public: + explicit BrowserContext(bool shaderF16) { + ContextOptions options; + if (shaderF16) options.requiredFeatures = {wgpu::FeatureName::ShaderF16}; + context = createContext(options); + } + + void run(const val &spec) { + RunGuard guard(running); + const auto code = spec["code"].as(); + const auto workgroupSize = shape(spec["workgroupSize"], "workgroupSize"); + const auto workgroups = shape(spec["workgroups"], "workgroups"); + const auto specs = spec["bindings"]; + + std::vector tensors; + std::vector bindings; + std::vector arrays; + std::vector downloads; + const auto count = specs["length"].as(); + tensors.reserve(count); + bindings.reserve(count); + arrays.reserve(count); + downloads.reserve(count); + + for (size_t i = 0; i < count; ++i) { + const auto binding = specs[i]; + const auto array = binding["data"]; + const auto type = numType(array); + const auto data = bytes(array); + if (data.empty() || data.size() % sizeBytes(type)) + throw std::invalid_argument("binding data has an invalid byte length"); + tensors.push_back(createTensor(context, {data.size() / sizeBytes(type)}, type)); + context.queue.WriteBuffer(tensors.back().buffer, 0, data.data(), data.size()); + + const auto access = binding["access"].as(); + if (access == "read") { + bindings.push_back(read(tensors.back())); + downloads.push_back(false); + } else if (access == "readWrite") { + bindings.push_back(readWrite(tensors.back())); + downloads.push_back(true); + } else { + throw std::invalid_argument("binding access must be read or readWrite"); + } + arrays.push_back(array); + } + + std::vector parameters; + const auto parameterValue = spec["parameters"]; + if (!parameterValue.isUndefined() && !parameterValue.isNull()) + parameters = bytes(parameterValue); + + auto kernel = createKernel(context, WGSL{code, workgroupSize}, bindings, + workgroups, parameters); + auto dispatched = dispatchKernel(context, kernel); + wait(context, dispatched); + + for (size_t i = 0; i < tensors.size(); ++i) { + if (!downloads[i]) continue; + switch (tensors[i].type) { + case kf16: download(tensors[i], arrays[i]); break; + case kf32: download(tensors[i], arrays[i]); break; + case ki32: download(tensors[i], arrays[i]); break; + } + } + } +}; + +BrowserContext *createBrowserContext(bool shaderF16) { + return new BrowserContext(shaderF16); +} + +} // namespace + +EMSCRIPTEN_BINDINGS(gpu_cpp_web) { + emscripten::class_("Context") + .function("run", &BrowserContext::run, emscripten::async()); + emscripten::function("createContext", &createBrowserContext, + emscripten::allow_raw_pointers(), emscripten::async()); +} diff --git a/build.py b/build.py deleted file mode 100644 index ffb5e0d..0000000 --- a/build.py +++ /dev/null @@ -1,32 +0,0 @@ -# Dictionary of header files and their relative paths -header_files = { - "#include \"webgpu/webgpu.h\"": "third_party/headers/webgpu/webgpu.h", - "#include \"numeric_types/half.hpp\"": "numeric_types/half.hpp", - "#include \"utils/logging.hpp\"": "utils/logging.hpp" -} - -def main(): - # File paths - source_file_path = "gpu.hpp" - output_file_path = "build/gpu.hpp" - - # Open source file and read contents - with open(source_file_path, "r") as source: - file_contents = source.read() - - # Ergodic over header files - for key, value in header_files.items(): - - # Replace header files - with open(value, "r") as header_file: - header_file_contents = header_file.read() - file_contents = file_contents.replace(key, header_file_contents) - - - # Open output file - with open(output_file_path, "w") as output: - # Write contents to output file - output.write(file_contents) - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/cmake/example.cmake b/cmake/example.cmake deleted file mode 100644 index eba8e7c..0000000 --- a/cmake/example.cmake +++ /dev/null @@ -1,68 +0,0 @@ -set(CMAKE_EXPORT_COMPILE_COMMANDS ON) # export compile_commands.json to use with - # LSP -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED ON) - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Include file finding utility script -include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/find_gpu.cmake") - -# Check if the file exists in the current directory -find_project_root(${CMAKE_CURRENT_SOURCE_DIR} ${FILENAME} - TARGET_FILE_PATH) -if("${TARGET_FILE_PATH}" STREQUAL "") - find_project_root(${FILEPATH_CURRENT_DIR} ${FILENAME} - TARGET_FILE_PATH) - if("${TARGET_FILE_PATH}" STREQUAL "") - message( - FATAL_ERROR - "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../" - ) - endif() -endif() - -# Ensure the build type is set -if(NOT CMAKE_BUILD_TYPE) - set(CMAKE_BUILD_TYPE - Release - CACHE STRING "Choose the type of build: Debug or Release" FORCE) -endif() - -# Define architecture and build type directories or file names -if(CMAKE_SIZEOF_VOID_P EQUAL 8) - set(ARCH "x64") -else() - set(ARCH "x86") -endif() - -if(CMAKE_BUILD_TYPE STREQUAL "Debug") - set(BUILD_TYPE "Debug") -else() - set(BUILD_TYPE "Release") -endif() - -if(NOT TARGET gpu) - message(STATUS "GPU_LIB not found") - include("${TARGET_FILE_PATH}/cmake/webgpu.cmake") - include("${TARGET_FILE_PATH}/cmake/gpu.cmake") -endif() - -add_executable(${PROJECT_NAME} run.cpp) -target_link_libraries(${PROJECT_NAME} PRIVATE gpu) -target_link_libraries(${PROJECT_NAME} PRIVATE wgpu) -target_link_libraries(${PROJECT_NAME} PRIVATE webgpu) - -if(WIN32) - # Ensure DLL is copied if on Windows - add_custom_command( - TARGET ${PROJECT_NAME} - POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different ${DLL_PATH} - $) -endif() diff --git a/cmake/find_gpu.cmake b/cmake/find_gpu.cmake deleted file mode 100644 index b6b7dad..0000000 --- a/cmake/find_gpu.cmake +++ /dev/null @@ -1,30 +0,0 @@ -# file name to find -set(FILENAME "gpu.hpp") - -# Function to check for file existence up the directory hierarchy -function(find_project_root current_dir filename result_var) - set(found FALSE) # Flag to indicate if the file is found - set(current_check_dir "${current_dir}") # Start from the given directory - # using 1 is jsut to supress the cmane-format warning - foreach(i RANGE 0 2 1) - set(filepath "${current_check_dir}/${filename}") - - if(EXISTS "${filepath}") - set(${result_var} - "${current_check_dir}" - PARENT_SCOPE) - set(found TRUE) - break() - endif() - - # Move one level up - get_filename_component(current_check_dir "${current_check_dir}" - DIRECTORY) - endforeach() - - if(NOT found) - set(${result_var} - "" - PARENT_SCOPE) # Set to empty if not found - endif() -endfunction() diff --git a/cmake/gpu.cmake b/cmake/gpu.cmake deleted file mode 100644 index 08db244..0000000 --- a/cmake/gpu.cmake +++ /dev/null @@ -1,69 +0,0 @@ -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Include file finding utility script -include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/find_gpu.cmake") - -# Check if the file exists in the current directory -find_project_root(${CMAKE_CURRENT_SOURCE_DIR} ${FILENAME} TARGET_FILE_PATH) -if("${TARGET_FILE_PATH}" STREQUAL "") - find_project_root(${FILEPATH_CURRENT_DIR} ${FILENAME} TARGET_FILE_PATH) - if("${TARGET_FILE_PATH}" STREQUAL "") - message( - FATAL_ERROR - "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../" - ) - endif() -endif() - -# Define architecture and build type directories or file names -if(CMAKE_SIZEOF_VOID_P EQUAL 8) - set(ARCH "x64") -else() - set(ARCH "x86") -endif() - -if(CMAKE_BUILD_TYPE STREQUAL "Debug") - set(BUILD_TYPE "Debug") -else() - set(BUILD_TYPE "Release") -endif() - -add_library(webgpulib SHARED IMPORTED) -add_library(gpu INTERFACE) -add_library(wgpu INTERFACE) -add_dependencies(gpu webgpulib) -# Define the header-only library -target_include_directories(gpu INTERFACE ${TARGET_FILE_PATH}) - -# Add headers webgpu.h -target_include_directories(wgpu - INTERFACE ${TARGET_FILE_PATH}/third_party/headers) -include(ExternalProject) - -set(DAWN_EXT_PREFIX "${TARGET_FILE_PATH}/third_party/local/dawn") - -ExternalProject_Add( - dawn_project - PREFIX ${DAWN_EXT_PREFIX} - GIT_REPOSITORY "https://dawn.googlesource.com/dawn" - GIT_TAG "main" - SOURCE_DIR "${DAWN_EXT_PREFIX}/source" - BINARY_DIR "${DAWN_EXT_PREFIX}/build" - INSTALL_DIR "${DAWN_EXT_PREFIX}/install" - GIT_SUBMODULES "" - # setting cmake args doesn't work and I don't know why - CONFIGURE_COMMAND - ${CMAKE_COMMAND} -S ${DAWN_EXT_PREFIX}/source -B - ${DAWN_EXT_PREFIX}/build -DDAWN_FETCH_DEPENDENCIES=ON - -DDAWN_ENABLE_INSTALL=ON -DDAWN_BUILD_MONOLITHIC_LIBRARY=ON - -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} -G ${CMAKE_GENERATOR} - INSTALL_COMMAND ${CMAKE_COMMAND} --install . --prefix - ${DAWN_EXT_PREFIX}/install - LOG_INSTALL ON) -find_library(LIBDAWN dawn PATHS "${DAWN_EXT_PREFIX}/install/lib") -target_link_libraries(webgpulib INTERFACE ${LIBDAWN}) diff --git a/cmake/webgpu.cmake b/cmake/webgpu.cmake deleted file mode 100644 index c63f1e2..0000000 --- a/cmake/webgpu.cmake +++ /dev/null @@ -1,61 +0,0 @@ -# Specify the filename to search for -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Include file finding utility script -include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/find_gpu.cmake") - -# Check if the file exists in the current directory -find_project_root(${CMAKE_CURRENT_SOURCE_DIR} ${FILENAME} TARGET_FILE_PATH) -if("${TARGET_FILE_PATH}" STREQUAL "") - find_project_root(${FILEPATH_CURRENT_DIR} ${FILENAME} TARGET_FILE_PATH) - if("${TARGET_FILE_PATH}" STREQUAL "") - message( - FATAL_ERROR - "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../" - ) - endif() -endif() - -include(FetchContent) - -set(FETCHCONTENT_BASE_DIR "${TARGET_FILE_PATH}/third_party/fetchcontent") -set(WEBGPU_DIST_LOCAL_PATH - "${TARGET_FILE_PATH}/third_party/local/WebGPU-distribution") - -if(USE_LOCAL_LIBS) - set(WEBGPU_DIST_GIT_REPO ${WEBGPU_DIST_LOCAL_PATH}) - message(STATUS "Using local WebGPU distribution: ${WEBGPU_DIST_LOCAL_PATH}") -else() - set(WEBGPU_DIST_GIT_REPO - "https://github.com/eliemichel/WebGPU-distribution") -endif() - -option(WEBGPU_TAG "WebGPU distribution tag to use") -if(NOT WEBGPU_TAG) - set(WEBGPU_TAG "dawn") -endif() -message(STATUS "Using WebGPU distribution tag: ${WEBGPU_TAG}") - -if(WEBGPU_TAG STREQUAL "dawn") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DWEBGPU_BACKEND_DAWN") - # use specific commit set(WEBGPU_TAG - # "1025b977e1927b6d0327e67352f90feb4bcf8274") set(WEBGPU_TAG - # "acf972b7b909f52e183bdae3971b93bb13d4a29e") - # add_compile_options(-UABSL_INTERNAL_AT_LEAST_CXX20) set(CMAKE_CXX_FLAGS - # "${CMAKE_CXX_FLAGS} -UABSL_INTERNAL_AT_LEAST_CXX20") - message(STATUS "Using Dawn backend") -endif() - -FetchContent_Declare( - webgpu - GIT_REPOSITORY ${WEBGPU_DIST_GIT_REPO} - GIT_TAG ${WEBGPU_TAG} - GIT_SHALLOW TRUE) -FetchContent_MakeAvailable(webgpu) diff --git a/examples/Makefile b/examples/Makefile deleted file mode 100644 index 3036e22..0000000 --- a/examples/Makefile +++ /dev/null @@ -1,66 +0,0 @@ -# List of targets (folders in your examples directory) -TARGETS := float16 gpu_puzzles hello_world matmul physics render shadertui transpose - -GPUCPP ?= $(shell pwd)/.. -CXX=clang++ -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -BUILD ?= debug - -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif - -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib -LFLAGS=-ldl -lwebgpu_dawn - -.PHONY: default all_release all_debug dawnlib run_setup check-python -.PHONY: $(addsuffix _release, $(TARGETS)) -.PHONY: $(addsuffix _debug, $(TARGETS)) -.PHONY: clean-all $(addprefix clean_, $(TARGETS)) - -all: all_$(BUILD) - -all_release: $(addsuffix _release, $(TARGETS)) - -all_debug: $(addsuffix _debug, $(TARGETS)) - -define BUILD_RULES -$(1)_release: dawnlib - mkdir -p $(1)/build - $(CXX) $(FLAGS) $(1)/run.cpp $(LFLAGS) -DNDEBUG -o $(1)/build/$(1) - -$(1)_debug: dawnlib - mkdir -p $(1)/build - $(CXX) $(FLAGS) $(1)/run.cpp $(LFLAGS) -o $(1)/build/$(1) - -$(1): $(1)_$(BUILD) -endef - -define RUN_RULES -run_$(1): $(1) - # note directory context is important in the case of shadertui which loads the shader from disk - $(LIBSPEC) && cd $(1) && ./build/$(1) -endef - -# Clean rules for cleaning specific targets -define CLEAN_RULES -clean_$(1): - rm -rf $(1)/build -endef - -$(foreach target,$(TARGETS),$(eval $(call BUILD_RULES,$(target)))) -$(foreach target,$(TARGETS),$(eval $(call RUN_RULES,$(target)))) -$(foreach target,$(TARGETS),$(eval $(call CLEAN_RULES,$(target)))) - -clean: $(addprefix clean_, $(TARGETS)) - -dawnlib: $(if $(wildcard $(GPUCPP)/third_party/lib/libdawn.so $(GPUCPP)/third_party/lib/libdawn.dylib),,run_setup) - -run_setup: check-python - cd $(GPUCPP) && (command -v python3 >/dev/null 2>&1 && python3 setup.py || python setup.py) - -check-python: - @command -v python3 >/dev/null 2>&1 || { echo >&2 "Python needs to be installed and in your path."; exit 1; } diff --git a/examples/README.md b/examples/README.md index bfd513e..fafcc27 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,28 +1,19 @@ # gpu.cpp examples -This directory contains examples of how to use gpu.cpp. +The root CMake project builds every example: -Each example is a standalone project that can be built and run independently by -running `make` from within the example directory. +```bash +./test +``` -Before running any of these examples, make sure you've downloaded the Dawn -native webgpu installation binary by running `make dawnlib` from the root -directory of the repository. +- `hello_world`: minimal WGSL upload, dispatch, and readback +- `float16`: device feature selection and f16 compute +- `gpu_puzzles`: exercises and solutions for Sasha Rush's GPU puzzles +- `matmul`: progressively optimized matrix multiplication kernels +- `physics`: an interactive double-pendulum simulation +- `render`: an interactive terminal SDF renderer +- `shadertui`: live-reloaded WGSL terminal shaders +- `transpose`: naive and tiled matrix transpose kernels -## Basic Examples - -| Example | Description | -|---------|-------------| -| [hello_world](hello_world) | Minimal example to get started with gpu.cpp, implements a GELU neural network activation function. | -| [gpu_puzzles](gpu_puzzles) | Implementation of Sasha Rush's GPU puzzles. | -| [shadertui](shadertui) | An example of runtime live reloading of WGSL - demonstrated using a terminal shadertoy-like scii rendering. | -| [render](render) | GPU ascii rendering of a signed distance function for two rotating 3D spheres. | -| [physics](physics) | Parallel physics simulation of a double pendulum with each thread starting at a different initial condition. | - -## Advanced Examples - -| Example | Description | -|---------|-------------| -| [float16](float16) | Hello World example using the float16 WebGPU extension, instead of the default float32. | -| [matmul](matmul) | Tiled matrix multiplication. | -| [transpose](transpose) | Tiled matrix transpose. | +`make run` runs `hello_world`. The interactive and benchmark examples are built +but not run by the test suite. diff --git a/examples/float16/CMakeLists.txt b/examples/float16/CMakeLists.txt deleted file mode 100644 index e4ef86a..0000000 --- a/examples/float16/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -cmake_minimum_required(VERSION 3.28) -project(float16) - -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Check if the file exists in the current directory -if(EXISTS ${FILEPATH_CURRENT_DIR}) - set(TARGET_FILE_PATH ${CMAKE_CURRENT_SOURCE_DIR}) -elseif(EXISTS ${FILEPATH_PROJECT_ROOT}) - set(TARGET_FILE_PATH ${PROJECT_ROOT}) -else() - message(FATAL_ERROR "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../") -endif() - -include("${TARGET_FILE_PATH}/cmake/example.cmake") diff --git a/examples/float16/Makefile b/examples/float16/Makefile deleted file mode 100644 index 51e895a..0000000 --- a/examples/float16/Makefile +++ /dev/null @@ -1,29 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -TARGET=float16 -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn - -run: ./build/$(TARGET) dawnlib - $(LIBSPEC) && ./build/$(TARGET) - -dawnlib: $(if $(wildcard $(GPUCPP)/third_party/lib/libwebgpu_dawn.so $(GPUCPP)/third_party/lib/libwebgpu_dawn.dylib),,run_setup) - -run_setup: check-python - cd $(GPUCPP) && python3 setup.py - -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -DNDEBUG -o ./build/$(TARGET) - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* - -check-python: - @command -v python3 >/dev/null 2>&1 || { echo >&2 "Python needs to be installed and in your path."; exit 1; } diff --git a/examples/float16/run.cpp b/examples/float16/run.cpp index 8f97210..20d3e58 100644 --- a/examples/float16/run.cpp +++ b/examples/float16/run.cpp @@ -1,6 +1,5 @@ #include #include -#include #include "gpu.hpp" #include "numeric_types/half.hpp" @@ -12,9 +11,8 @@ using namespace gpu; // createContext, createTensor, createKernel, static const char *kGelu = R"( const GELU_SCALING_FACTOR: f16 = 0.7978845608028654; // sqrt(2.0 / PI) -@group(0) @binding(0) var inp: array<{{precision}}>; +@group(0) @binding(0) var inp: array<{{precision}}>; @group(0) @binding(1) var out: array<{{precision}}>; -@group(0) @binding(1) var dummy: array<{{precision}}>; @compute @workgroup_size({{workgroupSize}}) fn main( @builtin(global_invocation_id) GlobalInvocationID: vec3) { @@ -33,26 +31,21 @@ int main(int argc, char **argv) { printf("--------------\n\n"); Context ctx = createContext( - {}, {}, - /*device descriptor, enabling f16 in WGSL*/ - { - .requiredFeatureCount = 1, - .requiredFeatures = std::array{WGPUFeatureName_ShaderF16}.data(), - }); + {.requiredFeatures = {wgpu::FeatureName::ShaderF16}}); static constexpr size_t N = 10000; std::array inputArr, outputArr; for (int i = 0; i < N; ++i) { inputArr[i] = half(static_cast(i) / 10.0f); // dummy input data } - Tensor input = createTensor(ctx, Shape{N}, kf16, inputArr.data()); + Tensor input = createTensor(ctx, Shape{N}, kf16, inputArr); Tensor output = createTensor(ctx, Shape{N}, kf16); - std::promise promise; - std::future future = promise.get_future(); - Kernel op = createKernel(ctx, {kGelu, 256, kf16}, Bindings{input, output}, - {cdiv(N, 256), 1, 1}); - dispatchKernel(ctx, op, promise); + Kernel op = createKernel(ctx, WGSL{kGelu, 256, kf16}, + Bindings{read(input), readWrite(output)}, + {ceilDiv(N, 256), 1, 1}); + auto future = dispatchKernel(ctx, op); wait(ctx, future); - toCPU(ctx, output, outputArr.data(), sizeof(outputArr)); + auto readback = toCPU(ctx, output, outputArr); + wait(ctx, readback); for (int i = 0; i < 12; ++i) { // Cast to float32 for printing to the screen @@ -63,5 +56,4 @@ int main(int argc, char **argv) { printf(" ...\n\n"); printf("Computed %zu float16 values of GELU(x: float16)\n\n", N); return 0; - } diff --git a/examples/gpu_puzzles/CMakeLists.txt b/examples/gpu_puzzles/CMakeLists.txt deleted file mode 100644 index b91e5c8..0000000 --- a/examples/gpu_puzzles/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -cmake_minimum_required(VERSION 3.28) -project(gpu_puzzles) - -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Check if the file exists in the current directory -if(EXISTS ${FILEPATH_CURRENT_DIR}) - set(TARGET_FILE_PATH ${CMAKE_CURRENT_SOURCE_DIR}) -elseif(EXISTS ${FILEPATH_PROJECT_ROOT}) - set(TARGET_FILE_PATH ${PROJECT_ROOT}) -else() - message(FATAL_ERROR "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../") -endif() - -include("${TARGET_FILE_PATH}/cmake/example.cmake") \ No newline at end of file diff --git a/examples/gpu_puzzles/Makefile b/examples/gpu_puzzles/Makefile deleted file mode 100644 index 90dfc2d..0000000 --- a/examples/gpu_puzzles/Makefile +++ /dev/null @@ -1,32 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -TARGET=gpu_puzzles -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn -FLAGS_KEY=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib key.cpp -ldl -lwebgpu_dawn - -run: ./build/$(TARGET) - $(LIBSPEC) && ./build/$(TARGET) - -run-key: ./build/key - $(LIBSPEC) && ./build/key - -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -o ./build/$(TARGET) - -build/key: key.cpp - mkdir -p build && $(CXX) $(FLAGS_KEY) -o ./build/key - -watch: - @command -v entr >/dev/null 2>&1 || { echo >&2 "Please install entr with 'brew install entr' or 'sudo apt-get install entr'"; exit 1; } - mkdir -p build && ls | entr -s "rm -f ./build/$(TARGET) && make -j$(NUM_JOBS) ./build/$(TARGET) && $(LIBSPEC) && ./build/$(TARGET)" - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* diff --git a/examples/gpu_puzzles/key.cpp b/examples/gpu_puzzles/key.cpp index 018ceac..1639263 100644 --- a/examples/gpu_puzzles/key.cpp +++ b/examples/gpu_puzzles/key.cpp @@ -7,7 +7,6 @@ #include "utils/array_utils.hpp" #include #include -#include using namespace gpu; @@ -21,14 +20,13 @@ template std::array makeData() { return inputArr; } -template void showResult(Context &ctx, Kernel &op, Tensor &output) { - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); +template +void showResult(Context &ctx, const Kernel &op, const Tensor &output) { + auto future = dispatchKernel(ctx, op); std::array outputArr; wait(ctx, future); - toCPU(ctx, output, outputArr.data(), sizeof(outputArr)); + auto readback = toCPU(ctx, output, outputArr); + wait(ctx, readback); printf("%s", show(outputArr, "output").c_str()); } @@ -36,7 +34,7 @@ template void showResult(Context &ctx, Ke // Implement a "kernel" (GPU function) that adds 10 to each position of vector // a and stores it in vector out. You have 1 thread per position. const char *kPuzzle1 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -51,9 +49,9 @@ fn main( void puzzle1(Context &ctx) { printf("\n\nPuzzle 1\n\n"); static constexpr size_t N = 4; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); - Kernel op = createKernel(ctx, {kPuzzle1, N}, Bindings{a, output}, + Kernel op = createKernel(ctx, WGSL{kPuzzle1, N}, Bindings{read(a), readWrite(output)}, /*totalWorkgroups */ {1, 1, 1}); showResult(ctx, op, output); } @@ -62,8 +60,8 @@ void puzzle1(Context &ctx) { // Implement a kernel that adds together each position of a and b and stores it // in out. You have 1 thread per position. const char *kPuzzle2 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output : array; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -78,10 +76,10 @@ fn main( void puzzle2(Context &ctx) { printf("\n\nPuzzle 2\n\n"); static constexpr size_t N = 4; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); + Tensor b = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); - Kernel op = createKernel(ctx, {kPuzzle2, N}, Bindings{a, b, output}, + Kernel op = createKernel(ctx, WGSL{kPuzzle2, N}, Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}); showResult(ctx, op, output); } @@ -90,7 +88,7 @@ void puzzle2(Context &ctx) { // Implement a kernel that adds 10 to each position of a and stores it in out. // You have more threads than positions. const char *kPuzzle3 = R"( -@group(0) @binding(0) var input: array; +@group(0) @binding(0) var input: array; @group(0) @binding(1) var output : array; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -104,10 +102,10 @@ fn main( void puzzle3(Context &ctx) { printf("\n\nPuzzle 3\n\n"); static constexpr size_t N = 8; - Tensor input = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor input = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); Kernel op = - createKernel(ctx, {kPuzzle3, N}, Bindings{input, output}, {1, 1, 1}); + createKernel(ctx, WGSL{kPuzzle3, N}, Bindings{read(input), readWrite(output)}, {1, 1, 1}); showResult(ctx, op, output); } @@ -115,7 +113,7 @@ void puzzle3(Context &ctx) { // Implement a kernel that adds 10 to each position of a and stores it in out. // Input a is 2D and square. You have more threads than positions. const char *kPuzzle4 = R"( -@group(0) @binding(0) var input: array; +@group(0) @binding(0) var input: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -139,14 +137,14 @@ void puzzle4(Context &ctx) { static constexpr size_t Wx = 3; static constexpr size_t Wy = 3; static constexpr size_t N = 2; - Tensor input = createTensor(ctx, {N, N}, kf32, makeData().data()); + Tensor input = createTensor(ctx, {N, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle4, /*workgroup size*/ {Wx, Wy, 1}}, - Bindings{input, output}, /* totalWorkgroups */ {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle4, /*workgroup size*/ {Wx, Wy, 1}}, + Bindings{read(input), readWrite(output)}, /* totalWorkgroups */ {1, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -154,8 +152,8 @@ void puzzle4(Context &ctx) { // Implement a kernel that adds a and b and stores it in out. Inputs a and b // are vectors. You have more threads than positions. const char *kPuzzle5 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b : array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b : array; @group(0) @binding(2) var output : array; @group(0) @binding(3) var params: Params; struct Params { @@ -178,16 +176,16 @@ void puzzle5(Context &ctx) { static constexpr size_t N = 2; static constexpr size_t Wx = 3; static constexpr size_t Wy = 3; - Tensor a = createTensor(ctx, {N, 1}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {1, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N, 1}, kf32, makeData()); + Tensor b = createTensor(ctx, {1, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle5, /*workgroup size*/ {Wx, Wy, 1}}, - Bindings{a, b, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle5, /*workgroup size*/ {Wx, Wy, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -196,7 +194,7 @@ void puzzle5(Context &ctx) { // You have fewer threads per block than the size of a. const char *kPuzzle6 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -218,15 +216,15 @@ void puzzle6(Context &ctx) { static constexpr size_t N = 9; static constexpr size_t Wx = 4; static constexpr size_t Bx = 3; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle6, {Wx, 1, 1}}, - Bindings{a, output}, {Bx, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle6, {Wx, 1, 1}}, + Bindings{read(a), readWrite(output)}, {Bx, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -235,7 +233,7 @@ void puzzle6(Context &ctx) { // You have fewer threads per block than the size of a in both directions. const char *kPuzzle7 = R"( -@group(0)@binding(0) var a: array; +@group(0)@binding(0) var a: array; @group(0)@binding(1) var output : array; @group(0)@binding(2) var params: Params; struct Params { @@ -261,15 +259,15 @@ void puzzle7(Context &ctx) { static constexpr size_t Wy = 3; static constexpr size_t Bx = 2; static constexpr size_t By = 2; - Tensor a = createTensor(ctx, {N, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle7, {Wx, Wy, 1}}, - Bindings{a, output}, {Bx, By, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle7, {Wx, Wy, 1}}, + Bindings{read(a), readWrite(output)}, {Bx, By, 1}, Params{N}); showResult(ctx, op, output); } @@ -280,7 +278,7 @@ void puzzle7(Context &ctx) { // (This example does not really need shared memory or syncthreads, but it is a demo.) const char *kPuzzle8 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -310,7 +308,7 @@ void puzzle8(Context &ctx) { static constexpr size_t N = 8; static constexpr size_t Wx = 4; static constexpr size_t Bx = 2; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; @@ -318,8 +316,8 @@ void puzzle8(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle8, {Wx, 1, 1}}, - Bindings{a, output}, {Bx, 1, 1}, Params{N, 8}); + createKernel(ctx, WGSL{kPuzzle8, {Wx, 1, 1}}, + Bindings{read(a), readWrite(output)}, {Bx, 1, 1}, Params{N, 8}); showResult(ctx, op, output); } @@ -328,7 +326,7 @@ void puzzle8(Context &ctx) { // You have 1 thread per position. You only need 1 global read and 1 global write per thread. const char *kPuzzle9 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -363,15 +361,15 @@ void puzzle9(Context &ctx) { printf("\n\nPuzzle 9\n\n"); static constexpr size_t N = 8; static constexpr size_t Wx = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle9, {Wx, 1, 1}}, - Bindings{a, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle9, {Wx, 1, 1}}, + Bindings{read(a), readWrite(output)}, {1, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -380,8 +378,8 @@ void puzzle9(Context &ctx) { // You have 1 thread per position. You only need 2 global reads and 1 global write per thread. const char *kPuzzle10 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output : array; @group(0) @binding(3) var params: Params; struct Params { @@ -414,16 +412,16 @@ void puzzle10(Context &ctx) { printf("\n\nPuzzle 10\n\n"); static constexpr size_t N = 8; static constexpr size_t Wx = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); + Tensor b = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {1}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle10, {Wx, 1, 1}}, - Bindings{a, b, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle10, {Wx, 1, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}, Params{N}); showResult<1>(ctx, op, output); } @@ -432,8 +430,8 @@ void puzzle10(Context &ctx) { // You need to handle the general case. You only need 2 global reads and 1 global write per thread. const char *kPuzzle11 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output : array; @group(0) @binding(3) var params: Params; struct Params { @@ -484,8 +482,8 @@ void puzzle11(Context &ctx) { static constexpr size_t N = 6; static constexpr size_t CONV = 3; static constexpr size_t Wx = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {CONV}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); + Tensor b = createTensor(ctx, {CONV}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; @@ -493,8 +491,8 @@ void puzzle11(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle11, {N, 1, 1}}, - Bindings{a, b, output}, {Wx, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle11, {N, 1, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {Wx, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -505,7 +503,7 @@ void puzzle11(Context &ctx) { // That is, each step of the algorithm should sum together half the remaining numbers. const char *kPuzzle12 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -545,15 +543,15 @@ fn main( void puzzle12(Context &ctx) { printf("\n\nPuzzle 12\n\n"); static constexpr size_t N = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {1}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle12, {N, 1, 1}}, - Bindings{a, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle12, {N, 1, 1}}, + Bindings{read(a), readWrite(output)}, {1, 1, 1}, Params{N}); showResult<1>(ctx, op, output); } @@ -562,7 +560,7 @@ void puzzle12(Context &ctx) { // Implement a kernel that computes a sum over each column of a and stores it in out. const char *kPuzzle13 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output: array; @group(0) @binding(2) var params: Params; @@ -613,7 +611,7 @@ void puzzle13(Context &ctx) { static constexpr size_t N = 6; static constexpr size_t TPB = 8; static constexpr size_t BATCH = 4; - Tensor a = createTensor(ctx, {BATCH, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {BATCH, N}, kf32, makeData()); Tensor output = createTensor(ctx, {BATCH}, kf32); struct Params { uint32_t TPB = TPB; @@ -621,8 +619,8 @@ void puzzle13(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle13, {TPB, 1, 1}}, - Bindings{a, output}, {1, BATCH, 1}, Params{TPB, N}); + createKernel(ctx, WGSL{kPuzzle13, {TPB, 1, 1}}, + Bindings{read(a), readWrite(output)}, {1, BATCH, 1}, Params{TPB, N}); showResult(ctx, op, output); } @@ -634,8 +632,8 @@ void puzzle13(Context &ctx) { // partial dot-product and iteratively move the part you copied into shared memory. You // should be able to do the hard case in 6 global reads. const char *kPuzzle14 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output: array; @group(0) @binding(3) var params: Params; @@ -686,8 +684,8 @@ void puzzle14(Context &ctx) { printf("\n\nPuzzle 14\n\n"); static constexpr size_t N = 2; static constexpr size_t TPB = 3; - Tensor a = createTensor(ctx, {N, N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {N, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N, N}, kf32, makeData()); + Tensor b = createTensor(ctx, {N, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t TPB = TPB; @@ -695,8 +693,8 @@ void puzzle14(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle14, {TPB, TPB, 1}}, - Bindings{a, b, output}, {1, 1, 1}, Params{TPB, N}); + createKernel(ctx, WGSL{kPuzzle14, {TPB, TPB, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}, Params{TPB, N}); showResult(ctx, op, output); } diff --git a/examples/gpu_puzzles/run.cpp b/examples/gpu_puzzles/run.cpp index e337688..2093353 100644 --- a/examples/gpu_puzzles/run.cpp +++ b/examples/gpu_puzzles/run.cpp @@ -7,7 +7,6 @@ #include "utils/array_utils.hpp" #include #include -#include using namespace gpu; @@ -21,14 +20,13 @@ template std::array makeData() { return inputArr; } -template void showResult(Context &ctx, Kernel &op, Tensor &output) { - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); +template +void showResult(Context &ctx, const Kernel &op, const Tensor &output) { + auto future = dispatchKernel(ctx, op); std::array outputArr; wait(ctx, future); - toCPU(ctx, output, outputArr.data(), sizeof(outputArr)); + auto readback = toCPU(ctx, output, outputArr); + wait(ctx, readback); printf("%s", show(outputArr, "output").c_str()); } @@ -36,7 +34,7 @@ template void showResult(Context &ctx, Ke // Implement a "kernel" (GPU function) that adds 10 to each position of vector // a and stores it in vector out. You have 1 thread per position. const char *kPuzzle1 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -48,9 +46,9 @@ fn main( void puzzle1(Context &ctx) { printf("\n\nPuzzle 1\n\n"); static constexpr size_t N = 4; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); - Kernel op = createKernel(ctx, {kPuzzle1, N}, Bindings{a, output}, + Kernel op = createKernel(ctx, WGSL{kPuzzle1, N}, Bindings{read(a), readWrite(output)}, /*nWorkgroups */ {1, 1, 1}); showResult(ctx, op, output); } @@ -59,8 +57,8 @@ void puzzle1(Context &ctx) { // Implement a kernel that adds together each position of a and b and stores it // in out. You have 1 thread per position. const char *kPuzzle2 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output : array; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -72,10 +70,10 @@ fn main( void puzzle2(Context &ctx) { printf("\n\nPuzzle 2\n\n"); static constexpr size_t N = 4; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); + Tensor b = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); - Kernel op = createKernel(ctx, {kPuzzle2, N}, Bindings{a, b, output}, + Kernel op = createKernel(ctx, WGSL{kPuzzle2, N}, Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}); showResult(ctx, op, output); } @@ -84,7 +82,7 @@ void puzzle2(Context &ctx) { // Implement a kernel that adds 10 to each position of a and stores it in out. // You have more threads than positions. const char *kPuzzle3 = R"( -@group(0) @binding(0) var input: array; +@group(0) @binding(0) var input: array; @group(0) @binding(1) var output : array; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -95,10 +93,10 @@ fn main( void puzzle3(Context &ctx) { printf("\n\nPuzzle 3\n\n"); static constexpr size_t N = 8; - Tensor input = createTensor(ctx, {N/2}, kf32, makeData().data()); + Tensor input = createTensor(ctx, {N/2}, kf32, makeData()); Tensor output = createTensor(ctx, {N/2}, kf32); Kernel op = - createKernel(ctx, {kPuzzle3, N}, Bindings{input, output}, {1, 1, 1}); + createKernel(ctx, WGSL{kPuzzle3, N}, Bindings{read(input), readWrite(output)}, {1, 1, 1}); showResult(ctx, op, output); } @@ -106,7 +104,7 @@ void puzzle3(Context &ctx) { // Implement a kernel that adds 10 to each position of a and stores it in out. // Input a is 2D and square. You have more threads than positions. const char *kPuzzle4 = R"( -@group(0) @binding(0) var input: array; +@group(0) @binding(0) var input: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -124,14 +122,14 @@ void puzzle4(Context &ctx) { static constexpr size_t Wx = 3; static constexpr size_t Wy = 3; static constexpr size_t N = 2; - Tensor input = createTensor(ctx, {N, N}, kf32, makeData().data()); + Tensor input = createTensor(ctx, {N, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle4, /*workgroup size*/ {Wx, Wy, 1}}, - Bindings{input, output}, /* nWorkgroups */ {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle4, /*workgroup size*/ {Wx, Wy, 1}}, + Bindings{read(input), readWrite(output)}, /* nWorkgroups */ {1, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -139,8 +137,8 @@ void puzzle4(Context &ctx) { // Implement a kernel that adds a and b and stores it in out. Inputs a and b // are vectors. You have more threads than positions. const char *kPuzzle5 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b : array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b : array; @group(0) @binding(2) var output : array; @group(0) @binding(3) var params: Params; struct Params { @@ -158,16 +156,16 @@ void puzzle5(Context &ctx) { static constexpr size_t N = 2; static constexpr size_t Wx = 3; static constexpr size_t Wy = 3; - Tensor a = createTensor(ctx, {N, 1}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {1, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N, 1}, kf32, makeData()); + Tensor b = createTensor(ctx, {1, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle5, /*workgroup size*/ {Wx, Wy, 1}}, - Bindings{a, b, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle5, /*workgroup size*/ {Wx, Wy, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -176,7 +174,7 @@ void puzzle5(Context &ctx) { // You have fewer threads per block than the size of a. const char *kPuzzle6 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -194,15 +192,15 @@ void puzzle6(Context &ctx) { static constexpr size_t N = 9; static constexpr size_t Wx = 4; static constexpr size_t Bx = 3; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle6, {Wx, 1, 1}}, - Bindings{a, output}, {Bx, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle6, {Wx, 1, 1}}, + Bindings{read(a), readWrite(output)}, {Bx, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -211,7 +209,7 @@ void puzzle6(Context &ctx) { // You have fewer threads per block than the size of a in both directions. const char *kPuzzle7 = R"( -@group(0)@binding(0) var a: array; +@group(0)@binding(0) var a: array; @group(0)@binding(1) var output : array; @group(0)@binding(2) var params: Params; struct Params { @@ -231,15 +229,15 @@ void puzzle7(Context &ctx) { static constexpr size_t Wy = 3; static constexpr size_t Bx = 2; static constexpr size_t By = 2; - Tensor a = createTensor(ctx, {N, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle7, {Wx, Wy, 1}}, - Bindings{a, output}, {Bx, By, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle7, {Wx, Wy, 1}}, + Bindings{read(a), readWrite(output)}, {Bx, By, 1}, Params{N}); showResult(ctx, op, output); } @@ -250,7 +248,7 @@ void puzzle7(Context &ctx) { // (This example does not really need shared memory or syncthreads, but it is a demo.) const char *kPuzzle8 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -269,7 +267,7 @@ void puzzle8(Context &ctx) { static constexpr size_t N = 8; static constexpr size_t Wx = 4; static constexpr size_t Bx = 2; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; @@ -277,8 +275,8 @@ void puzzle8(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle8, {Wx, 1, 1}}, - Bindings{a, output}, {Bx, 1, 1}, Params{N, 8}); + createKernel(ctx, WGSL{kPuzzle8, {Wx, 1, 1}}, + Bindings{read(a), readWrite(output)}, {Bx, 1, 1}, Params{N, 8}); showResult(ctx, op, output); } @@ -287,7 +285,7 @@ void puzzle8(Context &ctx) { // You have 1 thread per position. You only need 1 global read and 1 global write per thread. const char *kPuzzle9 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -305,15 +303,15 @@ void puzzle9(Context &ctx) { printf("\n\nPuzzle 9\n\n"); static constexpr size_t N = 8; static constexpr size_t Wx = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle9, {Wx, 1, 1}}, - Bindings{a, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle9, {Wx, 1, 1}}, + Bindings{read(a), readWrite(output)}, {1, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -322,8 +320,8 @@ void puzzle9(Context &ctx) { // You have 1 thread per position. You only need 2 global reads and 1 global write per thread. const char *kPuzzle10 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output : array; @group(0) @binding(3) var params: Params; struct Params { @@ -341,16 +339,16 @@ void puzzle10(Context &ctx) { printf("\n\nPuzzle 10\n\n"); static constexpr size_t N = 8; static constexpr size_t Wx = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); + Tensor b = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {1}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle10, {Wx, 1, 1}}, - Bindings{a, b, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle10, {Wx, 1, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}, Params{N}); showResult<1>(ctx, op, output); } @@ -359,8 +357,8 @@ void puzzle10(Context &ctx) { // You need to handle the general case. You only need 2 global reads and 1 global write per thread. const char *kPuzzle11 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output : array; @group(0) @binding(3) var params: Params; struct Params { @@ -381,8 +379,8 @@ void puzzle11(Context &ctx) { static constexpr size_t N = 6; static constexpr size_t CONV = 3; static constexpr size_t Wx = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {CONV}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); + Tensor b = createTensor(ctx, {CONV}, kf32, makeData()); Tensor output = createTensor(ctx, {N}, kf32); struct Params { uint32_t size = N; @@ -390,8 +388,8 @@ void puzzle11(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle11, {N, 1, 1}}, - Bindings{a, b, output}, {Wx, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle11, {N, 1, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {Wx, 1, 1}, Params{N}); showResult(ctx, op, output); } @@ -402,7 +400,7 @@ void puzzle11(Context &ctx) { // That is, each step of the algorithm should sum together half the remaining numbers. const char *kPuzzle12 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output : array; @group(0) @binding(2) var params: Params; struct Params { @@ -419,15 +417,15 @@ fn main( void puzzle12(Context &ctx) { printf("\n\nPuzzle 12\n\n"); static constexpr size_t N = 8; - Tensor a = createTensor(ctx, {N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N}, kf32, makeData()); Tensor output = createTensor(ctx, {1}, kf32); struct Params { uint32_t size = N; }; Kernel op = - createKernel(ctx, {kPuzzle12, {N, 1, 1}}, - Bindings{a, output}, {1, 1, 1}, Params{N}); + createKernel(ctx, WGSL{kPuzzle12, {N, 1, 1}}, + Bindings{read(a), readWrite(output)}, {1, 1, 1}, Params{N}); showResult<1>(ctx, op, output); } @@ -436,7 +434,7 @@ void puzzle12(Context &ctx) { // Implement a kernel that computes a sum over each column of a and stores it in out. const char *kPuzzle13 = R"( -@group(0) @binding(0) var a: array; +@group(0) @binding(0) var a: array; @group(0) @binding(1) var output: array; @group(0) @binding(2) var params: Params; @@ -459,7 +457,7 @@ void puzzle13(Context &ctx) { static constexpr size_t N = 6; static constexpr size_t TPB = 8; static constexpr size_t BATCH = 4; - Tensor a = createTensor(ctx, {BATCH, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {BATCH, N}, kf32, makeData()); Tensor output = createTensor(ctx, {BATCH}, kf32); struct Params { uint32_t TPB = TPB; @@ -467,8 +465,8 @@ void puzzle13(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle13, {TPB, 1, 1}}, - Bindings{a, output}, {1, BATCH, 1}, Params{TPB, N}); + createKernel(ctx, WGSL{kPuzzle13, {TPB, 1, 1}}, + Bindings{read(a), readWrite(output)}, {1, BATCH, 1}, Params{TPB, N}); showResult(ctx, op, output); } @@ -480,8 +478,8 @@ void puzzle13(Context &ctx) { // partial dot-product and iteratively move the part you copied into shared memory. You // should be able to do the hard case in 6 global reads. const char *kPuzzle14 = R"( -@group(0) @binding(0) var a: array; -@group(0) @binding(1) var b: array; +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; @group(0) @binding(2) var output: array; @group(0) @binding(3) var params: Params; @@ -504,8 +502,8 @@ void puzzle14(Context &ctx) { printf("\n\nPuzzle 14\n\n"); static constexpr size_t N = 2; static constexpr size_t TPB = 3; - Tensor a = createTensor(ctx, {N, N}, kf32, makeData().data()); - Tensor b = createTensor(ctx, {N, N}, kf32, makeData().data()); + Tensor a = createTensor(ctx, {N, N}, kf32, makeData()); + Tensor b = createTensor(ctx, {N, N}, kf32, makeData()); Tensor output = createTensor(ctx, {N, N}, kf32); struct Params { uint32_t TPB = TPB; @@ -513,8 +511,8 @@ void puzzle14(Context &ctx) { }; Kernel op = - createKernel(ctx, {kPuzzle14, {TPB, TPB, 1}}, - Bindings{a, b, output}, {1, 1, 1}, Params{TPB, N}); + createKernel(ctx, WGSL{kPuzzle14, {TPB, TPB, 1}}, + Bindings{read(a), read(b), readWrite(output)}, {1, 1, 1}, Params{TPB, N}); showResult(ctx, op, output); } diff --git a/examples/hello_world/CMakeLists.txt b/examples/hello_world/CMakeLists.txt deleted file mode 100644 index 72791ec..0000000 --- a/examples/hello_world/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -cmake_minimum_required(VERSION 3.28) -project(hello_world) - -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Check if the file exists in the current directory -if(EXISTS ${FILEPATH_CURRENT_DIR}) - set(TARGET_FILE_PATH ${CMAKE_CURRENT_SOURCE_DIR}) -elseif(EXISTS ${FILEPATH_PROJECT_ROOT}) - set(TARGET_FILE_PATH ${PROJECT_ROOT}) -else() - message(FATAL_ERROR "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../") -endif() - -include("${TARGET_FILE_PATH}/cmake/example.cmake") \ No newline at end of file diff --git a/examples/hello_world/Makefile b/examples/hello_world/Makefile deleted file mode 100644 index 7e64553..0000000 --- a/examples/hello_world/Makefile +++ /dev/null @@ -1,32 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -TARGET=hello_world -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn - -run: ./build/$(TARGET) dawnlib - $(LIBSPEC) && ./build/$(TARGET) - -dawnlib: $(if $(wildcard $(GPUCPP)/third_party/lib/libwebgpu_dawn.so $(GPUCPP)/third_party/lib/libwebgpu_dawn.dylib),,run_setup) - -run_setup: check-python - cd $(GPUCPP) && python3 setup.py - -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -DNO_LOG -o ./build/$(TARGET) - -debug: run.cpp - mkdir -p build && $(CXX) $(FLAGS) -g -o ./build/$(TARGET) - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* - -check-python: - @command -v python3 >/dev/null 2>&1 || { echo >&2 "Python needs to be installed and in your path."; exit 1; } diff --git a/examples/hello_world/run.cpp b/examples/hello_world/run.cpp index 7453869..71ba7ab 100644 --- a/examples/hello_world/run.cpp +++ b/examples/hello_world/run.cpp @@ -1,15 +1,13 @@ #include "gpu.hpp" #include #include -#include using namespace gpu; static const char *kGelu = R"( const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI) -@group(0) @binding(0) var inp: array<{{precision}}>; +@group(0) @binding(0) var inp: array<{{precision}}>; @group(0) @binding(1) var out: array<{{precision}}>; -@group(0) @binding(1) var dummy: array<{{precision}}>; @compute @workgroup_size({{workgroupSize}}) fn main( @builtin(global_invocation_id) GlobalInvocationID: vec3) { @@ -27,23 +25,22 @@ int main(int argc, char **argv) { printf("\nHello gpu.cpp!\n"); printf("--------------\n\n"); - // std::unique_ptr ctx = createContext(); Context ctx = createContext(); static constexpr size_t N = 10000; std::array inputArr, outputArr; for (int i = 0; i < N; ++i) { inputArr[i] = static_cast(i) / 10.0; // dummy input data } - Tensor input = createTensor(ctx, Shape{N}, kf32, inputArr.data()); + Tensor input = createTensor( + ctx, Shape{N}, kf32, std::span(inputArr)); Tensor output = createTensor(ctx, Shape{N}, kf32); - std::promise promise; - std::future future = promise.get_future(); - Kernel op = createKernel(ctx, {kGelu, 256, kf32}, - Bindings{input, output}, - {cdiv(N, 256), 1, 1}); - dispatchKernel(ctx, op, promise); + Kernel op = createKernel(ctx, WGSL{kGelu, 256, kf32}, + Bindings{read(input), readWrite(output)}, + {(N + 255) / 256, 1, 1}); + auto future = dispatchKernel(ctx, op); wait(ctx, future); - toCPU(ctx, output, outputArr.data(), sizeof(outputArr)); + auto readback = toCPU(ctx, output, std::span(outputArr)); + wait(ctx, readback); for (int i = 0; i < 12; ++i) { printf(" gelu(%.2f) = %.2f\n", inputArr[i], outputArr[i]); } diff --git a/examples/matmul/CMakeLists.txt b/examples/matmul/CMakeLists.txt deleted file mode 100644 index a0c5a48..0000000 --- a/examples/matmul/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -cmake_minimum_required(VERSION 3.28) -project(matmul) - -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Check if the file exists in the current directory -if(EXISTS ${FILEPATH_CURRENT_DIR}) - set(TARGET_FILE_PATH ${CMAKE_CURRENT_SOURCE_DIR}) -elseif(EXISTS ${FILEPATH_PROJECT_ROOT}) - set(TARGET_FILE_PATH ${PROJECT_ROOT}) -else() - message(FATAL_ERROR "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../") -endif() - -include("${TARGET_FILE_PATH}/cmake/example.cmake") \ No newline at end of file diff --git a/examples/matmul/Makefile b/examples/matmul/Makefile deleted file mode 100644 index 03cd20e..0000000 --- a/examples/matmul/Makefile +++ /dev/null @@ -1,38 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -CODEPATH = find . ../../utils ../../ -maxdepth 1 -type f -TARGET=matmul -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn - -run: ./build/$(TARGET) - $(LIBSPEC) && ./build/$(TARGET) - -run_with_metal_profiler: ./build/$(TARGET)_with_metal_profiler - $(LIBSPEC) && export METAL_CAPTURE_ENABLED=1 && ./build/$(TARGET)_with_metal_profiler - -run_with_time_profiler: ./build/$(TARGET)_with_metal_profiler - $(LIBSPEC) && xcrun xctrace record --template 'Time Profiler' --launch -- ./build/$(TARGET)_with_metal_profiler - -# Use clang -v to see the include paths -# Note in this example optimization is turned on -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -o ./build/$(TARGET) - -build/$(TARGET)_with_metal_profiler: run.cpp - mkdir -p build && $(CXX) $(FLAGS) -o ./build/$(TARGET)_with_metal_profiler $(GPUCPP)/experimental/profiler/metal.mm -framework metal -framework Foundation -DMETAL_PROFILER -g - install_name_tool -change @rpath/libwebgpu_dawn.dylib $(GPUCPP)/third_party/lib/libwebgpu_dawn.dylib ./build/$(TARGET)_with_metal_profiler - -watch: - @command -v entr >/dev/null 2>&1 || { echo >&2 "Please install entr with 'brew install entr' or 'sudo apt-get install entr'"; exit 1; } - mkdir -p build && $(CODEPATH) | entr -s "$(LIBSPEC) && rm -f ./build/$(TARGET) && make -j$(NUM_JOBS) ./build/$(TARGET) && ./build/$(TARGET)" - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* diff --git a/examples/matmul/run.cpp b/examples/matmul/run.cpp index 42d7009..634e881 100644 --- a/examples/matmul/run.cpp +++ b/examples/matmul/run.cpp @@ -4,19 +4,14 @@ #include #include -#include "gpu.hpp" // createContext, createTensor, createKernel, dispatchKernel, - // wait, resetCommandBuffer, toCPU +#include "gpu.hpp" #include "llmc/reference_impls.h" // for CPU reference implementation #include "utils/array_utils.hpp" // show, isclose, randn, randint #include "utils/logging.hpp" // LOG -#include "experimental/wgsl.h" // loopUnrolling +#include "utils/wgsl.hpp" #include "numeric_types/half.hpp" -#ifdef METAL_PROFILER -#include "experimental/profiler/metal.hpp" -#endif - using namespace gpu; const std::string versionToStr(int version); @@ -33,10 +28,10 @@ void matmulf16_forward_cpu(half* out, half* out_bt = out + b * T * OC + t * OC; const half* inp_bt = inp + b * T * C + t * C; for (int o = 0; o < OC; o++) { - float val = (bias != NULL) ? halfToFloat(bias[o]) : 0.0f; + float val = (bias != NULL) ? static_cast(bias[o]) : 0.0f; const half* wrow = weight + o*C; for (int i = 0; i < C; i++) { - val += halfToFloat(inp_bt[i]) * halfToFloat(wrow[i]); + val += static_cast(inp_bt[i]) * static_cast(wrow[i]); } out_bt[o] = val; } @@ -45,8 +40,8 @@ void matmulf16_forward_cpu(half* out, } static const char *kShaderMatmul1 = R"( -@group(0) @binding(0) var A: array<{{precision}}>; -@group(0) @binding(1) var B: array<{{precision}}>; +@group(0) @binding(0) var A: array<{{precision}}>; +@group(0) @binding(1) var B: array<{{precision}}>; @group(0) @binding(2) var C: array<{{precision}}>; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -65,7 +60,7 @@ fn main( } )"; -inline KernelCode createMatmul1(const char *shaderTemplate, const size_t M, +inline WGSL createMatmul1(const char *shaderTemplate, const size_t M, const size_t K, const size_t N, const Shape &workgroupSize = {256, 1, 1}, NumType precision = kf32) { @@ -80,8 +75,8 @@ inline KernelCode createMatmul1(const char *shaderTemplate, const size_t M, // Shared memory cache-blocking static const char *kShaderMatmul2 = R"( -@group(0) @binding(0) var A: array<{{precision}}>; -@group(0) @binding(1) var B: array<{{precision}}>; +@group(0) @binding(0) var A: array<{{precision}}>; +@group(0) @binding(1) var B: array<{{precision}}>; @group(0) @binding(2) var C: array<{{precision}}>; var As: array<{{precision}}, {{tileSize}} * {{tileSize}}>; var Bs: array<{{precision}}, {{tileSize}} * {{tileSize}}>; @@ -123,7 +118,7 @@ fn main( } )"; -inline KernelCode createMatmul2(const char *shaderTemplate, const size_t M, +inline WGSL createMatmul2(const char *shaderTemplate, const size_t M, const size_t K, const size_t N, const Shape &workgroupSize = {256, 1, 1}, NumType precision = kf32) { @@ -161,8 +156,8 @@ inline KernelCode createMatmul2(const char *shaderTemplate, const size_t M, * */ static const char *kShaderMatmul3 = R"( -@group(0) @binding(0) var a: array<{{precision}}>; -@group(0) @binding(1) var b: array<{{precision}}>; +@group(0) @binding(0) var a: array<{{precision}}>; +@group(0) @binding(1) var b: array<{{precision}}>; @group(0) @binding(2) var c: array<{{precision}}>; var tileA: array<{{precision}}, {{BM}} * {{BK}}>; var tileB: array<{{precision}}, {{BN}} * {{BK}}>; @@ -225,7 +220,7 @@ fn main( } )"; -inline KernelCode createMatmul3(const char *shaderTemplate, const size_t M, +inline WGSL createMatmul3(const char *shaderTemplate, const size_t M, const size_t K, const size_t N, const size_t BM, const size_t BK, const size_t BN, const size_t TM, @@ -250,7 +245,7 @@ inline KernelCode createMatmul3(const char *shaderTemplate, const size_t M, {"{{BN}}", toString(BN)}, {"{{TM}}", toString(TM)}}); if (unrolling) { - std::string unrolledCode = loopUnrolling(codeString); + std::string unrolledCode = unrollLoops(codeString); // LOG(kDefLog, kInfo, "Unrolled code:\n%s", unrolledCode.c_str()); return {unrolledCode, workgroupSize, precision}; } else { @@ -262,8 +257,8 @@ inline KernelCode createMatmul3(const char *shaderTemplate, const size_t M, * */ static const char *kShaderMatmul4 = R"( -@group(0) @binding(0) var a: array<{{precision}}>; -@group(0) @binding(1) var b: array<{{precision}}>; +@group(0) @binding(0) var a: array<{{precision}}>; +@group(0) @binding(1) var b: array<{{precision}}>; @group(0) @binding(2) var c: array<{{precision}}>; var tileA: array<{{precision}}, {{BM}} * {{BK}}>; var tileB: array<{{precision}}, {{BN}} * {{BK}}>; @@ -337,7 +332,7 @@ fn main( } )"; -inline KernelCode createMatmul4(const char *shaderTemplate, const size_t M, +inline WGSL createMatmul4(const char *shaderTemplate, const size_t M, const size_t K, const size_t N, const size_t BM, const size_t BK, const size_t BN, const size_t TM, const size_t TN, @@ -366,7 +361,7 @@ inline KernelCode createMatmul4(const char *shaderTemplate, const size_t M, {"{{NUM_TILEB}}", toString(BN * BK / num_threads)} }); if (unrolling) { - std::string unrolledCode = loopUnrolling(codeString); + std::string unrolledCode = unrollLoops(codeString); // LOG(kDefLog, kInfo, "Unrolled code:\n%s", unrolledCode.c_str()); return {unrolledCode, workgroupSize, precision}; } else { @@ -378,8 +373,8 @@ inline KernelCode createMatmul4(const char *shaderTemplate, const size_t M, * */ static const char *kShaderMatmulWithVectorization = R"( -@group(0) @binding(0) var a: array<{{precision}}>; -@group(0) @binding(1) var b: array<{{precision}}>; +@group(0) @binding(0) var a: array<{{precision}}>; +@group(0) @binding(1) var b: array<{{precision}}>; @group(0) @binding(2) var c: array>; var tileA: array<{{precision}}, {{BM}} * {{BK}}>; var tileB: array<{{precision}}, {{BN}} * {{BK}}>; @@ -456,7 +451,7 @@ fn main( } )"; -inline KernelCode createMatmulWithVectorization(const char *shaderTemplate, const size_t M, +inline WGSL createMatmulWithVectorization(const char *shaderTemplate, const size_t M, const size_t K, const size_t N, const size_t BM, const size_t BK, const size_t BN, const size_t TM, const size_t TN, @@ -488,7 +483,7 @@ inline KernelCode createMatmulWithVectorization(const char *shaderTemplate, cons {"{{BN4}}", toString(BN / 4)}, }); if (unrolling) { - std::string unrolledCode = loopUnrolling(codeString); + std::string unrolledCode = unrollLoops(codeString); // LOG(kDefLog, kInfo, "Unrolled code:\n%s", unrolledCode.c_str()); return {unrolledCode, workgroupSize, precision}; } else { @@ -500,8 +495,8 @@ inline KernelCode createMatmulWithVectorization(const char *shaderTemplate, cons * */ static const char *kShaderMatmulWithTranspose = R"( -@group(0) @binding(0) var a: array<{{precision}}>; -@group(0) @binding(1) var b: array<{{precision}}>; +@group(0) @binding(0) var a: array<{{precision}}>; +@group(0) @binding(1) var b: array<{{precision}}>; @group(0) @binding(2) var c: array>; var tileA: array<{{precision}}, {{BM}} * {{BK}}>; var tileB: array<{{precision}}, {{BK}} * {{BN}}>; @@ -578,7 +573,7 @@ fn main( } )"; -inline KernelCode createMatmulWithTranspose(const char *shaderTemplate, const size_t M, +inline WGSL createMatmulWithTranspose(const char *shaderTemplate, const size_t M, const size_t K, const size_t N, const size_t BM, const size_t BK, const size_t BN, const size_t TM, const size_t TN, @@ -608,7 +603,7 @@ inline KernelCode createMatmulWithTranspose(const char *shaderTemplate, const si {"{{N4}}", toString(N / 4)}, {"{{BN4}}", toString(BN / 4)}, }); - std::string unrolledCode = loopUnrolling(codeString); + std::string unrolledCode = unrollLoops(codeString); // LOG(kDefLog, kInfo, "Unrolled code:\n%s", unrolledCode.c_str()); return {unrolledCode, workgroupSize, precision}; } @@ -617,8 +612,8 @@ inline KernelCode createMatmulWithTranspose(const char *shaderTemplate, const si * @brief No-Op shader with matmul bindings for performance testing */ static const char *kShaderNoOp = R"( -@group(0) @binding(0) var A: array<{{precision}}>; -@group(0) @binding(1) var B: array<{{precision}}>; +@group(0) @binding(0) var A: array<{{precision}}>; +@group(0) @binding(1) var B: array<{{precision}}>; @group(0) @binding(2) var C: array<{{precision}}>; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -626,7 +621,7 @@ fn main( } )"; -inline KernelCode createNoOp(const char *shaderTemplate, +inline WGSL createNoOp(const char *shaderTemplate, const Shape &workgroupSize = {256, 1, 1}, NumType precision = kf32) { std::string codeString(shaderTemplate); @@ -685,24 +680,24 @@ Kernel selectMatmul(Context &ctx, int version, Kernel kernel; if (version == 1) { Shape wgSize = {256, 1, 1}; - Shape nWorkgroups = cdiv({M, N, 1}, {16, 16, 1}); - KernelCode matmul = createNoOp(kShaderNoOp, /*wgsize*/ wgSize); + Shape nWorkgroups = ceilDiv({M, N, 1}, {16, 16, 1}); + WGSL matmul = createNoOp(kShaderNoOp, /*wgsize*/ wgSize); kernel = createKernel(ctx, matmul, bindings, /*nWorkgroups*/ nWorkgroups); } else if (version == 2) { Shape wgSize = {16, 16, 1}; LOG(kDefLog, kInfo, "wgSize: %s", toString(wgSize).c_str()); - KernelCode matmul = + WGSL matmul = createMatmul1(kShaderMatmul1, M, K, N, /*wgsize*/ wgSize, numtype); kernel = createKernel(ctx, matmul, bindings, - /*nWorkgroups*/ cdiv({M, N, 1}, wgSize)); + /*nWorkgroups*/ ceilDiv({M, N, 1}, wgSize)); } else if (version == 3) { static constexpr size_t tileSize = 16; - KernelCode matmul = createMatmul2(kShaderMatmul2, M, K, N, + WGSL matmul = createMatmul2(kShaderMatmul2, M, K, N, /*wgSize*/ {tileSize * tileSize, 1, 1}, numtype); kernel = createKernel(ctx, matmul, bindings, - /* nWorkgroups*/ cdiv({M, N, 1}, {tileSize, tileSize, 1})); + /* nWorkgroups*/ ceilDiv({M, N, 1}, {tileSize, tileSize, 1})); } else if (version == 4 || version == 6) { static constexpr size_t BM = 64; static constexpr size_t BK = 4; @@ -711,12 +706,12 @@ Kernel selectMatmul(Context &ctx, int version, BN / BK; // BM * BN / TM == BM * BK, therefore TM == BN / BK Shape wgSize = {BM * BN / TM, 1, 1}; // BM * BN values per workgroup, TM values per thread - Shape nWorkgroups = {cdiv(M, BM), cdiv(N, BN), 1}; + Shape nWorkgroups = {ceilDiv(M, BM), ceilDiv(N, BN), 1}; LOG(kDefLog, kInfo, "M: %d, K: %d, N: %d", M, K, N); LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d", BM, BK, BN, TM); LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); - KernelCode matmul = createMatmul3(kShaderMatmul3, M, K, N, BM, BK, BN, TM, + WGSL matmul = createMatmul3(kShaderMatmul3, M, K, N, BM, BK, BN, TM, /*wgSize*/ wgSize, numtype, /*Loop unrolling*/ version == 6 ? true: false); @@ -729,12 +724,12 @@ Kernel selectMatmul(Context &ctx, int version, static constexpr size_t TM = BM / BK; static constexpr size_t TN = BN / BK; Shape wgSize = {(BM / TM) * (BN / TN), 1, 1}; // This is the same as BK * BK. - Shape nWorkgroups = {cdiv(M, BM), cdiv(N, BN), 1}; + Shape nWorkgroups = {ceilDiv(M, BM), ceilDiv(N, BN), 1}; LOG(kDefLog, kInfo, "M: %d, K: %d, N: %d", M, K, N); LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d, TN: %d", BM, BK, BN, TM, TN); LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); - KernelCode matmul = createMatmul4(kShaderMatmul4, M, K, N, BM, BK, BN, TM, TN, + WGSL matmul = createMatmul4(kShaderMatmul4, M, K, N, BM, BK, BN, TM, TN, /*wgSize*/ wgSize, numtype, /*Loop unrolling*/ version == 7 ? true: false); @@ -747,12 +742,12 @@ Kernel selectMatmul(Context &ctx, int version, static constexpr size_t TM = BM / BK; static constexpr size_t TN = BN / BK; Shape wgSize = {(BM / TM) * (BN / TN), 1, 1}; // This is the same as BK * BK. - Shape nWorkgroups = {cdiv(M, BM), cdiv(N, BN), 1}; + Shape nWorkgroups = {ceilDiv(M, BM), ceilDiv(N, BN), 1}; LOG(kDefLog, kInfo, "M: %d, K: %d, N: %d", M, K, N); LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d, TN: %d", BM, BK, BN, TM, TN); LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); - KernelCode matmul = createMatmulWithVectorization(kShaderMatmulWithVectorization, M, K, N, BM, BK, BN, TM, TN, + WGSL matmul = createMatmulWithVectorization(kShaderMatmulWithVectorization, M, K, N, BM, BK, BN, TM, TN, /*wgSize*/ wgSize, numtype, /*Loop unrolling*/ true); @@ -765,12 +760,12 @@ Kernel selectMatmul(Context &ctx, int version, static constexpr size_t TM = BM / BK; static constexpr size_t TN = BN / BK; Shape wgSize = {(BM / TM) * (BN / TN), 1, 1}; // This is the same as BK * BK. - Shape nWorkgroups = {cdiv(M, BM), cdiv(N, BN), 1}; + Shape nWorkgroups = {ceilDiv(M, BM), ceilDiv(N, BN), 1}; LOG(kDefLog, kInfo, "M: %d, K: %d, N: %d", M, K, N); LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d, TN: %d", BM, BK, BN, TM, TN); LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); - KernelCode matmul = createMatmulWithTranspose(kShaderMatmulWithTranspose, M, K, N, BM, BK, BN, TM, TN, + WGSL matmul = createMatmulWithTranspose(kShaderMatmulWithTranspose, M, K, N, BM, BK, BN, TM, TN, /*wgSize*/ wgSize, numtype); kernel = createKernel(ctx, matmul, bindings, @@ -791,61 +786,29 @@ void runTest(int version, size_t M, size_t K, size_t N, assert(numtype == kf16); } - // Allocate GPU buffers and copy data - WGPUDeviceDescriptor devDescriptor = {}; - devDescriptor.requiredFeatureCount = 1; - devDescriptor.requiredFeatures = std::array{WGPUFeatureName_ShaderF16}.data(); - - Context ctx; - if (numtype == kf16) { - ctx = createContext( - {}, {}, - /*device descriptor, enabling f16 in WGSL*/ - { - .requiredFeatureCount = 1, - .requiredFeatures = std::array{WGPUFeatureName_ShaderF16}.data() - }); - if (ctx.adapterStatus != WGPURequestAdapterStatus_Success) { - LOG(kDefLog, kError, "Failed to create adapter with f16 support, try running an f32 test instead (`export MATMUL_VERSION=9)."); - exit(1); - } - if (ctx.deviceStatus != WGPURequestDeviceStatus_Success) { - LOG(kDefLog, kError, "Failed to create device with f16 support, try running an f32 test instead. (`export MATMUL_VERSION=9)"); - exit(1); - } - } - - if (numtype == kf32) { - ctx = createContext({}, {}, {}); - if (ctx.adapterStatus != WGPURequestAdapterStatus_Success || - ctx.deviceStatus != WGPURequestDeviceStatus_Success) { - LOG(kDefLog, kError, "Failed to create adapter or device"); - // stop execution - exit(1); - } else { - LOG(kDefLog, kInfo, "Successfully created adapter and device"); - } - } + ContextOptions options; + if (numtype == kf16) + options.requiredFeatures = {wgpu::FeatureName::ShaderF16}; + Context ctx = createContext(options); - Tensor input = createTensor(ctx, Shape{M, K}, numtype, inputPtr.get()); - Tensor weights = createTensor(ctx, Shape{N, K}, numtype, weightsPtr.get()); // column-major + Tensor input = createTensor( + ctx, Shape{M, K}, numtype, + std::span(inputPtr.get(), M * K)); + Tensor weights = createTensor( + ctx, Shape{N, K}, numtype, + std::span(weightsPtr.get(), N * K)); // column-major -#ifdef METAL_PROFILER - constexpr size_t nIter = 1; -#else constexpr size_t nIter = 30; -#endif - // Initialize Kernel and bind GPU buffers - // pre-allocate for async dispatch - std::array, nIter> promises; - std::array, nIter> futures; + std::array futures; std::array kernels; std::array outputs; for (int i = 0; i < nIter; i++) { - futures[i] = promises[i].get_future(); outputs[i] = createTensor(ctx, Shape{M, N}, numtype); - kernels[i] = selectMatmul(ctx, version, {input, weights, outputs[i]}, M, K, N, numtype); + kernels[i] = selectMatmul( + ctx, version, + Bindings{read(input), read(weights), readWrite(outputs[i])}, M, K, N, + numtype); } LOG(kDefLog, kInfo, "Dispatching Kernel version %d: %s, %d iterations ...", @@ -854,7 +817,7 @@ void runTest(int version, size_t M, size_t K, size_t N, // Dispatch kernel nIter times auto start = std::chrono::high_resolution_clock::now(); for (int i = 0; i < nIter; i++) { - dispatchKernel(ctx, kernels[i], promises[i]); + futures[i] = dispatchKernel(ctx, kernels[i]); } for (int i = 0; i < nIter; i++) { wait(ctx, futures[i]); @@ -871,7 +834,9 @@ void runTest(int version, size_t M, size_t K, size_t N, 1000000000.0 * static_cast(nIter); LOG(kDefLog, kInfo, "Copying result to CPU"); - toCPU(ctx, outputs[0], outputPtr.get(), M * N * sizeof(precision)); + auto readback = toCPU( + ctx, outputs[0], std::span(outputPtr.get(), M * N)); + wait(ctx, readback); LOG(kDefLog, kInfo, "%s", show(outputPtr.get(), M, N, "Output[0]").c_str()); @@ -961,17 +926,11 @@ int main() { N = 2 * 4096; } -#ifdef METAL_PROFILER - startCapture(); -#endif if (enableF16) { runTestWithCheck(version, M, K, N, transposedInput, kTestSize, numtype); } else { runTestWithCheck(version, M, K, N, transposedInput, kTestSize, numtype); } -#ifdef METAL_PROFILER - stopCapture(); -#endif LOG(kDefLog, kInfo, "Done."); return 0; diff --git a/examples/physics/CMakeLists.txt b/examples/physics/CMakeLists.txt deleted file mode 100644 index b0b9d9e..0000000 --- a/examples/physics/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -cmake_minimum_required(VERSION 3.28) -project(physics) - -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Check if the file exists in the current directory -if(EXISTS ${FILEPATH_CURRENT_DIR}) - set(TARGET_FILE_PATH ${CMAKE_CURRENT_SOURCE_DIR}) -elseif(EXISTS ${FILEPATH_PROJECT_ROOT}) - set(TARGET_FILE_PATH ${PROJECT_ROOT}) -else() - message(FATAL_ERROR "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../") -endif() - -include("${TARGET_FILE_PATH}/cmake/example.cmake") \ No newline at end of file diff --git a/examples/physics/Makefile b/examples/physics/Makefile deleted file mode 100644 index 10cfb13..0000000 --- a/examples/physics/Makefile +++ /dev/null @@ -1,24 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -TARGET=physics -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn - -run: ./build/$(TARGET) - $(LIBSPEC) && ./build/$(TARGET) - -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -o ./build/$(TARGET) - -watch: - mkdir -p build && ls | entr -s "rm -f ./build/$(TARGET) && make -j$(NUM_JOBS) ./build/$(TARGET) && $(LIBSPEC) && ./build/$(TARGET)" - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* diff --git a/examples/physics/run.cpp b/examples/physics/run.cpp index 02b7e9f..5b16015 100644 --- a/examples/physics/run.cpp +++ b/examples/physics/run.cpp @@ -1,11 +1,10 @@ #include #include #include -#include #include -#include "experimental/tui.h" // rasterize #include "gpu.hpp" +#include "utils/tui.hpp" using namespace gpu; @@ -16,7 +15,7 @@ const dt: f32 = 0.03; @group(0) @binding(1) var theta2: array; @group(0) @binding(2) var thetaVel1: array; @group(0) @binding(3) var thetaVel2: array; -@group(0) @binding(4) var length: array; +@group(0) @binding(4) var length: array; @group(0) @binding(5) var pos: array; // x1, y1 for each pendulum //@group(0) @binding(6) var pos2: array; // x2, y2 for each pendulum @compute @workgroup_size({{workgroupSize}}) @@ -64,31 +63,32 @@ int main() { } // GPU buffers - Tensor theta1 = createTensor(ctx, Shape{N}, kf32, theta1Arr.data()); - Tensor theta2 = createTensor(ctx, Shape{N}, kf32, theta2Arr.data()); - Tensor vel1 = createTensor(ctx, Shape{N}, kf32, v1Arr.data()); - Tensor vel2 = createTensor(ctx, Shape{N}, kf32, v2Arr.data()); - Tensor length = createTensor(ctx, Shape{N}, kf32, lengthArr.data()); + Tensor theta1 = createTensor(ctx, Shape{N}, kf32, theta1Arr); + Tensor theta2 = createTensor(ctx, Shape{N}, kf32, theta2Arr); + Tensor vel1 = createTensor(ctx, Shape{N}, kf32, v1Arr); + Tensor vel2 = createTensor(ctx, Shape{N}, kf32, v2Arr); + Tensor length = createTensor(ctx, Shape{N}, kf32, lengthArr); std::array posArr; // x, y outputs for each pendulum std::string screen(80 * 40, ' '); Tensor pos = createTensor(ctx, Shape{N * 4}, kf32); // Prepare computation - KernelCode kernel{kUpdateSim, 256, kf32}; - printf("WGSL code: %s\n", kernel.data.c_str()); + WGSL kernel{kUpdateSim, 256, kf32}; + printf("WGSL code: %s\n", kernel.code.c_str()); Kernel update = createKernel( - ctx, kernel, Bindings{theta1, theta2, vel1, vel2, length, pos}, - /* nWorkgroups */ cdiv({N, 1, 1}, kernel.workgroupSize)); + ctx, kernel, + Bindings{readWrite(theta1), readWrite(theta2), readWrite(vel1), + readWrite(vel2), read(length), readWrite(pos)}, + ceilDiv(Shape{N, 1, 1}, kernel.workgroupSize)); // Main simulation update loop printf("\033[2J\033[H"); while (true) { auto start = std::chrono::high_resolution_clock::now(); - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, update, promise); + auto future = dispatchKernel(ctx, update); wait(ctx, future); - toCPU(ctx, pos, posArr.data(), sizeof(posArr)); + auto readback = toCPU(ctx, pos, posArr); + wait(ctx, readback); auto end = std::chrono::high_resolution_clock::now(); std::chrono::duration elapsed = end - start; // N * 2 because there's two objects per pendulum @@ -96,8 +96,7 @@ int main() { printf("\033[1;1H" // reset cursor "# simulations: %lu\n%s", N, screen.c_str()); - resetCommandBuffer(ctx.device, update); // Prepare kernel command - // buffer for nxt iteration - std::this_thread::sleep_for(std::chrono::milliseconds(8) - elapsed); + auto delay = std::chrono::duration(8) - elapsed; + if (delay.count() > 0) std::this_thread::sleep_for(delay); } } diff --git a/examples/render/CMakeLists.txt b/examples/render/CMakeLists.txt deleted file mode 100644 index bdb5867..0000000 --- a/examples/render/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -cmake_minimum_required(VERSION 3.28) -project(render) - -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Check if the file exists in the current directory -if(EXISTS ${FILEPATH_CURRENT_DIR}) - set(TARGET_FILE_PATH ${CMAKE_CURRENT_SOURCE_DIR}) -elseif(EXISTS ${FILEPATH_PROJECT_ROOT}) - set(TARGET_FILE_PATH ${PROJECT_ROOT}) -else() - message(FATAL_ERROR "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../") -endif() - -include("${TARGET_FILE_PATH}/cmake/example.cmake") \ No newline at end of file diff --git a/examples/render/Makefile b/examples/render/Makefile deleted file mode 100644 index d07048c..0000000 --- a/examples/render/Makefile +++ /dev/null @@ -1,25 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -TARGET=render -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn - -run: ./build/$(TARGET) - $(LIBSPEC) && ./build/$(TARGET) - -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -o ./build/$(TARGET) - -watch: - @command -v entr >/dev/null 2>&1 || { echo >&2 "Please install entr with 'brew install entr' or 'sudo apt-get install entr'"; exit 1; } - mkdir -p build && ls | entr -s "rm -f ./build/$(TARGET) && make -j$(NUM_JOBS) ./build/$(TARGET) && $(LIBSPEC) && ./build/$(TARGET)" - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* diff --git a/examples/render/run.cpp b/examples/render/run.cpp index f2c6bec..d3a6167 100644 --- a/examples/render/run.cpp +++ b/examples/render/run.cpp @@ -115,24 +115,23 @@ int main(int argc, char **argv) { std::fill(begin(screen), end(screen), 0.0f); Context ctx = createContext(); - Tensor devScreen = createTensor(ctx, {NROWS, NCOLS}, kf32, screen.data()); + Tensor devScreen = createTensor(ctx, {NROWS, NCOLS}, kf32, screen); uint32_t zeroTime = getCurrentTimeInMilliseconds(); Shape wgSize = {16, 16, 1}; - KernelCode code = {kSDF, wgSize}; - Kernel renderKernel = createKernel(ctx, code, Bindings{devScreen}, - cdiv({NCOLS, NROWS, 1}, wgSize), params); + WGSL code = {kSDF, wgSize}; + Kernel renderKernel = + createKernel(ctx, code, Bindings{readWrite(devScreen)}, + ceilDiv(Shape{NCOLS, NROWS, 1}, wgSize), params); printf("\033[2J\033[H"); while (true) { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, renderKernel, promise); + auto future = dispatchKernel(ctx, renderKernel); wait(ctx, future); - toCPU(ctx, devScreen, screen.data(), sizeof(screen)); + auto readback = toCPU(ctx, devScreen, screen); + wait(ctx, readback); params.time = getCurrentTimeInMilliseconds() - zeroTime; toGPU(ctx, params, renderKernel); - resetCommandBuffer(ctx.device, renderKernel); static const char intensity[] = "@B%8&WM#$Z0OQLCJUYX/" diff --git a/examples/shadertui/CMakeLists.txt b/examples/shadertui/CMakeLists.txt deleted file mode 100644 index 0938023..0000000 --- a/examples/shadertui/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -cmake_minimum_required(VERSION 3.28) -project(shadertui) - -set(FILENAME "gpu.hpp") - -get_filename_component(PROJECT_ROOT ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) -get_filename_component(PROJECT_ROOT ${PROJECT_ROOT} DIRECTORY) - -# Construct potential paths -set(FILEPATH_CURRENT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/${FILENAME}") -set(FILEPATH_PROJECT_ROOT "${PROJECT_ROOT}/${FILENAME}") - -# Check if the file exists in the current directory -if(EXISTS ${FILEPATH_CURRENT_DIR}) - set(TARGET_FILE_PATH ${CMAKE_CURRENT_SOURCE_DIR}) -elseif(EXISTS ${FILEPATH_PROJECT_ROOT}) - set(TARGET_FILE_PATH ${PROJECT_ROOT}) -else() - message(FATAL_ERROR "File ${FILENAME} not found in either ${CMAKE_CURRENT_SOURCE_DIR} or ${CMAKE_CURRENT_SOURCE_DIR}/../../") -endif() - -include("${TARGET_FILE_PATH}/cmake/example.cmake") \ No newline at end of file diff --git a/examples/shadertui/Makefile b/examples/shadertui/Makefile deleted file mode 100644 index 81c740b..0000000 --- a/examples/shadertui/Makefile +++ /dev/null @@ -1,28 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -TARGET=shadertui -CODEPATH = find . ../../utils ../../ -maxdepth 1 -type f -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn - - -run: ./build/$(TARGET) - $(LIBSPEC) && ./build/$(TARGET) - -# Use clang -v to see the include paths -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -o ./build/$(TARGET) - -watch: - @command -v entr >/dev/null 2>&1 || { echo >&2 "Please install entr with 'brew install entr' or 'sudo apt-get install entr'"; exit 1; } - mkdir -p build && $(CODEPATH) | entr -s "$(LIBSPEC) && rm -f ./build/$(TARGET) && make -j$(NUM_JOBS) run" - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* diff --git a/examples/shadertui/run.cpp b/examples/shadertui/run.cpp index 943180b..f8f80b0 100644 --- a/examples/shadertui/run.cpp +++ b/examples/shadertui/run.cpp @@ -1,7 +1,6 @@ #include #include #include -#include #include #include #include @@ -45,7 +44,6 @@ void loadKernelCode(const std::string &filename, std::string &codeString) { FILE *file = fopen(filename.c_str(), "r"); int nTries = 0; while (!file) { - fclose(file); std::this_thread::sleep_for(std::chrono::milliseconds(10)); file = fopen(filename.c_str(), "r"); if (++nTries > 5) { @@ -76,10 +74,7 @@ int main() { // std::fill(begin(screenArr), end(screenArr), 0.0); auto gen = std::mt19937{std::random_device{}()}; randint(screenArr, gen, 0, 1); - Tensor screen = createTensor(ctx, {kRows, kCols}, kf32, screenArr.data()); - - std::promise promise; - std::future future = promise.get_future(); + Tensor screen = createTensor(ctx, {kRows, kCols}, kf32, screenArr); std::string codeString; struct Params { @@ -91,10 +86,10 @@ int main() { LOG(kDefLog, kInfo, "Loading shader code from shader.wgsl"); loadKernelCode("shader.wgsl", codeString); - KernelCode shader{codeString.c_str(), Shape{16, 16, 1}}; + WGSL shader{codeString, Shape{16, 16, 1}}; Kernel renderKernel = - createKernel(ctx, shader, Bindings{screen}, - cdiv({kCols, kRows, 1}, shader.workgroupSize), params); + createKernel(ctx, shader, Bindings{readWrite(screen)}, + ceilDiv(Shape{kCols, kRows, 1}, shader.workgroupSize), params); LOG(kDefLog, kInfo, "Starting render loop"); @@ -109,15 +104,16 @@ int main() { while (true) { if (frame % framesPerLoad == 0) { loadKernelCode("shader.wgsl", codeString); - if (codeString != shader.data) { + if (codeString != shader.code) { // TODO(avh): Use a better way to avoid write/read race conditions // and recover from partial write errors std::this_thread::sleep_for(std::chrono::milliseconds(20)); loadKernelCode("shader.wgsl", codeString); - shader = {codeString.c_str(), Shape{16, 16, 1}}; + shader = {codeString, Shape{16, 16, 1}}; renderKernel = - createKernel(ctx, shader, Bindings{screen}, - cdiv({kCols, kRows, 1}, shader.workgroupSize), params); + createKernel(ctx, shader, Bindings{readWrite(screen)}, + ceilDiv(Shape{kCols, kRows, 1}, shader.workgroupSize), + params); ticks++; start = std::chrono::high_resolution_clock::now(); } @@ -126,17 +122,16 @@ int main() { params.time = getCurrentTimeInMilliseconds(start); toGPU(ctx, params, renderKernel); auto frameStart = std::chrono::high_resolution_clock::now(); - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, renderKernel, promise); + auto future = dispatchKernel(ctx, renderKernel); wait(ctx, future); - resetCommandBuffer(ctx.device, renderKernel); - toCPU(ctx, screen, screenArr); + auto readback = toCPU(ctx, screen, screenArr); + wait(ctx, readback); rasterize(screenArr, raster); auto frameEnd = std::chrono::high_resolution_clock::now(); std::chrono::duration frameElapsed = frameEnd - frameStart; elapsed = frameEnd - start; - std::this_thread::sleep_for(std::chrono::milliseconds(10) - frameElapsed); + auto delay = std::chrono::duration(10) - frameElapsed; + if (delay.count() > 0) std::this_thread::sleep_for(delay); printf("\033[H%s\nRender loop running (full screen recommended) ...\nEdit and save shader.wgsl to see changes here.\nReloaded shader.wgsl %zu times\n", raster.data(), ticks); fflush(stdout); } diff --git a/examples/shadertui/script.py b/examples/shadertui/script.py index 13c9c28..162b7d0 100644 --- a/examples/shadertui/script.py +++ b/examples/shadertui/script.py @@ -1,50 +1,17 @@ -import shutil -import os +from pathlib import Path +from shutil import copyfile -# List of shader files to copy -shader_files = [ - "default.wgsl", - "roundrect1.wgsl", - "roundrect2.wgsl", - "shapes.wgsl", - "boat.wgsl", - "gradient_flow.wgsl", - "wave_interference.wgsl", - "reaction_diffusion.wgsl", - "voronoi.wgsl", - "fluid.wgsl", - "aurora.wgsl", - "julia.wgsl", - "mandelbrot.wgsl", - "particles.wgsl", - "default.wgsl", -] # Add more file names as needed +shaders = ('default.wgsl roundrect1.wgsl roundrect2.wgsl shapes.wgsl boat.wgsl gradient_flow.wgsl wave_interference.wgsl ' + 'reaction_diffusion.wgsl voronoi.wgsl fluid.wgsl aurora.wgsl julia.wgsl mandelbrot.wgsl particles.wgsl default.wgsl').split() -def copy_file(src, dst): - try: - shutil.copy(src, dst) - print(f"Copied {src} to {dst}") - except IOError as e: - print(f"Unable to copy file. {e}") - except: - print("Unexpected error:", sys.exc_info()) +def main(): + root = Path(__file__).parent + print(f'Run `cd {root} && ../../build/shadertui_gpu` in another terminal.') + input('Press return to begin...') + for shader in shaders: + copyfile(root / shader, root / 'shader.wgsl') + input(f'Loaded {shader}. Press return to continue...') -if __name__ == "__main__": - # clear screen - os.system("cls" if os.name == "nt" else "clear") - print( - "\nThis script is meant to run alongside the shadertui runner. To start shadertui open a separate terminal and run `make` from this directory.\n" - ) - input("Press return/enter to continue...") - - for shader in shader_files: - if os.path.exists(shader): - copy_file(shader, "shader.wgsl") - input("Press return/enter to continue...") - else: - print(f"File {shader} does not exist.") - break - - print("All files processed.") +if __name__ == '__main__': main() diff --git a/examples/transpose/Makefile b/examples/transpose/Makefile deleted file mode 100644 index 1495c96..0000000 --- a/examples/transpose/Makefile +++ /dev/null @@ -1,27 +0,0 @@ -CXX=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -CODEPATH = find . ../../utils ../../ -maxdepth 1 -type f -TARGET=transpose -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif -FLAGS=-std=c++17 $(STDLIB) -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -L$(GPUCPP)/third_party/lib run.cpp -ldl -lwebgpu_dawn - -run: ./build/$(TARGET) - $(LIBSPEC) && ./build/$(TARGET) - -# Use clang -v to see the include paths -build/$(TARGET): run.cpp - mkdir -p build && $(CXX) $(FLAGS) -o ./build/$(TARGET) - -watch: - @command -v entr >/dev/null 2>&1 || { echo >&2 "Please install entr with 'brew install entr' or 'sudo apt-get install entr'"; exit 1; } - mkdir -p build && $(CODEPATH) | entr -s "$(LIBSPEC) && rm -f ./build/$(TARGET) && make -j$(NUM_JOBS) ./build/$(TARGET) && ./build/$(TARGET)" - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* diff --git a/examples/transpose/run.cpp b/examples/transpose/run.cpp index 4b0a28a..be658f8 100644 --- a/examples/transpose/run.cpp +++ b/examples/transpose/run.cpp @@ -4,20 +4,19 @@ #include #include -#include "gpu.hpp" // createContext, createTensor, createKernel, dispatchKernel, - // wait, resetCommandBuffer, toCPU +#include "gpu.hpp" #include "llmc/reference_impls.h" // for CPU reference implementation #include "utils/array_utils.hpp" // show, isclose, randn, randint #include "utils/logging.hpp" // LOG -#include "experimental/wgsl.h" // loopUnrolling +#include "utils/wgsl.hpp" using namespace gpu; // This implements the tranpose kernels in https://developer.nvidia.com/blog/efficient-matrix-transpose-cuda-cc . static const char *kShaderTranspose1 = R"( -@group(0) @binding(0) var A: array<{{precision}}>; +@group(0) @binding(0) var A: array<{{precision}}>; @group(0) @binding(1) var B: array<{{precision}}>; @compute @workgroup_size({{workgroupSize}}) fn main( @@ -28,7 +27,7 @@ fn main( } )"; -inline KernelCode createTranspose1(const char *shaderTemplate, +inline WGSL createTranspose1(const char *shaderTemplate, const size_t M, const size_t N, const Shape &workgroupSize = {256, 1, 1}, NumType precision = kf32) { @@ -42,7 +41,7 @@ inline KernelCode createTranspose1(const char *shaderTemplate, // Shared memory cache-blocking static const char *kShaderTranspose2 = R"( -@group(0) @binding(0) var A: array<{{precision}}>; +@group(0) @binding(0) var A: array<{{precision}}>; @group(0) @binding(1) var B: array<{{precision}}>; var tile: array<{{precision}}, {{BN}} * {{BM}}>; @compute @workgroup_size({{workgroupSize}}) @@ -78,7 +77,7 @@ fn main( } )"; -inline KernelCode createTranspose2(const char *shaderTemplate, +inline WGSL createTranspose2(const char *shaderTemplate, const size_t M, const size_t N, const size_t BM, const size_t BN, const size_t TM, const size_t TN, @@ -98,8 +97,7 @@ inline KernelCode createTranspose2(const char *shaderTemplate, {"{{TM}}", toString(TM)}, {"{{TN}}", toString(TN)} }); - std::string unrolledCode = codeString ;// loopUnrolling(codeString); - return {unrolledCode, workgroupSize}; + return {codeString, workgroupSize}; } void initData(size_t M, size_t N, std::unique_ptr &inputPtr) { @@ -115,10 +113,10 @@ Kernel selectTranspose(Context &ctx, int version, if (version == 1) { Shape wgSize = {16, 16, 1}; LOG(kDefLog, kInfo, "wgSize: %s", toString(wgSize).c_str()); - KernelCode transpose = + WGSL transpose = createTranspose1(kShaderTranspose1, M, N, /*wgsize*/ wgSize); // The shape of input == M x N kernel = createKernel(ctx, transpose, bindings, - /*nWorkgroups*/ cdiv({N, M, 1}, wgSize)); // The shape of output == N x M + /*nWorkgroups*/ ceilDiv({N, M, 1}, wgSize)); // The shape of output == N x M } else if (version == 2) { static constexpr size_t BM = 64; static constexpr size_t BK = 16; @@ -126,12 +124,12 @@ Kernel selectTranspose(Context &ctx, int version, static constexpr size_t TM = BM / BK; static constexpr size_t TN = BN / BK; Shape wgSize = {(BM / TM) * (BN / TN), 1, 1}; // This is the same as BK * BK. - Shape nWorkgroups = {cdiv(N, BN), cdiv(M, BM), 1}; + Shape nWorkgroups = {ceilDiv(N, BN), ceilDiv(M, BM), 1}; LOG(kDefLog, kInfo, "M: %d, N: %d", M, N); LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d, TN: %d", BM, BK, BN, TM, TN); LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); - KernelCode transpose = createTranspose2(kShaderTranspose2, M, N, BM, BN, TM, TN, + WGSL transpose = createTranspose2(kShaderTranspose2, M, N, BM, BN, TM, TN, /*wgSize*/ wgSize, kf32); kernel = createKernel(ctx, transpose, bindings, @@ -149,34 +147,28 @@ void runTest(int version, size_t M, size_t N, // Allocate GPU buffers and copy data Context ctx = createContext(); - Tensor input = createTensor(ctx, Shape{M, N}, kf32, inputPtr.get()); + Tensor input = createTensor( + ctx, Shape{M, N}, kf32, + std::span(inputPtr.get(), M * N)); Tensor output = createTensor(ctx, Shape{N, M}, kf32); constexpr size_t nIter = 50; // Initialize Kernel and bind GPU buffers LOG(kDefLog, kInfo, "Creating Kernel"); - Kernel kernel = selectTranspose(ctx, version, {input, output}, M, N); + Kernel kernel = selectTranspose( + ctx, version, Bindings{read(input), readWrite(output)}, M, N); // Dispatch kernel execution LOG(kDefLog, kInfo, "Dispatching Kernel version %d, %d iterations ...", version, nIter); - // pre-allocate promises and futures for async dispatch - // TODO(avh): implement a pooling mechanism for promises/futures in gpu.h - std::array, nIter> promises; - std::array, nIter> futures; - for (int i = 0; i < nIter; i++) { - futures[i] = promises[i].get_future(); - } - // Dispatch kernel nIter times auto start = std::chrono::high_resolution_clock::now(); for (int i = 0; i < nIter; i++) { if (!isCPU) { - dispatchKernel(ctx, kernel, promises[i]); - wait(ctx, futures[i]); - resetCommandBuffer(ctx.device, kernel); + auto future = dispatchKernel(ctx, kernel); + wait(ctx, future); } else { transpose(inputPtr.get(), outputPtr.get(), M, N); } @@ -193,7 +185,9 @@ void runTest(int version, size_t M, size_t N, LOG(kDefLog, kInfo, "Copying result to CPU"); if (!isCPU) { - toCPU(ctx, output, outputPtr.get(), M * N * sizeof(float)); + auto readback = + toCPU(ctx, output, std::span(outputPtr.get(), M * N)); + wait(ctx, readback); } LOG(kDefLog, kInfo, "%s", show(outputPtr.get(), N, M, "Output").c_str()); diff --git a/experimental/kernels/Makefile b/experimental/kernels/Makefile deleted file mode 100644 index e2d89b1..0000000 --- a/experimental/kernels/Makefile +++ /dev/null @@ -1,154 +0,0 @@ -CXX=clang++ -CC=clang++ -GPUCPP ?= $(PWD)/../.. -LIBDIR ?= $(GPUCPP)/third_party/lib -LIBSPEC ?= . $(GPUCPP)/source -NUM_JOBS?=$(shell nproc) -ifeq ($(shell $(CXX) -std=c++17 -x c++ -E -include array - < /dev/null > /dev/null 2>&1 ; echo $$?),0) - STDLIB := -else - STDLIB := -stdlib=libc++ -endif - -# ASYNCIFY allows emscripten to sleep -EMFLAGS=-std=c++17 -I$(GPUCPP) -I$(GPUCPP)/third_party/headers/wasm -I. -Iunittest_llmc -I$(GPUCPP)/third_party/llm.c -s USE_WEBGPU=1 -s ASYNCIFY=1 -s STACK_SIZE=100000 -s MEMORY64=1 -s ALLOW_MEMORY_GROWTH=1 -CXXFLAGS=-std=c++17 -I$(GPUCPP) -I$(GPUCPP)/third_party/headers -I. -Iunittest_llmc -CFLAGS=-Ofast -march=native -I. -Iunittest_llmc -# CFLAGS=-O2 -march=native -I. -Iunittest_llmc - -LDFLAGS=$(STDLIB) -L$(GPUCPP)/third_party/lib -ldl -lwebgpu_dawn -fsanitize=address -FLAGS=$(CXXFLAGS) $(LDFLAGS) - -ifeq ($(shell [ -d /opt/homebrew/opt/libomp/lib ] && echo "exists"), exists) - CFLAGS += -Xclang -fopenmp -DOMP -I/opt/homebrew/opt/libomp/include - LDFLAGS += -L/opt/homebrew/opt/libomp/lib -lomp - $(info ✓ OpenMP found) -else - $(info ✗ OpenMP not found) -endif - -default: run-native - -build/reduce: reduce.cpp kernels.h - $(CC) $(CFLAGS) $(CXXFLAGS) $(LDFLAGS) -o $@ $< - $(LIBSPEC) && build/reduce - -run_llm.c: ./build/test_gpt2 dawnlib - $(LIBSPEC) && $< - -run_llm.c_with_metal_profiler: ./build/test_gpt2_with_metal_profiler dawnlib - $(LIBSPEC) && export METAL_CAPTURE_ENABLED=1 && $< - -run_llm.c_with_time_profiler: ./build/test_gpt2_with_metal_profiler dawnlib - $(LIBSPEC) && xcrun xctrace record --template 'Time Profiler' --launch -- $< - -run_llm.c_train: ./build/train_gpt2 dawnlib - if [ ! -d dev ] ; then ln -s $(GPUCPP)/third_party/llm.c/dev ; fi - if [ ! -f gpt2_tokenizer.bin ] ; then ln -s $(GPUCPP)/third_party/llm.c/gpt2_tokenizer.bin ; fi - $(LIBSPEC) && $< - -llm.c: - # if [ ! -d llm.c ]; then git clone git@github.com:karpathy/llm.c.git ; fi - ln -s $(GPUCPP)/third_party/llm.c - -gpt2_124M.bin: llm.c - if [ ! -f $@ ]; then ./llm.c/dev/download_starter_pack.sh ; \ - ln -s ./llm.c/gpt2_124M.bin ; \ - ln -s ./llm.c/gpt2_124M_debug_state.bin ; \ - ln -s ./llm.c/gpt2_tokenizer.bin ; \ - fi - -define preprocess_file - sed -i -e 's/int main(/int MAIN(/g' llm.c/test_gpt2.c - sed -i -e 's/int main(/int MAIN(/g' llm.c/train_gpt2.c - sed -i -e 's/void encoder_forward(/void ENCODER_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void layernorm_forward(/void LAYERNORM_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void matmul_forward(/void MATMUL_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void attention_forward(/void ATTENTION_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void gelu_forward(/void GELU_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void residual_forward(/void RESIDUAL_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void softmax_forward(/void SOFTMAX_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void crossentropy_forward(/void CROSSENTROPY_FORWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void encoder_backward(/void ENCODER_BACKWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void layernorm_backward(/void LAYERNORM_BACKWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void matmul_backward(/void MATMUL_BACKWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void attention_backward(/void ATTENTION_BACKWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void gelu_backward(/void GELU_BACKWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void residual_backward(/void RESIDUAL_BACKWARD_CPU(/g' llm.c/train_gpt2.c - sed -i -e 's/void crossentropy_softmax_backward(/void CROSSENTROPY_SOFTMAX_BACKWARD_CPU(/g' llm.c/train_gpt2.c - grep -q "^#include \"unittest_kernels.h\"" llm.c/train_gpt2.c || \ - printf '1i\n#include "unittest_kernels.h"\n.\nw\nq\n' | ed -s llm.c/train_gpt2.c -endef - -build/test_gpt2: llm.c build/unittest_kernels.o gpt2_124M.bin - mkdir -p build - $(call preprocess_file) - $(CC) $(CFLAGS) $(LDFLAGS) -o $@ llm.c/test_gpt2.c build/unittest_kernels.o - -build/test_gpt2_with_metal_profiler: llm.c build/unittest_kernels.o gpt2_124M.bin - mkdir -p build - $(call preprocess_file) - $(CC) $(CFLAGS) $(LDFLAGS) -o $@ llm.c/test_gpt2.c build/unittest_kernels.o -I$(GPUCPP) $(GPUCPP)/experimental/profiler/metal.mm -framework metal -framework Foundation -DMETAL_PROFILER -g - install_name_tool -change @rpath/libwebgpu_dawn.dylib $(GPUCPP)/third_party/lib/libwebgpu_dawn.dylib $@ - -build/train_gpt2: llm.c build/unittest_kernels.o gpt2_124M.bin - mkdir -p build - $(call preprocess_file) - $(CC) $(CFLAGS) $(LDFLAGS) -o $@ llm.c/train_gpt2.c build/unittest_kernels.o - -build/ops.o: ops.cpp ops.hpp kernels.h llm.c - mkdir -p build && $(CXX) $(CXXFLAGS) -c -o $@ $< - -build/gpt2_webgpu: llm.c gpt2_124M.bin llm.c gpt2_webgpu.cpp ops.cpp - mkdir -p build - $(CC) $(CXXFLAGS) -Illm.c $(LDFLAGS) -o $@ gpt2_webgpu.cpp ops.cpp - -build/gpt2_webgpu_aot: llm.c gpt2_124M.bin llm.c gpt2_webgpu_aot.cpp ops_aot.cpp - mkdir -p build - $(CC) $(CXXFLAGS) -Illm.c $(LDFLAGS) -o $@ gpt2_webgpu_aot.cpp ops_aot.cpp - -build/gpt2_webgpu.html: check-emsdk gpt2_webgpu.cpp term.html llm.c - em++ gpt2_webgpu.cpp ops.cpp \ - --preload-file gpt2_tokenizer.bin@/gpt2_tokenizer.bin \ - --preload-file gpt2_124M.bin@/gpt2_124M.bin \ - --preload-file gpt2_124M_debug_state.bin@/gpt2_124M_debug_state.bin \ - --preload-file llm.c/dev/data/tinyshakespeare/tiny_shakespeare_train.bin@dev/data/tinyshakespeare/tiny_shakespeare_train.bin \ - --preload-file llm.c/dev/data/tinyshakespeare/tiny_shakespeare_val.bin@dev/data/tinyshakespeare/tiny_shakespeare_val.bin \ - -o build/gpt2_webgpu.html \ - $(EMFLAGS) \ - --shell-file term.html \ - -watch-web: - ls *.cpp *.c *.hpp *.h | entr -s make build/gpt2_webgpu.html - -watch-native: - ls *.cpp *.c *.hpp *.h | entr -s "rm -f build/gpt2_webgpu && rm -f build/ops.o && make build/gpt2_webgpu" - -run-native: build/gpt2_webgpu_aot - . $(GPUCPP)/source && ./build/gpt2_webgpu_aot - -# server: build/train_gpt2.html build/test_gpt2.html build/gpt2_gpucpp.html -server: build/gpt2_webgpu.html - @echo "\n┌───────────────────────────────────────────────────────────────────────────────────┐" - @echo "│ Open http://localhost:8000/build/run.html in your browser to see the output. │" - @echo "│ │" - @echo "│ Press Ctrl+C to stop the server. │" - @echo "└───────────────────────────────────────────────────────────────────────────────────┘\n\n" - python3 -m http.server --directory . - -build/unittest_kernels.o: unittest_llmc/unittest_kernels.cpp unittest_llmc/unittest_kernels.h kernels.h - mkdir -p build && $(CXX) $(CXXFLAGS) -DNDEBUG -c -o $@ $< - -dawnlib: $(if $(wildcard $(GPUCPP)/third_party/lib/libwebgpu_dawn.so $(GPUCPP)/third_party/lib/libwebgpu_dawn.dylib),,run_setup) - -run_setup: check-python - cd $(GPUCPP) && python3 setup.py - -clean: - read -r -p "This will delete the contents of build/*. Are you sure? [CTRL-C to abort] " response && rm -rf build/* - -check-python: - @command -v python3 >/dev/null 2>&1 || { echo >&2 "Python needs to be installed and in your path."; exit 1; } - -check-emsdk: - @which em++ > /dev/null || (echo "emsdk not found. Please install emsdk and run 'source emsdk_env.sh' in the emsdk directory." && exit 1) diff --git a/experimental/kernels/dev b/experimental/kernels/dev deleted file mode 120000 index 1162e4b..0000000 --- a/experimental/kernels/dev +++ /dev/null @@ -1 +0,0 @@ -../../third_party/llm.c/dev \ No newline at end of file diff --git a/experimental/kernels/gpt2_webgpu.cpp b/experimental/kernels/gpt2_webgpu.cpp deleted file mode 100644 index 8e71541..0000000 --- a/experimental/kernels/gpt2_webgpu.cpp +++ /dev/null @@ -1,769 +0,0 @@ -#include "gpu.hpp" -#include "ops.hpp" -/* -This file trains the GPT-2 model. -This version is the clean, minimal, reference. As such: -- it runs on CPU. -- it does not make the code too complex; it is readable. -- it does not use any processor-specific instructions, intrinsics and such. -- it _does_ use a few OpenMP pragmas because this is a large speedup at very low cost -There will be other versions of this code that specialize it and make it fast. -*/ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#ifdef OMP -#include -#endif -// our own utilities -// defines: fopenCheck, freadCheck, fcloseCheck, fseekCheck, mallocCheck -#include "llmc/utils.h" -// defines: tokenizer_init, tokenizer_decode, tokenizer_free -#include "llmc/tokenizer.h" -// defines: dataloader_init, dataloader_reset, dataloader_next_batch, dataloader_free -#include "llmc/dataloader.h" - -using namespace gpu; - -// ---------------------------------------------------------------------------- -// GPT-2 model definition - -typedef struct { - int max_seq_len; // max sequence length, e.g. 1024 - int vocab_size; // vocab size, e.g. 50257 - int padded_vocab_size; // padded to e.g. %128==0, 50304 - int num_layers; // number of layers, e.g. 12 - int num_heads; // number of heads in attention, e.g. 12 - int channels; // number of channels, e.g. 768 -} GPT2Config; - -// the parameters of the model -#define NUM_PARAMETER_TENSORS 16 -typedef struct { - float* wte; // (V, C) - float* wpe; // (maxT, C) - float* ln1w; // (L, C) - float* ln1b; // (L, C) - float* qkvw; // (L, 3*C, C) - float* qkvb; // (L, 3*C) - float* attprojw; // (L, C, C) - float* attprojb; // (L, C) - float* ln2w; // (L, C) - float* ln2b; // (L, C) - float* fcw; // (L, 4*C, C) - float* fcb; // (L, 4*C) - float* fcprojw; // (L, C, 4*C) - float* fcprojb; // (L, C) - float* lnfw; // (C) - float* lnfb; // (C) -} ParameterTensors; - -void fill_in_parameter_sizes(size_t* param_sizes, GPT2Config config) { - size_t Vp = config.padded_vocab_size; - size_t C = config.channels; - size_t maxT = config.max_seq_len; - size_t L = config.num_layers; - param_sizes[0] = Vp * C; // wte - param_sizes[1] = maxT * C; // wpe - param_sizes[2] = L * C; // ln1w - param_sizes[3] = L * C; // ln1b - param_sizes[4] = L * (3 * C) * C; // qkvw - param_sizes[5] = L * (3 * C); // qkvb - param_sizes[6] = L * C * C; // attprojw - param_sizes[7] = L * C; // attprojb - param_sizes[8] = L * C; // ln2w - param_sizes[9] = L * C; // ln2b - param_sizes[10] = L * (4 * C) * C; // fcw - param_sizes[11] = L * (4 * C); // fcb - param_sizes[12] = L * C * (4 * C); // fcprojw - param_sizes[13] = L * C; // fcprojb - param_sizes[14] = C; // lnfw - param_sizes[15] = C; // lnfb -} - -// allocate memory for the parameters and point the individual tensors to the right places -float* malloc_and_point_parameters(ParameterTensors* params, size_t* param_sizes) { - size_t num_parameters = 0; - for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { - num_parameters += param_sizes[i]; - } - // malloc all parameters all at once - float* params_memory = (float*)mallocCheck(num_parameters * sizeof(float)); - // assign all the tensors - float** ptrs[] = { - ¶ms->wte, ¶ms->wpe, ¶ms->ln1w, ¶ms->ln1b, ¶ms->qkvw, ¶ms->qkvb, - ¶ms->attprojw, ¶ms->attprojb, ¶ms->ln2w, ¶ms->ln2b, ¶ms->fcw, ¶ms->fcb, - ¶ms->fcprojw, ¶ms->fcprojb, ¶ms->lnfw, ¶ms->lnfb - }; - float* params_memory_iterator = params_memory; - for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { - *(ptrs[i]) = params_memory_iterator; - params_memory_iterator += param_sizes[i]; - } - return params_memory; -} - - -#define NUM_ACTIVATION_TENSORS 23 -typedef struct { - float* encoded; // (B, T, C) - float* ln1; // (L, B, T, C) - float* ln1_mean; // (L, B, T) - float* ln1_rstd; // (L, B, T) - float* qkv; // (L, B, T, 3*C) - float* atty; // (L, B, T, C) - float* preatt; // (L, B, NH, T, T) - float* att; // (L, B, NH, T, T) - float* attproj; // (L, B, T, C) - float* residual2; // (L, B, T, C) - float* ln2; // (L, B, T, C) - float* ln2_mean; // (L, B, T) - float* ln2_rstd; // (L, B, T) - float* fch; // (L, B, T, 4*C) - float* fch_gelu; // (L, B, T, 4*C) - float* fcproj; // (L, B, T, C) - float* residual3; // (L, B, T, C) - float* lnf; // (B, T, C) - float* lnf_mean; // (B, T) - float* lnf_rstd; // (B, T) - float* logits; // (B, T, V) - float* probs; // (B, T, V) - float* losses; // (B, T) -} ActivationTensors; - - - -void fill_in_activation_sizes(size_t* act_sizes, GPT2Config config, int B, int T) { - size_t C = config.channels; - size_t NH = config.num_heads; - size_t L = config.num_layers; - size_t Vp = config.padded_vocab_size; - act_sizes[0] = B * T * C; // encoded - act_sizes[1] = L * B * T * C; // ln1 - act_sizes[2] = L * B * T; // ln1_mean - act_sizes[3] = L * B * T; // ln1_rstd - act_sizes[4] = L * B * T * 3 * C; // qkv - act_sizes[5] = L * B * T * C; // atty - act_sizes[6] = L * B * NH * T * T; // preatt - act_sizes[7] = L * B * NH * T * T; // att - act_sizes[8] = L * B * T * C; // attproj - act_sizes[9] = L * B * T * C; // residual2 - act_sizes[10] = L * B * T * C; // ln2 - act_sizes[11] = L * B * T; // ln2_mean - act_sizes[12] = L * B * T; // ln2_rstd - act_sizes[13] = L * B * T * 4 * C; // fch - act_sizes[14] = L * B * T * 4 * C; // fch_gelu - act_sizes[15] = L * B * T * C; // fcproj - act_sizes[16] = L * B * T * C; // residual3 - act_sizes[17] = B * T * C; // lnf - act_sizes[18] = B * T; // lnf_mean - act_sizes[19] = B * T; // lnf_rstd - act_sizes[20] = B * T * Vp; // logits - act_sizes[21] = B * T * Vp; // probs - act_sizes[22] = B * T; // losses -} - -float* malloc_and_point_activations(ActivationTensors* acts, size_t* act_sizes) { - size_t num_activations = 0; - for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) { - num_activations += act_sizes[i]; - } - float* acts_memory = (float*)mallocCheck(num_activations * sizeof(float)); - float** ptrs[] = { - &acts->encoded, &acts->ln1, &acts->ln1_mean, &acts->ln1_rstd, &acts->qkv, &acts->atty, - &acts->preatt, &acts->att, &acts->attproj, &acts->residual2, &acts->ln2, &acts->ln2_mean, - &acts->ln2_rstd, &acts->fch, &acts->fch_gelu, &acts->fcproj, &acts->residual3, &acts->lnf, - &acts->lnf_mean, &acts->lnf_rstd, &acts->logits, &acts->probs, &acts->losses - }; - float* acts_memory_iterator = acts_memory; - for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) { - *(ptrs[i]) = acts_memory_iterator; - acts_memory_iterator += act_sizes[i]; - } - return acts_memory; -} - -struct GPUParameters { - Tensor data[NUM_PARAMETER_TENSORS]; -}; - -struct GPUActivations { - Tensor data[NUM_ACTIVATION_TENSORS]; -}; - - -void gpu_alloc(Context& ctx, Tensor* tensors, size_t* sizes, size_t n) { - for (size_t i = 0; i < n; i++) { - tensors[i] = createTensor(ctx, Shape{sizes[i]}, kf32); - } -} - -typedef struct { - GPT2Config config; - // the weights (parameters) of the model, and their sizes - ParameterTensors params; - GPUParameters params_; // TODO(avh): eventually this replaces params - size_t param_sizes[NUM_PARAMETER_TENSORS]; - float* params_memory; - size_t num_parameters; - // gradients of the weights - ParameterTensors grads; - float* grads_memory; - // buffers for the AdamW optimizer - float* m_memory; - float* v_memory; - // the activations of the model, and their sizes - ActivationTensors acts; - GPUActivations acts_; // TODO(avh): eventually this replaces params - size_t act_sizes[NUM_ACTIVATION_TENSORS]; - float* acts_memory; - size_t num_activations; - // gradients of the activations - ActivationTensors grads_acts; - float* grads_acts_memory; - // other run state configuration - int batch_size; // the batch size (B) of current forward pass - int seq_len; // the sequence length (T) of current forward pass - int* inputs; // the input tokens for the current forward pass - int* targets; // the target tokens for the current forward pass - float mean_loss; // after a forward pass with targets, will be populated with the mean loss -} GPT2; - -void gpt2_build_from_checkpoint(Context& ctx, GPT2 *model, const char* checkpoint_path) { - - // read in model from a checkpoint file - FILE *model_file = fopenCheck(checkpoint_path, "rb"); - int model_header[256]; - freadCheck(model_header, sizeof(int), 256, model_file); - if (model_header[0] != 20240326) { printf("Bad magic model file\n"); exit(1); } - if (model_header[1] != 3) { - printf("Bad version in model file\n"); - printf("---> HINT: try to re-run `python train_gpt2.py`\n"); - exit(1); - } - - // read in hyperparameters - size_t maxT, V, Vp, L, NH, C; // size_t to prevent int overflow - model->config.max_seq_len = maxT = model_header[2]; - model->config.vocab_size = V = model_header[3]; -#ifdef __EMSCRIPTEN__ - model->config.num_layers = L = 12; // TODO(avh): Debugging only hack - revert this -#else - model->config.num_layers = L = model_header[4]; -#endif - model->config.num_heads = NH = model_header[5]; - model->config.channels = C = model_header[6]; - model->config.padded_vocab_size = Vp = model_header[7]; - printf("[GPT-2]\n"); - printf("max_seq_len: %zu\n", maxT); - printf("vocab_size: %zu\n", V); - printf("padded_vocab_size: %zu\n", Vp); - printf("num_layers: %zu\n", L); - printf("num_heads: %zu\n", NH); - printf("channels: %zu\n", C); - - // allocate space for all the parameters and read them in - fill_in_parameter_sizes(model->param_sizes, model->config); - - // count the number of parameters - size_t num_parameters = 0; - for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { - num_parameters += model->param_sizes[i]; - } - printf("num_parameters: %zu\n", num_parameters); - model->num_parameters = num_parameters; - - // read in all the parameters from file - model->params_memory = malloc_and_point_parameters(&model->params, model->param_sizes); - freadCheck(model->params_memory, sizeof(float), num_parameters, model_file); - fcloseCheck(model_file); - - // other inits - model->acts_memory = NULL; - model->grads_memory = NULL; - model->m_memory = NULL; - model->v_memory = NULL; - model->grads_acts_memory = NULL; - model->inputs = NULL; - model->targets = NULL; - model->batch_size = 0; - model->seq_len = 0; - model->mean_loss = -1.0f; // -1.0f will designate no loss - - // TODO(avh): this is just a resource test for now, eventually deprecate CPU allocations - gpu_alloc(ctx, model->params_.data, model->param_sizes, NUM_PARAMETER_TENSORS); - -} - - -void gpt2_forward(Context& ctx, GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { - // targets are optional and could be NULL - - // ensure the model was initialized or error out - if (model->params_memory == NULL) { - printf("Error: model was not initialized properly.\n"); - exit(1); - } - - // convenience parameters (size_t to help prevent int overflow) - size_t V = model->config.vocab_size; - size_t Vp = model->config.padded_vocab_size; - size_t L = model->config.num_layers; - size_t NH = model->config.num_heads; - size_t C = model->config.channels; - - // validate inputs, all indices must be in the range [0, V) - for(int i = 0; i < B * T; i++) { - assert(0 <= inputs[i] && inputs[i] < V); - if (targets != NULL) { - assert(0 <= targets[i] && targets[i] < V); - } - } - - // allocate space for all the activations if needed (done here, lazily) - if(model->acts_memory == NULL) { - // record the current B,T as well - model->batch_size = B; - model->seq_len = T; - // and now allocate the space - fill_in_activation_sizes(model->act_sizes, model->config, B, T); - // TODO(avh): this is just a resource test for now, eventually deprecate CPU allocations - gpu_alloc(ctx, model->acts_.data, model->act_sizes, NUM_PARAMETER_TENSORS); - size_t num_activations = 0; - for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) { - num_activations += model->act_sizes[i]; - } - printf("num_activations: %zu\n", num_activations); - model->num_activations = num_activations; - printf("Allocating %.2f MB for activations\n", num_activations * sizeof(float) / (1024.0f * 1024.0f)); - model->acts_memory = malloc_and_point_activations(&model->acts, model->act_sizes); - // also create memory for caching inputs and targets - model->inputs = (int*)mallocCheck(B * T * sizeof(int)); - model->targets = (int*)mallocCheck(B * T * sizeof(int)); // might be unused if we never have targets but it's small - } else { - // validate B,T is consistent with how we've allocated the memory before - // in principle we could get more clever here in the future, for now this is safest - if (B != model->batch_size || T != model->seq_len) { - printf("Model: B=%d T=%d, Desired: B=%d T=%d\n", model->batch_size, model->seq_len, (int)B, (int)T); - exit(EXIT_FAILURE); - } - } - - printf("Cache inputs/targets\n"); - // cache the inputs/targets - memcpy(model->inputs, inputs, B * T * sizeof(int)); - if (targets != NULL) { - memcpy(model->targets, targets, B * T * sizeof(int)); - } - - printf("Forward pass\n"); - // forward pass - ParameterTensors params = model->params; // for brevity - ActivationTensors acts = model->acts; - float* residual; - printf("Encoding\n"); - printf("inputs[0] = %d\n", inputs[0]); - encoder_forward(ctx, acts.encoded, inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0] - for (int l = 0; l < L; l++) { - printf("Forward Pass Layer %d\n", l); - - residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; - - // get the pointers of the weights for this layer - float* l_ln1w = params.ln1w + l * C; - float* l_ln1b = params.ln1b + l * C; - float* l_qkvw = params.qkvw + l * 3*C * C; - float* l_qkvb = params.qkvb + l * 3*C; - float* l_attprojw = params.attprojw + l * C * C; - float* l_attprojb = params.attprojb + l * C; - float* l_ln2w = params.ln2w + l * C; - float* l_ln2b = params.ln2b + l * C; - float* l_fcw = params.fcw + l * 4*C * C; - float* l_fcb = params.fcb + l * 4*C; - float* l_fcprojw = params.fcprojw + l * C * 4*C; - float* l_fcprojb = params.fcprojb + l * C; - - // get the pointers of the activations for this layer - float* l_ln1 = acts.ln1 + l * B * T * C; - float* l_ln1_mean = acts.ln1_mean + l * B * T; - float* l_ln1_rstd = acts.ln1_rstd + l * B * T; - float* l_qkv = acts.qkv + l * B * T * 3*C; - float* l_atty = acts.atty + l * B * T * C; - float* l_preatt = acts.preatt + l * B * NH * T * T; - float* l_att = acts.att + l * B * NH * T * T; - float* l_attproj = acts.attproj + l * B * T * C; - float* l_residual2 = acts.residual2 + l * B * T * C; - float* l_ln2 = acts.ln2 + l * B * T * C; - float* l_ln2_mean = acts.ln2_mean + l * B * T; - float* l_ln2_rstd = acts.ln2_rstd + l * B * T; - float* l_fch = acts.fch + l * B * T * 4*C; - float* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C; - float* l_fcproj = acts.fcproj + l * B * T * C; - float* l_residual3 = acts.residual3 + l * B * T * C; - - // now do the forward pass - printf(" [Forward] : LayerNorm1\n"); - layernorm_forward(ctx, l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C); - printf(" [Forward] : QKV Projection\n"); - matmul_forward(ctx, l_qkv, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C); - printf(" [Forward] : Attention\n"); - attention_forward(ctx, l_atty, l_preatt, l_att, l_qkv, B, T, C, NH); - printf(" [Forward] : Attention Projection\n"); - matmul_forward(ctx, l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C); - printf(" [Forward] : Residual1\n"); - residual_forward(ctx, l_residual2, residual, l_attproj, B*T*C); - printf(" [Forward] : LayerNorm2\n"); - layernorm_forward(ctx, l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C); - printf(" [Forward] : FF Up\n"); - matmul_forward(ctx, l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C); - printf(" [Forward] : GELU\n"); - gelu_forward(ctx, l_fch_gelu, l_fch, B*T*4*C); - printf(" [Forward] : FF Down\n"); - matmul_forward(ctx, l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C); - printf(" [Forward] : Residual2\n"); - residual_forward(ctx, l_residual3, l_residual2, l_fcproj, B*T*C); - } - residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 - layernorm_forward(ctx, acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, params.lnfw, params.lnfb, B, T, C); - matmul_forward(ctx, acts.logits, acts.lnf, params.wte, NULL, B, T, C, Vp); - softmax_forward(ctx, acts.probs, acts.logits, B, T, V, Vp); - - printf("Crossentropy\n"); - // also forward the cross-entropy loss function if we have the targets - if (targets != NULL) { - crossentropy_forward(ctx, model->acts.losses, model->acts.probs, targets, B, T, Vp); - // for convenience also evaluate the mean loss - float mean_loss = 0.0f; - for (int i=0; iacts.losses[i]; } - mean_loss /= B*T; - model->mean_loss = mean_loss; - } else { - // if we don't have targets, we don't have a loss - model->mean_loss = -1.0f; - } - printf("Forward pass done\n"); -} - -void gpt2_zero_grad(GPT2 *model) { - if(model->grads_memory != NULL) { memset(model->grads_memory, 0, model->num_parameters * sizeof(float)); } - if(model->grads_acts_memory != NULL) { memset(model->grads_acts_memory, 0, model->num_activations * sizeof(float)); } -} - -void gpt2_backward(Context& ctx, GPT2 *model) { - printf("Backward pass\n"); - - // double check we forwarded previously, with targets - if (model->mean_loss == -1.0f) { - printf("Error: must forward with targets before backward\n"); - exit(1); - } - - // lazily allocate the memory for gradients of the weights and activations, if needed - if (model->grads_memory == NULL) { - printf("Allocating %.2f MB for gradients\n", model->num_parameters * sizeof(float) / (1024.0f * 1024.0f)); - model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_sizes); - model->grads_acts_memory = malloc_and_point_activations(&model->grads_acts, model->act_sizes); - gpt2_zero_grad(model); - } - - // convenience shortcuts (and size_t to help prevent int overflow) - size_t B = model->batch_size; - size_t T = model->seq_len; - size_t V = model->config.vocab_size; - size_t Vp = model->config.padded_vocab_size; - size_t L = model->config.num_layers; - size_t NH = model->config.num_heads; - size_t C = model->config.channels; - - // backward pass: go in the reverse order of the forward pass, and call backward() functions - ParameterTensors params = model->params; // for brevity - ParameterTensors grads = model->grads; - ActivationTensors acts = model->acts; - ActivationTensors grads_acts = model->grads_acts; - - // we kick off the chain rule by filling in dlosses with 1.0f/(B*T) - // technically this is a small, inline backward() pass of calculating - // total, final loss as the mean over all losses over all (B,T) positions in the batch - float dloss_mean = 1.0f / (B*T); - for (int i = 0; i < B*T; i++) { grads_acts.losses[i] = dloss_mean; } - - crossentropy_softmax_backward(ctx, grads_acts.logits, grads_acts.losses, acts.probs, model->targets, B, T, V, Vp); - matmul_backward(ctx, grads_acts.lnf, grads.wte, NULL, grads_acts.logits, acts.lnf, params.wte, B, T, C, Vp); - float* residual = acts.residual3 + (L-1) * B * T * C; // last layer's residual - float* dresidual = grads_acts.residual3 + (L-1) * B * T * C; // write to last layer's residual - layernorm_backward(ctx, dresidual, grads.lnfw, grads.lnfb, grads_acts.lnf, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C); - - for (int l = L-1; l >= 0; l--) { - printf("Backward Pass Layer %d\n", l); - - residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; - dresidual = l == 0 ? grads_acts.encoded : grads_acts.residual3 + (l-1) * B * T * C; - - // get the pointers of the weights for this layer - float* l_ln1w = params.ln1w + l * C; - float* l_qkvw = params.qkvw + l * 3*C * C; - float* l_attprojw = params.attprojw + l * C * C; - float* l_ln2w = params.ln2w + l * C; - float* l_fcw = params.fcw + l * 4*C * C; - float* l_fcprojw = params.fcprojw + l * C * 4*C; - // get the pointers of the gradients of the weights for this layer - float* dl_ln1w = grads.ln1w + l * C; - float* dl_ln1b = grads.ln1b + l * C; - float* dl_qkvw = grads.qkvw + l * 3*C * C; - float* dl_qkvb = grads.qkvb + l * 3*C; - float* dl_attprojw = grads.attprojw + l * C * C; - float* dl_attprojb = grads.attprojb + l * C; - float* dl_ln2w = grads.ln2w + l * C; - float* dl_ln2b = grads.ln2b + l * C; - float* dl_fcw = grads.fcw + l * 4*C * C; - float* dl_fcb = grads.fcb + l * 4*C; - float* dl_fcprojw = grads.fcprojw + l * C * 4*C; - float* dl_fcprojb = grads.fcprojb + l * C; - // get the pointers of the activations for this layer - float* l_ln1 = acts.ln1 + l * B * T * C; - float* l_ln1_mean = acts.ln1_mean + l * B * T; - float* l_ln1_rstd = acts.ln1_rstd + l * B * T; - float* l_qkv = acts.qkv + l * B * T * 3*C; - float* l_atty = acts.atty + l * B * T * C; - float* l_att = acts.att + l * B * NH * T * T; - float* l_residual2 = acts.residual2 + l * B * T * C; - float* l_ln2 = acts.ln2 + l * B * T * C; - float* l_ln2_mean = acts.ln2_mean + l * B * T; - float* l_ln2_rstd = acts.ln2_rstd + l * B * T; - float* l_fch = acts.fch + l * B * T * 4*C; - float* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C; - // get the pointers of the gradients of the activations for this layer - float* dl_ln1 = grads_acts.ln1 + l * B * T * C; - float* dl_qkv = grads_acts.qkv + l * B * T * 3*C; - float* dl_atty = grads_acts.atty + l * B * T * C; - float* dl_preatt = grads_acts.preatt + l * B * NH * T * T; - float* dl_att = grads_acts.att + l * B * NH * T * T; - float* dl_attproj = grads_acts.attproj + l * B * T * C; - float* dl_residual2 = grads_acts.residual2 + l * B * T * C; - float* dl_ln2 = grads_acts.ln2 + l * B * T * C; - float* dl_fch = grads_acts.fch + l * B * T * 4*C; - float* dl_fch_gelu = grads_acts.fch_gelu + l * B * T * 4*C; - float* dl_fcproj = grads_acts.fcproj + l * B * T * C; - float* dl_residual3 = grads_acts.residual3 + l * B * T * C; - - // backprop this layer - printf(" [Backward] : Residual2\n"); - residual_backward(ctx, dl_residual2, dl_fcproj, dl_residual3, B*T*C); - printf(" [Backward] : FF Down \n"); - matmul_backward(ctx, dl_fch_gelu, dl_fcprojw, dl_fcprojb, dl_fcproj, l_fch_gelu, l_fcprojw, B, T, 4*C, C); - printf(" [Backward] : GELU\n"); - gelu_backward(ctx, dl_fch, l_fch, dl_fch_gelu, B*T*4*C); - printf(" [Backward] : FF Up\n"); - matmul_backward(ctx, dl_ln2, dl_fcw, dl_fcb, dl_fch, l_ln2, l_fcw, B, T, C, 4*C); - printf(" [Backward] : LayerNorm2\n"); - layernorm_backward(ctx, dl_residual2, dl_ln2w, dl_ln2b, dl_ln2, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C); - printf(" [Backward] : Residual1\n"); - residual_backward(ctx, dresidual, dl_attproj, dl_residual2, B*T*C); - printf(" [Backward] : Attention Projection\n"); - matmul_backward(ctx, dl_atty, dl_attprojw, dl_attprojb, dl_attproj, l_atty, l_attprojw, B, T, C, C); - printf(" [Backward] : Attention\n"); - attention_backward(ctx, dl_qkv, dl_preatt, dl_att, dl_atty, l_qkv, l_att, B, T, C, NH); - printf(" [Backward] : QKV Projection\n"); - matmul_backward(ctx, dl_ln1, dl_qkvw, dl_qkvb, dl_qkv, l_ln1, l_qkvw, B, T, C, 3*C); - printf(" [Backward] : LayerNorm1\n"); - layernorm_backward(ctx, dresidual, dl_ln1w, dl_ln1b, dl_ln1, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); - } - encoder_backward(ctx, grads.wte, grads.wpe, grads_acts.encoded, model->inputs, B, T, C); -} - -void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) { - // reference: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html - - // lazily allocate the memory for m_memory and v_memory - if (model->m_memory == NULL) { - model->m_memory = (float*)calloc(model->num_parameters, sizeof(float)); - model->v_memory = (float*)calloc(model->num_parameters, sizeof(float)); - } - - for (size_t i = 0; i < model->num_parameters; i++) { - float param = model->params_memory[i]; - float grad = model->grads_memory[i]; - - // update the first moment (momentum) - float m = beta1 * model->m_memory[i] + (1.0f - beta1) * grad; - // update the second moment (RMSprop) - float v = beta2 * model->v_memory[i] + (1.0f - beta2) * grad * grad; - // bias-correct both moments - float m_hat = m / (1.0f - powf(beta1, t)); - float v_hat = v / (1.0f - powf(beta2, t)); - - // update - model->m_memory[i] = m; - model->v_memory[i] = v; - model->params_memory[i] -= learning_rate * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * param); - } -} - -void gpt2_free(GPT2 *model) { - free(model->params_memory); - free(model->grads_memory); - free(model->m_memory); - free(model->v_memory); - free(model->acts_memory); - free(model->grads_acts_memory); - free(model->inputs); - free(model->targets); -} - -#ifndef TESTING -// if we are TESTING (see test_gpt2.c), we'll skip the int main below -// ---------------------------------------------------------------------------- -// sampler - -unsigned int random_u32(uint64_t *state) { - // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A - *state ^= *state >> 12; - *state ^= *state << 25; - *state ^= *state >> 27; - return (*state * 0x2545F4914F6CDD1Dull) >> 32; -} -float random_f32(uint64_t *state) { // random float32 in [0,1) - return (random_u32(state) >> 8) / 16777216.0f; -} - -int sample_mult(float* probabilities, int n, float coin) { - // sample index from probabilities (they must sum to 1!) - // coin is a random number in [0, 1), usually from random_f32() - float cdf = 0.0f; - for (int i = 0; i < n; i++) { - cdf += probabilities[i]; - if (coin < cdf) { - return i; - } - } - return n - 1; // in case of rounding errors -} - -// ---------------------------------------------------------------------------- -// main training loop -int main() { - - setLogLevel(kWarn); - - printf("Creating GPU context\n"); - WGPURequiredLimits requiredLimits = LIMITS_BUFFER_SIZE_1GB; - gpu::Context ctx = gpu::createContext({}, {}, { - .requiredLimits = &requiredLimits - }); - // gpu::Context ctx = gpu::createContext(); - - // build the GPT-2 model from a checkpoint - GPT2 model; - gpt2_build_from_checkpoint(ctx, &model, "gpt2_124M.bin"); - - // build the DataLoaders from tokens files. for now use tiny_shakespeare if available, else tiny_stories - const char* tiny_stories_train = "dev/data/tinystories/TinyStories_train.bin"; - const char* tiny_stories_val = "dev/data/tinystories/TinyStories_val.bin"; - const char* tiny_shakespeare_train = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; - const char* tiny_shakespeare_val = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; - const char* train_tokens = access(tiny_shakespeare_train, F_OK) != -1 ? tiny_shakespeare_train : tiny_stories_train; - const char* val_tokens = access(tiny_shakespeare_val, F_OK) != -1 ? tiny_shakespeare_val : tiny_stories_val; - constexpr int B = 4; // batch size 4 (i.e. 4 independent token sequences will be trained on) - constexpr int T = 64; // sequence length 64 (i.e. each sequence is 64 tokens long). must be <= maxT, which is 1024 for GPT-2 - DataLoader train_loader, val_loader; - dataloader_init(&train_loader, train_tokens, B, T, 0, 1, 1); - dataloader_init(&val_loader, val_tokens, B, T, 0, 1, 0); - printf("train dataset num_batches: %zu\n", train_loader.num_tokens / (B*T)); - printf("val dataset num_batches: %zu\n", val_loader.num_tokens / (B*T)); - int val_num_batches = 5; - - // build the Tokenizer - Tokenizer tokenizer; - tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); - - // some memory for generating samples from the model - uint64_t rng_state = 1337; - int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); - const int genT = 64; // number of steps of inference we will do - - - // train - struct timespec start, end; - printf("Starting training\n"); - for (int step = 0; step <= 40; step++) { - printf("Step %d\n", step); - - // once in a while estimate the validation loss - if (step % 10 == 0) { - float val_loss = 0.0f; - dataloader_reset(&val_loader); - for (int i = 0; i < val_num_batches; i++) { - dataloader_next_batch(&val_loader); - gpt2_forward(ctx, &model, val_loader.inputs, val_loader.targets, B, T); - val_loss += model.mean_loss; - } - val_loss /= val_num_batches; - printf("val loss %f\n", val_loss); - } - - // once in a while do model inference to print generated text - if (step > 0 && step % 20 == 0) { - // fill up gen_tokens with the GPT2_EOT, which kicks off the generation - for(int i = 0; i < B * T; ++i) { - gen_tokens[i] = tokenizer.eot_token; - } - // now sample from the model autoregressively - printf("generating:\n---\n"); - for (int t = 1; t < genT; t++) { - // note that inference is very wasteful here because for each token - // we re-calculate the forward pass for all of (B,T) positions from scratch - // but the inference here is just for sanity checking anyway - // and we can maybe optimize a bit more later, with careful tests - gpt2_forward(ctx, &model, gen_tokens, NULL, B, T); - // furthermore, below we're only using b=0 (i.e. the first row) of all B rows - // we're in principle running B "inference streams" in parallel here - // but only using position 0 - // get the Vp-dimensional vector probs[0, t-1, :] - float* probs = model.acts.probs + (t-1) * model.config.padded_vocab_size; - float coin = random_f32(&rng_state); - // note we're only sampling from the first V elements, ignoring padding - // (the probabilities in the padded region should be zero anyway) - int next_token = sample_mult(probs, model.config.vocab_size, coin); - gen_tokens[t] = next_token; - // print the generated token, either using the Tokenizer or a fallback - if (tokenizer.init_ok) { - const char* token_str = tokenizer_decode(&tokenizer, next_token); - safe_printf(token_str); - } else { - // fall back to printing the token id - printf("%d ", next_token); - } - fflush(stdout); - } - printf("\n---\n"); - } - - // do a training step - clock_gettime(CLOCK_MONOTONIC, &start); - dataloader_next_batch(&train_loader); - gpt2_forward(ctx, &model, train_loader.inputs, train_loader.targets, B, T); - gpt2_zero_grad(&model); - gpt2_backward(ctx, &model); - gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, step+1); - clock_gettime(CLOCK_MONOTONIC, &end); - double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9; - printf("step %d: train loss %f (took %f ms)\n", step, model.mean_loss, time_elapsed_s * 1000); - } - - // free - dataloader_free(&train_loader); - dataloader_free(&val_loader); - tokenizer_free(&tokenizer); - gpt2_free(&model); - free(gen_tokens); - return 0; -} -#endif diff --git a/experimental/kernels/gpt2_webgpu_aot.cpp b/experimental/kernels/gpt2_webgpu_aot.cpp deleted file mode 100644 index 966fb7a..0000000 --- a/experimental/kernels/gpt2_webgpu_aot.cpp +++ /dev/null @@ -1,1109 +0,0 @@ -#include "gpu.hpp" -#include "ops_aot.hpp" -/* -This file trains the GPT-2 model. -This version is the clean, minimal, reference. As such: -- it runs on CPU. -- it does not make the code too complex; it is readable. -- it does not use any processor-specific instructions, intrinsics and such. -- it _does_ use a few OpenMP pragmas because this is a large speedup at very low cost -There will be other versions of this code that specialize it and make it fast. -*/ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#ifdef OMP -#include -#endif -// our own utilities -// defines: fopenCheck, freadCheck, fcloseCheck, fseekCheck, mallocCheck -#include "llmc/utils.h" -// defines: tokenizer_init, tokenizer_decode, tokenizer_free -#include "llmc/tokenizer.h" -// defines: dataloader_init, dataloader_reset, dataloader_next_batch, dataloader_free -#include "llmc/dataloader.h" - -using namespace gpu; - -// ---------------------------------------------------------------------------- -// GPT-2 model definition - -typedef struct { - int max_seq_len; // max sequence length, e.g. 1024 - int vocab_size; // vocab size, e.g. 50257 - int padded_vocab_size; // padded to e.g. %128==0, 50304 - int num_layers; // number of layers, e.g. 12 - int num_heads; // number of heads in attention, e.g. 12 - int channels; // number of channels, e.g. 768 -} GPT2Config; - -// the parameters of the model -#define NUM_PARAMETER_TENSORS 16 -typedef struct { - Tensor wte; // (V, C) - Tensor wpe; // (maxT, C) - std::vector ln1w; // (L, C) - std::vector ln1b; // (L, C) - std::vector qkvw; // (L, 3*C, C) - std::vector qkvb; // (L, 3*C) - std::vector attprojw; // (L, C, C) - std::vector attprojb; // (L, C) - std::vector ln2w; // (L, C) - std::vector ln2b; // (L, C) - std::vector fcw; // (L, 4*C, C) - std::vector fcb; // (L, 4*C) - std::vector fcprojw; // (L, C, 4*C) - std::vector fcprojb; // (L, C) - Tensor lnfw; // (C) - Tensor lnfb; // (C) -} ParameterTensors; - -void fill_in_parameter_sizes(size_t* param_sizes, GPT2Config config) { - size_t Vp = config.padded_vocab_size; - size_t C = config.channels; - size_t maxT = config.max_seq_len; - size_t L = config.num_layers; - param_sizes[0] = Vp * C; // wte - param_sizes[1] = maxT * C; // wpe - param_sizes[2] = L * C; // ln1w - param_sizes[3] = L * C; // ln1b - param_sizes[4] = L * (3 * C) * C; // qkvw - param_sizes[5] = L * (3 * C); // qkvb - param_sizes[6] = L * C * C; // attprojw - param_sizes[7] = L * C; // attprojb - param_sizes[8] = L * C; // ln2w - param_sizes[9] = L * C; // ln2b - param_sizes[10] = L * (4 * C) * C; // fcw - param_sizes[11] = L * (4 * C); // fcb - param_sizes[12] = L * C * (4 * C); // fcprojw - param_sizes[13] = L * C; // fcprojb - param_sizes[14] = C; // lnfw - param_sizes[15] = C; // lnfb -} - -// allocate memory for the parameters and point the individual tensors to the right places -void malloc_and_point_parameters(Context& ctx, GPT2Config config, ParameterTensors* params, size_t* param_sizes) { - size_t L = config.num_layers; - params->wte = createTensor(ctx, Shape{param_sizes[0]}, kf32); - params->wpe = createTensor(ctx, Shape{param_sizes[1]}, kf32); - - params->ln1w.resize(L); - params->ln1b.resize(L); - params->qkvw.resize(L); - params->qkvb.resize(L); - params->attprojw.resize(L); - params->attprojb.resize(L); - params->ln2w.resize(L); - params->ln2b.resize(L); - params->fcw.resize(L); - params->fcb.resize(L); - params->fcprojw.resize(L); - params->fcprojb.resize(L); - for(int l = 0; l < L ; l++) { - params->ln1w[l] = createTensor(ctx, Shape{param_sizes[2]/config.num_layers}, kf32); - params->ln1b[l] = createTensor(ctx, Shape{param_sizes[3]/config.num_layers}, kf32); - params->qkvw[l] = createTensor(ctx, Shape{param_sizes[4]/config.num_layers}, kf32); - params->qkvb[l] = createTensor(ctx, Shape{param_sizes[5]/config.num_layers}, kf32); - params->attprojw[l] = createTensor(ctx, Shape{param_sizes[6]/config.num_layers}, kf32); - params->attprojb[l] = createTensor(ctx, Shape{param_sizes[7]/config.num_layers}, kf32); - params->ln2w[l] = createTensor(ctx, Shape{param_sizes[8]/config.num_layers}, kf32); - params->ln2b[l] = createTensor(ctx, Shape{param_sizes[9]/config.num_layers}, kf32); - params->fcw[l] = createTensor(ctx, Shape{param_sizes[10]/config.num_layers}, kf32); - params->fcb[l] = createTensor(ctx, Shape{param_sizes[11]/config.num_layers}, kf32); - params->fcprojw[l] = createTensor(ctx, Shape{param_sizes[12]/config.num_layers}, kf32); - params->fcprojb[l] = createTensor(ctx, Shape{param_sizes[13]/config.num_layers}, kf32); - } - params->lnfw = createTensor(ctx, Shape{param_sizes[14]}, kf32); - params->lnfb = createTensor(ctx, Shape{param_sizes[15]}, kf32); -} - - -#define NUM_ACTIVATION_TENSORS 23 -typedef struct { - Tensor encoded; // (B, T, C) - std::vector ln1; // (L, B, T, C) - std::vector ln1_mean; // (L, B, T) - std::vector ln1_rstd; // (L, B, T) - std::vector qkv; // (L, B, T, 3*C) - std::vector atty; // (L, B, T, C) - std::vector preatt; // (L, B, NH, T, T) - std::vector att; // (L, B, NH, T, T) - std::vector attproj; // (L, B, T, C) - std::vector residual2; // (L, B, T, C) - std::vector ln2; // (L, B, T, C) - std::vector ln2_mean; // (L, B, T) - std::vector ln2_rstd; // (L, B, T) - std::vector fch; // (L, B, T, 4*C) - std::vector fch_gelu; // (L, B, T, 4*C) - std::vector fcproj; // (L, B, T, C) - std::vector residual3; // (L, B, T, C) - Tensor lnf; // (B, T, C) - Tensor lnf_mean; // (B, T) - Tensor lnf_rstd; // (B, T) - Tensor logits; // (B, T, V) - Tensor probs; // (B, T, V) - Tensor losses; // (B, T) -} ActivationTensors; - -typedef struct { - Kernel encoder_forward; - std::vector layernorm_forward; - std::vector qkv_projection_forward; - std::vector attention_forward; - std::vector attention_projection_forward; - std::vector residual_forward; - std::vector ff_up_forward; - std::vector gelu_forward; - std::vector ff_down_forward; - std::vector residual2_forward; - Kernel layernorm_final_forward; - Kernel matmul_final_forward; - Kernel softmax_final_forward; - Kernel crossentropy_forward; - - Kernel crossentropy_softmax_backward; - Kernel matmul_final_backward; - Kernel layernorm_final_backward; - std::vector residual2_backward; - std::vector ff_down_backward; - std::vector gelu_backward; - std::vector ff_up_backward; - std::vector layernorm2_backward; - std::vector attention_projection_backward; - std::vector attention_backward; - std::vector qkv_projection_backward; - std::vector layernorm1_backward; - Kernel encoder_backward; -} Kernels; - -void fill_in_activation_sizes(size_t* act_sizes, GPT2Config config, int B, int T) { - size_t C = config.channels; - size_t NH = config.num_heads; - size_t L = config.num_layers; - size_t Vp = config.padded_vocab_size; - act_sizes[0] = B * T * C; // encoded - act_sizes[1] = L * B * T * C; // ln1 - act_sizes[2] = L * B * T; // ln1_mean - act_sizes[3] = L * B * T; // ln1_rstd - act_sizes[4] = L * B * T * 3 * C; // qkv - act_sizes[5] = L * B * T * C; // atty - act_sizes[6] = L * B * NH * T * T; // preatt - act_sizes[7] = L * B * NH * T * T; // att - act_sizes[8] = L * B * T * C; // attproj - act_sizes[9] = L * B * T * C; // residual2 - act_sizes[10] = L * B * T * C; // ln2 - act_sizes[11] = L * B * T; // ln2_mean - act_sizes[12] = L * B * T; // ln2_rstd - act_sizes[13] = L * B * T * 4 * C; // fch - act_sizes[14] = L * B * T * 4 * C; // fch_gelu - act_sizes[15] = L * B * T * C; // fcproj - act_sizes[16] = L * B * T * C; // residual3 - act_sizes[17] = B * T * C; // lnf - act_sizes[18] = B * T; // lnf_mean - act_sizes[19] = B * T; // lnf_rstd - act_sizes[20] = B * T * Vp; // logits - act_sizes[21] = B * T * Vp; // probs - act_sizes[22] = B * T; // losses -} - -void malloc_and_point_activations(Context& ctx, GPT2Config config, ActivationTensors* acts, size_t* act_sizes) { - size_t L = config.num_layers; - acts->encoded = createTensor(ctx, Shape{act_sizes[0]}, kf32); - acts->ln1.resize(L); - acts->ln1_mean.resize(L); - acts->ln1_rstd.resize(L); - acts->qkv.resize(L); - acts->atty.resize(L); - acts->preatt.resize(L); - acts->att.resize(L); - acts->attproj.resize(L); - acts->residual2.resize(L); - acts->ln2.resize(L); - acts->ln2_mean.resize(L); - acts->ln2_rstd.resize(L); - acts->fch.resize(L); - acts->fch_gelu.resize(L); - acts->fcproj.resize(L); - acts->residual3.resize(L); - for (int l = 0; l < L; l++) { - acts->ln1[l] = createTensor(ctx, Shape{act_sizes[1]/config.num_layers}, kf32); - acts->ln1_mean[l] = createTensor(ctx, Shape{act_sizes[2]/config.num_layers}, kf32); - acts->ln1_rstd[l] = createTensor(ctx, Shape{act_sizes[3]/config.num_layers}, kf32); - acts->qkv[l] = createTensor(ctx, Shape{act_sizes[4]/config.num_layers}, kf32); - acts->atty[l] = createTensor(ctx, Shape{act_sizes[5]/config.num_layers}, kf32); - acts->preatt[l] = createTensor(ctx, Shape{act_sizes[6]/config.num_layers}, kf32); - acts->att[l] = createTensor(ctx, Shape{act_sizes[7]/config.num_layers}, kf32); - acts->attproj[l] = createTensor(ctx, Shape{act_sizes[8]/config.num_layers}, kf32); - acts->residual2[l] = createTensor(ctx, Shape{act_sizes[9]/config.num_layers}, kf32); - acts->ln2[l] = createTensor(ctx, Shape{act_sizes[10]/config.num_layers}, kf32); - acts->ln2_mean[l] = createTensor(ctx, Shape{act_sizes[11]/config.num_layers}, kf32); - acts->ln2_rstd[l] = createTensor(ctx, Shape{act_sizes[12]/config.num_layers}, kf32); - acts->fch[l] = createTensor(ctx, Shape{act_sizes[13]/config.num_layers}, kf32); - acts->fch_gelu[l] = createTensor(ctx, Shape{act_sizes[14]/config.num_layers}, kf32); - acts->fcproj[l] = createTensor(ctx, Shape{act_sizes[15]/config.num_layers}, kf32); - acts->residual3[l] = createTensor(ctx, Shape{act_sizes[16]/config.num_layers}, kf32); - } - acts->lnf = createTensor(ctx, Shape{act_sizes[17]}, kf32); - acts->lnf_mean = createTensor(ctx, Shape{act_sizes[18]}, kf32); - acts->lnf_rstd = createTensor(ctx, Shape{act_sizes[19]}, kf32); - acts->logits = createTensor(ctx, Shape{act_sizes[20]}, kf32); - acts->probs = createTensor(ctx, Shape{act_sizes[21]}, kf32); - acts->losses = createTensor(ctx, Shape{act_sizes[22]}, kf32); -} - -void gpu_alloc(Context& ctx, Tensor* tensors, size_t* sizes, size_t n) { - for (size_t i = 0; i < n; i++) { - tensors[i] = createTensor(ctx, Shape{sizes[i]}, kf32); - } -} - -typedef struct { - GPT2Config config; - // the weights (parameters) of the model, and their sizes - ParameterTensors params; - size_t param_sizes[NUM_PARAMETER_TENSORS]; - float* params_memory; - size_t num_parameters; - // gradients of the weights - ParameterTensors grads; - float* grads_memory; - // buffers for the AdamW optimizer - float* m_memory; - float* v_memory; - // the activations of the model, and their sizes - ActivationTensors acts; - size_t act_sizes[NUM_ACTIVATION_TENSORS]; - float* acts_memory; - size_t num_activations; - // gradients of the activations - ActivationTensors grads_acts; - float* grads_acts_memory; - // other run state configuration - int batch_size; // the batch size (B) of current forward pass - int seq_len; // the sequence length (T) of current forward pass - Tensor inputs; // the input tokens for the current forward pass - Tensor targets; // the target tokens for the current forward pass - float mean_loss; // after a forward pass with targets, will be populated with the mean loss - float* mean_loss_buffer; - float* probs_buffer; - - Tensor nullTensor; - - // kernels - Kernels kernels; - bool backward_enabled; -} GPT2; - -void gpt2_build_from_checkpoint(Context& ctx, GPT2 *model, const char* checkpoint_path) { - printf("Building GPT-2 model from checkpoint '%s'\n", checkpoint_path); - // read in model from a checkpoint file - FILE *model_file = fopenCheck(checkpoint_path, "rb"); - int model_header[256]; - freadCheck(model_header, sizeof(int), 256, model_file); - if (model_header[0] != 20240326) { printf("Bad magic model file\n"); exit(1); } - if (model_header[1] != 3) { - printf("Bad version in model file\n"); - printf("---> HINT: try to re-run `python train_gpt2.py`\n"); - exit(1); - } - - // read in hyperparameters - size_t maxT, V, Vp, L, NH, C; // size_t to prevent int overflow - model->config.max_seq_len = maxT = model_header[2]; - model->config.vocab_size = V = model_header[3]; -#ifdef __EMSCRIPTEN__ - model->config.num_layers = L = 12; // TODO(avh): Debugging only hack - revert this -#else - model->config.num_layers = L = model_header[4]; -#endif - model->config.num_heads = NH = model_header[5]; - model->config.channels = C = model_header[6]; - model->config.padded_vocab_size = Vp = model_header[7]; - printf("[GPT-2]\n"); - printf("max_seq_len: %zu\n", maxT); - printf("vocab_size: %zu\n", V); - printf("padded_vocab_size: %zu\n", Vp); - printf("num_layers: %zu\n", L); - printf("num_heads: %zu\n", NH); - printf("channels: %zu\n", C); - - // allocate space for all the parameters and read them in - fill_in_parameter_sizes(model->param_sizes, model->config); - // count the number of parameters - size_t num_parameters = 0; - for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { - num_parameters += model->param_sizes[i]; - } - printf("num_parameters: %zu\n", num_parameters); - model->num_parameters = num_parameters; - - // read in all the parameters from file - malloc_and_point_parameters(ctx, model->config, &model->params, model->param_sizes); - model->params_memory = (float*)mallocCheck(num_parameters * sizeof(float)); - freadCheck(model->params_memory, sizeof(float), num_parameters, model_file); - fcloseCheck(model_file); - - // transfer to GPU memory - float* iter = model->params_memory; - toGPU(ctx, iter, model->params.wte); - iter += model->param_sizes[0]; - toGPU(ctx, iter, model->params.wpe); - iter += model->param_sizes[1]; - for (int l = 0; l < L; l++) { - toGPU(ctx, iter, model->params.ln1w[l]); - iter += model->param_sizes[2]/L; - toGPU(ctx, iter, model->params.ln1b[l]); - iter += model->param_sizes[3]/L; - toGPU(ctx, iter, model->params.qkvw[l]); - iter += model->param_sizes[4]/L; - toGPU(ctx, iter, model->params.qkvb[l]); - iter += model->param_sizes[5]/L; - toGPU(ctx, iter, model->params.attprojw[l]); - iter += model->param_sizes[6]/L; - toGPU(ctx, iter, model->params.attprojb[l]); - iter += model->param_sizes[7]/L; - toGPU(ctx, iter, model->params.ln2w[l]); - iter += model->param_sizes[8]/L; - toGPU(ctx, iter, model->params.ln2b[l]); - iter += model->param_sizes[9]/L; - toGPU(ctx, iter, model->params.fcw[l]); - iter += model->param_sizes[10]/L; - toGPU(ctx, iter, model->params.fcb[l]); - iter += model->param_sizes[11]/L; - toGPU(ctx, iter, model->params.fcprojw[l]); - iter += model->param_sizes[12]/L; - toGPU(ctx, iter, model->params.fcprojb[l]); - iter += model->param_sizes[13]/L; - } - toGPU(ctx, iter, model->params.lnfw); - iter += model->param_sizes[14]; - toGPU(ctx, iter, model->params.lnfb); - iter += model->param_sizes[15]; - - - // other inits - model->acts_memory = NULL; - model->grads_memory = NULL; - model->m_memory = NULL; - model->v_memory = NULL; - model->grads_acts_memory = NULL; - model->batch_size = 0; - model->seq_len = 0; - model->mean_loss = -1.0f; // -1.0f will designate no loss - model->mean_loss_buffer = NULL; - model->probs_buffer = NULL; - model->backward_enabled = false; - - printf("Model build complete\n"); - -} - - -void gpt2_forward(Context& ctx, GPT2 *model, Tensor& inputs, Tensor& targets, size_t B, size_t T) { - // targets are optional and could be NULL - - // ensure the model was initialized or error out - if (model->params_memory == NULL) { - printf("Error: model was not initialized properly.\n"); - exit(1); - } - - // convenience parameters (size_t to help prevent int overflow) - size_t V = model->config.vocab_size; - size_t Vp = model->config.padded_vocab_size; - size_t L = model->config.num_layers; - size_t NH = model->config.num_heads; - size_t C = model->config.channels; - - // // validate inputs, all indices must be in the range [0, V) - // for(int i = 0; i < B * T; i++) { - // assert(0 <= inputs[i] && inputs[i] < V); - // if (targets != NULL) { - // assert(0 <= targets[i] && targets[i] < V); - // } - // } - - // allocate space for all the activations if needed (done here, lazily) - if(model->acts_memory == NULL) { - // record the current B,T as well - model->batch_size = B; - model->seq_len = T; - // and now allocate the space - fill_in_activation_sizes(model->act_sizes, model->config, B, T); - model->mean_loss_buffer = (float*)mallocCheck(sizeof(float) * model->batch_size * model->seq_len); - model->probs_buffer = (float*)mallocCheck(sizeof(float) * model->batch_size * model->seq_len * Vp); - - // TODO(avh): this is just a resource test for now, eventually deprecate CPU allocations - size_t num_activations = 0; - for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) { - num_activations += model->act_sizes[i]; - } - printf("num_activations: %zu\n", num_activations); - model->num_activations = num_activations; - printf("Allocating %.2f MB for activations\n", num_activations * sizeof(float) / (1024.0f * 1024.0f)); - malloc_and_point_activations(ctx, model->config, &model->acts, model->act_sizes); - // also create memory for caching inputs and targets - //model->inputs = (int*)mallocCheck(B * T * sizeof(int)); - //model->targets = (int*)mallocCheck(B * T * sizeof(int)); // might be unused if we never have targets but it's small - model->inputs = createTensor(ctx, Shape{B * T}, ki32); - model->targets = createTensor(ctx, Shape{B * T}, ki32); - } else { - // validate B,T is consistent with how we've allocated the memory before - // in principle we could get more clever here in the future, for now this is safest - if (B != model->batch_size || T != model->seq_len) { - printf("Model: B=%d T=%d, Desired: B=%d T=%d\n", model->batch_size, model->seq_len, (int)B, (int)T); - exit(EXIT_FAILURE); - } - } - // create all kernels ahead of time - if (model->kernels.encoder_forward == nullptr) { - printf("Creating Kernels\n"); - Kernels& kernels = model->kernels; - kernels.layernorm_forward.resize(L); - kernels.layernorm1_backward.resize(L); - kernels.qkv_projection_forward.resize(L); - kernels.qkv_projection_backward.resize(L); - kernels.attention_forward.resize(L); - kernels.attention_backward.resize(L); - kernels.attention_projection_forward.resize(L); - kernels.attention_projection_backward.resize(L); - kernels.residual_forward.resize(L); - kernels.residual2_forward.resize(L); - kernels.residual2_backward.resize(L); - kernels.ff_up_forward.resize(L); - kernels.ff_up_backward.resize(L); - kernels.gelu_forward.resize(L); - kernels.gelu_backward.resize(L); - kernels.ff_down_forward.resize(L); - kernels.ff_down_backward.resize(L); - for (int l = 0; l < L; ++l) { - kernels.layernorm_forward[l] = layernorm_forward(ctx, model->acts.ln1[l], model->acts.ln1_mean[l], model->acts.ln1_rstd[l], - /*input=*/ model->acts.residual3[l], /*weight=*/ model->params.ln1w[l], /*bias=*/ model->params.ln1b[l], - B, T, C); - kernels.qkv_projection_forward[l] = matmul_forward(ctx, model->acts.qkv[l], model->acts.ln1[l], model->params.qkvw[l], model->params.qkvb[l], B, T, C, 3*C); - kernels.attention_forward[l] = attention_forward(ctx, model->acts.atty[l], model->acts.preatt[l], model->acts.att[l], model->acts.qkv[l], B, T, C, NH); - kernels.attention_projection_forward[l] = matmul_forward(ctx, model->acts.attproj[l], model->acts.atty[l], model->params.attprojw[l], model->params.attprojb[l], B, T, C, C); - kernels.residual_forward[l] = residual_forward(ctx, model->acts.residual2[l], model->acts.residual3[l], model->acts.attproj[l], B*T*C); - kernels.ff_up_forward[l] = matmul_forward(ctx, model->acts.fch[l], model->acts.ln2[l], model->params.fcw[l], model->params.fcb[l], B, T, C, 4*C); - kernels.gelu_forward[l] = gelu_forward(ctx, model->acts.fch_gelu[l], model->acts.fch[l], B*T*4*C); - kernels.ff_down_forward[l] = matmul_forward(ctx, model->acts.fcproj[l], model->acts.fch_gelu[l], model->params.fcw[l], model->params.fcb[l], B, T, 4*C, C); - kernels.residual2_forward[l] = residual_forward(ctx, model->acts.residual3[l], model->acts.residual2[l], model->acts.fcproj[l], B*T*C); - } - kernels.crossentropy_forward = crossentropy_forward(ctx, model->acts.losses, model->acts.probs, targets, B, T, Vp); - - kernels.encoder_forward = encoder_forward(ctx, model->acts.encoded, inputs, model->params.wte, model->params.wpe, B, T, C); // encoding goes into residual[0] - if(model->backward_enabled) - kernels.encoder_backward = encoder_backward(ctx, model->params.wte, model->params.wpe, model->acts.encoded, inputs, B, T, C); - kernels.layernorm_final_forward = layernorm_forward(ctx, model->acts.lnf, model->acts.lnf_mean, model->acts.lnf_rstd, - /*input=*/ model->acts.residual3[L-1], /*weight=*/ model->params.lnfw, /*bias=*/ model->params.lnfb, - B, T, C); - Tensor nullTensor = createTensor(ctx, Shape{1}, kf32); - model->nullTensor = nullTensor; - kernels.matmul_final_forward = matmul_forward(ctx, model->acts.logits, model->acts.lnf, model->params.wte, nullTensor, B, T, C, Vp); - kernels.softmax_final_forward = softmax_forward(ctx, model->acts.probs, model->acts.logits, B, T, V, Vp); - if(model->backward_enabled) - kernels.crossentropy_softmax_backward = crossentropy_softmax_backward(ctx, model->acts.logits, model->acts.losses, model->acts.probs, targets, B, T, V, Vp); - if(model->backward_enabled) - kernels.matmul_final_backward = matmul_backward(ctx, model->acts.lnf, model->params.wte, nullTensor, model->acts.logits, - model->acts.lnf, model->params.wte, B, T, C, Vp); - if(model->backward_enabled) - kernels.layernorm_final_backward = layernorm_backward(ctx, model->acts.residual3[L-1], model->params.lnfw, model->params.lnfb, - model->acts.lnf, model->acts.residual3[L-1], model->params.lnfw, - model->acts.lnf_mean, model->acts.lnf_rstd, B, T, C); - printf("Created Kernels\n"); - } - - printf("Cache inputs/targets\n"); - printf("Forward pass\n"); - // forward pass - ParameterTensors params = model->params; // for brevity - ActivationTensors acts = model->acts; - float* residual; - printf("Encoding\n"); - //printf("inputs[0] = %d\n", inputs[0]); - // encoder_forward(ctx, acts.encoded, inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0] - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.encoder_forward, promise); - wait(ctx, future); - } - for (int l = 0; l < L; l++) { - printf("Forward Pass Layer %d\n", l); - - // now do the forward pass - printf(" [Forward] : LayerNorm1\n"); - // layernorm_forward(ctx, l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.layernorm_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : QKV Projection\n"); - // matmul_forward(ctx, l_qkv, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.qkv_projection_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : Attention\n"); - // attention_forward(ctx, l_atty, l_preatt, l_att, l_qkv, B, T, C, NH); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.attention_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : Attention Projection\n"); - // matmul_forward(ctx, l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.attention_projection_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : Residual1\n"); - // residual_forward(ctx, l_residual2, residual, l_attproj, B*T*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.residual_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : LayerNorm2\n"); - // layernorm_forward(ctx, l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.layernorm_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : FF Up\n"); - // matmul_forward(ctx, l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.ff_up_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : GELU\n"); - // gelu_forward(ctx, l_fch_gelu, l_fch, B*T*4*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.gelu_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : FF Down\n"); - // matmul_forward(ctx, l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.ff_down_forward[l], promise); - wait(ctx, future); - } - printf(" [Forward] : Residual2\n"); - // residual_forward(ctx, l_residual3, l_residual2, l_fcproj, B*T*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.residual2_forward[l], promise); - wait(ctx, future); - } - } - // residual = acts.residual3.data() + (L-1) * B * T * C; // last residual is in residual3 - // layernorm_forward(ctx, acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, params.lnfw, params.lnfb, B, T, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.layernorm_final_forward, promise); - wait(ctx, future); - } - // matmul_forward(ctx, acts.logits, acts.lnf, params.wte, NULL, B, T, C, Vp); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.matmul_final_forward, promise); - wait(ctx, future); - } - // softmax_forward(ctx, acts.probs, acts.logits, B, T, V, Vp); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.softmax_final_forward, promise); - wait(ctx, future); - } - - printf("Crossentropy\n"); - // also forward the cross-entropy loss function if we have the targets - // When targets's shape is (1), it means we don't have targets - if (targets.shape[0] != 1) { - // crossentropy_forward(ctx, model->acts.losses, model->acts.probs, targets, B, T, Vp); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.crossentropy_forward, promise); - wait(ctx, future); - } - // for convenience also evaluate the mean loss - float mean_loss = 0.0f; - toCPU(ctx, model->acts.losses, model->mean_loss_buffer, B*T * sizeof(float)); - for (int i=0; imean_loss_buffer[i]; } - mean_loss /= B*T; - model->mean_loss = mean_loss; - } else { - // if we don't have targets, we don't have a loss - model->mean_loss = -1.0f; - } - printf("Forward pass done\n"); -} - -void gpt2_zero_grad(GPT2 *model) { - if(model->grads_memory != NULL) { memset(model->grads_memory, 0, model->num_parameters * sizeof(float)); } - if(model->grads_acts_memory != NULL) { memset(model->grads_acts_memory, 0, model->num_activations * sizeof(float)); } -} - -void gpt2_backward(Context& ctx, GPT2 *model) { - printf("Backward pass\n"); - - // double check we forwarded previously, with targets - if (model->mean_loss == -1.0f) { - printf("Error: must forward with targets before backward\n"); - exit(1); - } - - // lazily allocate the memory for gradients of the weights and activations, if needed - if (model->grads_memory == NULL) { - printf("Allocating %.2f MB for gradients\n", model->num_parameters * sizeof(float) / (1024.0f * 1024.0f)); - malloc_and_point_parameters(ctx, model->config, &model->grads, model->param_sizes); - malloc_and_point_activations(ctx, model->config, &model->grads_acts, model->act_sizes); - gpt2_zero_grad(model); - } - - // convenience shortcuts (and size_t to help prevent int overflow) - size_t B = model->batch_size; - size_t T = model->seq_len; - size_t V = model->config.vocab_size; - size_t Vp = model->config.padded_vocab_size; - size_t L = model->config.num_layers; - size_t NH = model->config.num_heads; - size_t C = model->config.channels; - - // backward pass: go in the reverse order of the forward pass, and call backward() functions - ParameterTensors params = model->params; // for brevity - ParameterTensors grads = model->grads; - ActivationTensors acts = model->acts; - ActivationTensors grads_acts = model->grads_acts; - - // we kick off the chain rule by filling in dlosses with 1.0f/(B*T) - // technically this is a small, inline backward() pass of calculating - // total, final loss as the mean over all losses over all (B,T) positions in the batch - float dloss_mean = 1.0f / (B*T); - for (int i = 0; i < B*T; i++) { model->mean_loss_buffer[i] = dloss_mean; } - toGPU(ctx, model->mean_loss_buffer, model->acts.losses); - //toGPU(ctx, grads_acts.losses.data, model->acts_.data[22]); - - // crossentropy_softmax_backward(ctx, grads_acts.logits, grads_acts.losses, acts.probs, model->targets, B, T, V, Vp); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.crossentropy_softmax_backward, promise); - wait(ctx, future); - } - // matmul_backward(ctx, grads_acts.lnf, grads.wte, NULL, grads_acts.logits, acts.lnf, params.wte, B, T, C, Vp); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.matmul_final_backward, promise); - wait(ctx, future); - } - // layernorm_backward(ctx, dresidual, grads.lnfw, grads.lnfb, grads_acts.lnf, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.layernorm_final_backward, promise); - wait(ctx, future); - } - - for (int l = L-1; l >= 0; l--) { - printf("Backward Pass Layer %d\n", l); - // backprop this layer - printf(" [Backward] : Residual2\n"); - // residual_backward(ctx, dl_residual2, dl_fcproj, dl_residual3, B*T*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.residual2_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : FF Down \n"); - // matmul_backward(ctx, dl_fch_gelu, dl_fcprojw, dl_fcprojb, dl_fcproj, l_fch_gelu, l_fcprojw, B, T, 4*C, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.ff_down_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : GELU\n"); - // gelu_backward(ctx, dl_fch, l_fch, dl_fch_gelu, B*T*4*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.gelu_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : FF Up\n"); - // matmul_backward(ctx, dl_ln2, dl_fcw, dl_fcb, dl_fch, l_ln2, l_fcw, B, T, C, 4*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.ff_up_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : LayerNorm2\n"); - // layernorm_backward(ctx, dl_residual2, dl_ln2w, dl_ln2b, dl_ln2, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.layernorm2_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : Residual1\n"); - // residual_backward(ctx, dresidual, dl_attproj, dl_residual2, B*T*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.residual_forward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : Attention Projection\n"); - // matmul_backward(ctx, dl_atty, dl_attprojw, dl_attprojb, dl_attproj, l_atty, l_attprojw, B, T, C, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.attention_projection_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : Attention\n"); - // attention_backward(ctx, dl_qkv, dl_preatt, dl_att, dl_atty, l_qkv, l_att, B, T, C, NH); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.attention_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : QKV Projection\n"); - // matmul_backward(ctx, dl_ln1, dl_qkvw, dl_qkvb, dl_qkv, l_ln1, l_qkvw, B, T, C, 3*C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.qkv_projection_backward[l], promise); - wait(ctx, future); - } - printf(" [Backward] : LayerNorm1\n"); - // layernorm_backward(ctx, dresidual, dl_ln1w, dl_ln1b, dl_ln1, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.layernorm1_backward[l], promise); - wait(ctx, future); - } - } - // encoder_backward(ctx, grads.wte, grads.wpe, grads_acts.encoded, model->inputs, B, T, C); - { - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, model->kernels.encoder_backward, promise); - wait(ctx, future); - } - // toCPU(ctx, model->params_.data[0], model->grads.wte.data, model->param_sizes[0] * sizeof(float)); - // toCPU(ctx, model->params_.data[1], model->grads.wpe.data, model->param_sizes[1] * sizeof(float)); -} - -void gpt2_update(Context& ctx, GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) { - // reference: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html - - // lazily allocate the memory for m_memory and v_memory - if (model->m_memory == NULL) { - model->m_memory = (float*)calloc(model->num_parameters, sizeof(float)); - model->v_memory = (float*)calloc(model->num_parameters, sizeof(float)); - } - - // Copy the parameters to the CPU - float* iter = model->params_memory; - toCPU(ctx, model->params.wte, iter, model->param_sizes[0] * sizeof(float)); - iter += model->param_sizes[0]; - toCPU(ctx, model->params.wpe, iter, model->param_sizes[1] * sizeof(float)); - iter += model->param_sizes[1]; - size_t L = model->config.num_layers; - for (int l = 0; l < L; l++) { - toCPU(ctx, model->params.ln1w[l], iter, model->param_sizes[2]/L * sizeof(float)); - iter += model->param_sizes[2]/L; - toCPU(ctx, model->params.ln1b[l], iter, model->param_sizes[3]/L * sizeof(float)); - iter += model->param_sizes[3]/L; - toCPU(ctx, model->params.qkvw[l], iter, model->param_sizes[4]/L * sizeof(float)); - iter += model->param_sizes[4]/L; - toCPU(ctx, model->params.qkvb[l], iter, model->param_sizes[5]/L * sizeof(float)); - iter += model->param_sizes[5]/L; - toCPU(ctx, model->params.attprojw[l], iter, model->param_sizes[6]/L * sizeof(float)); - iter += model->param_sizes[6]/L; - toCPU(ctx, model->params.attprojb[l], iter, model->param_sizes[7]/L * sizeof(float)); - iter += model->param_sizes[7]/L; - toCPU(ctx, model->params.ln2w[l], iter, model->param_sizes[8]/L * sizeof(float)); - iter += model->param_sizes[8]/L; - toCPU(ctx, model->params.ln2b[l], iter, model->param_sizes[9]/L * sizeof(float)); - iter += model->param_sizes[9]/L; - toCPU(ctx, model->params.fcw[l], iter, model->param_sizes[10]/L * sizeof(float)); - iter += model->param_sizes[10]/L; - toCPU(ctx, model->params.fcb[l], iter, model->param_sizes[11]/L * sizeof(float)); - iter += model->param_sizes[11]/L; - toCPU(ctx, model->params.fcprojw[l], iter, model->param_sizes[12]/L * sizeof(float)); - iter += model->param_sizes[12]/L; - toCPU(ctx, model->params.fcprojb[l], iter, model->param_sizes[13]/L * sizeof(float)); - iter += model->param_sizes[13]/L; - } - toCPU(ctx, model->params.lnfw, iter, model->param_sizes[14] * sizeof(float)); - iter += model->param_sizes[14]; - toCPU(ctx, model->params.lnfb, iter, model->param_sizes[15] * sizeof(float)); - iter += model->param_sizes[15]; - - - for (size_t i = 0; i < model->num_parameters; i++) { - float param = model->params_memory[i]; - float grad = model->grads_memory[i]; - - // update the first moment (momentum) - float m = beta1 * model->m_memory[i] + (1.0f - beta1) * grad; - // update the second moment (RMSprop) - float v = beta2 * model->v_memory[i] + (1.0f - beta2) * grad * grad; - // bias-correct both moments - float m_hat = m / (1.0f - powf(beta1, t)); - float v_hat = v / (1.0f - powf(beta2, t)); - - // update - model->m_memory[i] = m; - model->v_memory[i] = v; - model->params_memory[i] -= learning_rate * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * param); - } - // toGPU(ctx, model->params_memory, model->params_.data[0]); - // toGPU(ctx, model->params_memory + model->param_sizes[0], model->params_.data[1]); - iter = model->params_memory; - toGPU(ctx, iter, model->params.wte); - iter += model->param_sizes[0]; - toGPU(ctx, iter, model->params.wpe); - iter += model->param_sizes[1]; - for (int l = 0; l < L; l++) { - toGPU(ctx, iter, model->params.ln1w[l]); - iter += model->param_sizes[2]/L; - toGPU(ctx, iter, model->params.ln1b[l]); - iter += model->param_sizes[3]/L; - toGPU(ctx, iter, model->params.qkvw[l]); - iter += model->param_sizes[4]/L; - toGPU(ctx, iter, model->params.qkvb[l]); - iter += model->param_sizes[5]/L; - toGPU(ctx, iter, model->params.attprojw[l]); - iter += model->param_sizes[6]/L; - toGPU(ctx, iter, model->params.attprojb[l]); - iter += model->param_sizes[7]/L; - toGPU(ctx, iter, model->params.ln2w[l]); - iter += model->param_sizes[8]/L; - toGPU(ctx, iter, model->params.ln2b[l]); - iter += model->param_sizes[9]/L; - toGPU(ctx, iter, model->params.fcw[l]); - iter += model->param_sizes[10]/L; - toGPU(ctx, iter, model->params.fcb[l]); - iter += model->param_sizes[11]/L; - toGPU(ctx, iter, model->params.fcprojw[l]); - iter += model->param_sizes[12]/L; - toGPU(ctx, iter, model->params.fcprojb[l]); - iter += model->param_sizes[13]/L; - } - toGPU(ctx, iter, model->params.lnfw); - iter += model->param_sizes[14]; - toGPU(ctx, iter, model->params.lnfb); - iter += model->param_sizes[15]; -} - -void gpt2_free(GPT2 *model) { - free(model->params_memory); - free(model->grads_memory); - free(model->m_memory); - free(model->v_memory); - free(model->acts_memory); - free(model->grads_acts_memory); - // free(model->inputs); - // free(model->targets); - free(model->mean_loss_buffer); -} - -#ifndef TESTING -// if we are TESTING (see test_gpt2.c), we'll skip the int main below -// ---------------------------------------------------------------------------- -// sampler - -unsigned int random_u32(uint64_t *state) { - // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A - *state ^= *state >> 12; - *state ^= *state << 25; - *state ^= *state >> 27; - return (*state * 0x2545F4914F6CDD1Dull) >> 32; -} -float random_f32(uint64_t *state) { // random float32 in [0,1) - return (random_u32(state) >> 8) / 16777216.0f; -} - -int sample_mult(float* probabilities, int n, float coin) { - // sample index from probabilities (they must sum to 1!) - // coin is a random number in [0, 1), usually from random_f32() - float cdf = 0.0f; - for (int i = 0; i < n; i++) { - cdf += probabilities[i]; - if (coin < cdf) { - return i; - } - } - return n - 1; // in case of rounding errors -} - -// ---------------------------------------------------------------------------- -// main training loop -int main() { - - setLogLevel(kWarn); - - printf("Creating GPU context\n"); - WGPURequiredLimits requiredLimits = LIMITS_BUFFER_SIZE_1GB; - gpu::Context ctx = gpu::createContext({}, {}, { - .requiredLimits = &requiredLimits - }); - - // build the GPT-2 model from a checkpoint - GPT2 model; - gpt2_build_from_checkpoint(ctx, &model, "gpt2_124M.bin"); - - // build the DataLoaders from tokens files. for now use tiny_shakespeare if available, else tiny_stories - const char* tiny_stories_train = "dev/data/tinystories/TinyStories_train.bin"; - const char* tiny_stories_val = "dev/data/tinystories/TinyStories_val.bin"; - const char* tiny_shakespeare_train = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; - const char* tiny_shakespeare_val = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; - const char* train_tokens = access(tiny_shakespeare_train, F_OK) != -1 ? tiny_shakespeare_train : tiny_stories_train; - const char* val_tokens = access(tiny_shakespeare_val, F_OK) != -1 ? tiny_shakespeare_val : tiny_stories_val; - constexpr int B = 4; // batch size 4 (i.e. 4 independent token sequences will be trained on) - constexpr int T = 64; // sequence length 64 (i.e. each sequence is 64 tokens long). must be <= maxT, which is 1024 for GPT-2 - DataLoader train_loader, val_loader; - dataloader_init(&train_loader, train_tokens, B, T, 0, 1, 1); - dataloader_init(&val_loader, val_tokens, B, T, 0, 1, 0); - printf("train dataset num_batches: %zu\n", train_loader.num_tokens / (B*T)); - printf("val dataset num_batches: %zu\n", val_loader.num_tokens / (B*T)); - int val_num_batches = 5; - - // build the Tokenizer - Tokenizer tokenizer; - tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); - - // some memory for generating samples from the model - uint64_t rng_state = 1337; - // int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); - const int genT = 64; // number of steps of inference we will do - - // train - struct timespec start, end; - Tensor inputs = createTensor(ctx, Shape{B, T}, ki32); - Tensor targets = createTensor(ctx, Shape{B, T}, ki32); - Tensor gen_tokens = createTensor(ctx, Shape{B, T}, ki32); - int* gen_tokens_cpu = (int*)mallocCheck(B * T * sizeof(int)); - printf("Starting training\n"); - for (int step = 0; step <= 40; step++) { - printf("Step %d\n", step); - - // once in a while estimate the validation loss - if (step % 10 == 0) { - float val_loss = 0.0f; - dataloader_reset(&val_loader); - for (int i = 0; i < val_num_batches; i++) { - dataloader_next_batch(&val_loader); - toGPU(ctx, val_loader.inputs, inputs); - toGPU(ctx, val_loader.targets, targets); - gpt2_forward(ctx, &model, inputs, targets, B, T); - val_loss += model.mean_loss; - } - val_loss /= val_num_batches; - printf("val loss %f\n", val_loss); - } - - // once in a while do model inference to print generated text - if (step > 0 && step % 20 == 0) { - // fill up gen_tokens with the GPT2_EOT, which kicks off the generation - for(int i = 0; i < B * T; ++i) { - gen_tokens_cpu[i] = tokenizer.eot_token; - } - toGPU(ctx, gen_tokens_cpu, gen_tokens); - // now sample from the model autoregressively - printf("generating:\n---\n"); - for (int t = 1; t < genT; t++) { - // note that inference is very wasteful here because for each token - // we re-calculate the forward pass for all of (B,T) positions from scratch - // but the inference here is just for sanity checking anyway - // and we can maybe optimize a bit more later, with careful tests - gpt2_forward(ctx, &model, gen_tokens, model.nullTensor, B, T); - // furthermore, below we're only using b=0 (i.e. the first row) of all B rows - // we're in principle running B "inference streams" in parallel here - // but only using position 0 - // get the Vp-dimensional vector probs[0, t-1, :] - toCPU(ctx, model.acts.probs, model.probs_buffer, B * T * model.config.padded_vocab_size * sizeof(float)); - float* probs = model.probs_buffer + (t-1) * model.config.padded_vocab_size; - - float coin = random_f32(&rng_state); - // note we're only sampling from the first V elements, ignoring padding - // (the probabilities in the padded region should be zero anyway) - int next_token = sample_mult(probs, model.config.vocab_size, coin); - gen_tokens_cpu[t] = next_token; - toGPU(ctx, gen_tokens_cpu, gen_tokens); - // print the generated token, either using the Tokenizer or a fallback - if (tokenizer.init_ok) { - const char* token_str = tokenizer_decode(&tokenizer, next_token); - safe_printf(token_str); - } else { - // fall back to printing the token id - printf("%d ", next_token); - } - fflush(stdout); - } - printf("\n---\n"); - } - - // do a training step - clock_gettime(CLOCK_MONOTONIC, &start); - dataloader_next_batch(&train_loader); - toGPU(ctx, train_loader.inputs, inputs); - toGPU(ctx, train_loader.targets, targets); - gpt2_forward(ctx, &model, inputs, targets, B, T); - if (model.backward_enabled) { - gpt2_zero_grad(&model); - gpt2_backward(ctx, &model); - gpt2_update(ctx, &model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, step+1); - } - clock_gettime(CLOCK_MONOTONIC, &end); - double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9; - printf("step %d: train loss %f (took %f ms)\n", step, model.mean_loss, time_elapsed_s * 1000); - } - - // free - dataloader_free(&train_loader); - dataloader_free(&val_loader); - tokenizer_free(&tokenizer); - gpt2_free(&model); - // free(gen_tokens); - return 0; -} -#endif diff --git a/experimental/kernels/kernels.h b/experimental/kernels/kernels.h deleted file mode 100644 index 62a461e..0000000 --- a/experimental/kernels/kernels.h +++ /dev/null @@ -1,858 +0,0 @@ -#ifndef KERNELS_H -#define KERNELS_H - -#include "gpu.hpp" - -namespace gpu { - - -static const char *kShaderGelu = R"( -const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI) -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&inp)) { - let x: {{precision}} = inp[i]; - // select is more stable for larger values of x - out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR - * (x + .044715 * x * x * x))), x, x > 10.0); - } -} -)"; - -static const char *kShaderGeluBackward = R"( -const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI) -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var dout: array<{{precision}}>; -@group(0) @binding(2) var dinp: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&inp)) { - let x: {{precision}} = inp[i]; - let cube: {{precision}} = 0.044715f * x * x * x; - let tanh_arg: {{precision}} = GELU_SCALING_FACTOR * (x + cube); - let tanh_out: {{precision}} = tanh(tanh_arg); - let cosh_out: {{precision}} = cosh(tanh_arg); - let sech_out: {{precision}} = 1.0f / (cosh_out * cosh_out); - let local_grad: {{precision}} = select(0.5f * (1.0f + tanh_out), 1, x > 10.0) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); - dinp[i] += local_grad * dout[i]; - } -} -)"; - -static const char *kShaderTanh = R"( -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&inp)) { - let x: f32 = inp[i]; - out[i] = tanh(x); - } -} -)"; - -static const char *kShaderHadamard = R"( -@group(0) @binding(0) var A: array<{{precision}}>; -@group(0) @binding(1) var B: array<{{precision}}>; -@group(0) @binding(2) var C: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let idx = GlobalInvocationID.x; - if (idx < arrayLength(&A)) { - C[idx] = A[idx] * B[idx]; - } -} -)"; - -static const char *kShaderResidual = R"( -@group(0) @binding(0) var A: array<{{precision}}>; -@group(0) @binding(1) var B: array<{{precision}}>; -@group(0) @binding(2) var C: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let idx = GlobalInvocationID.x; - if (idx < arrayLength(&A)) { - C[idx] = A[idx] + B[idx]; - } -} -)"; - -static const char *kShaderResidualBackward = R"( -@group(0) @binding(0) var dout: array<{{precision}}>; -@group(0) @binding(1) var dinp1: array<{{precision}}>; -@group(0) @binding(2) var dinp2: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&dout)) { - dinp1[i] += dout[i]; - dinp2[i] += dout[i]; - } -} -)"; - -/* LayerNorm - * v1: - * - No caching mean/std for backwards - * - No parallel reduction - * - Simple 1 thread for each 1..N - */ -// TODO(avh): Allow larger virtual 1D workgroups by making use of y / z -// dimensions and calculating the threadID accordingly. -static const char *kShaderLayerNorm1 = R"( -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var weight: array<{{precision}}>; -@group(0) @binding(2) var bias: array<{{precision}}>; -@group(0) @binding(3) var out: array<{{precision}}>; -@group(0) @binding(4) var params: Params; - -struct Params { - N: u32, - C: u32, -}; - -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) GlobalInvocationID: vec3, - @builtin(local_invocation_id) LocalInvocationID: vec3, - @builtin(workgroup_id) WorkgroupID: vec3) { - let idx: u32 = GlobalInvocationID.x; - - if (idx >= params.N) { return; } - - let C: u32 = params.C; - - // Calculate mean - var sum: f32 = 0.0; - for (var i: u32 = 0; i < C; i = i + 1) { - sum += inp[idx * C + i]; - } - let mean_val: f32 = sum / f32(C); - - // Calculate rstd - sum = 0.0; - for (var i: u32 = 0; i < C; i = i + 1) { - let diff: f32 = inp[idx * C + i] - mean_val; - sum += diff * diff; - } - let rstd_val: f32 = 1.0 / sqrt(sum / f32(C) + 1e-5); - - for (var i: u32 = 0; i < C; i = i + 1) { - let n: f32 = rstd_val * (inp[idx * C + i] - mean_val); - out[idx * C + i] = n * weight[i] + bias[i]; - } -} -)"; - -/* Softmax - * v1: - * - equivalent to naive softmax with one thread per row - */ -static const char *kShaderSoftmax1 = R"( -@group(0) @binding(0) var inp : array<{{precision}}>; -@group(0) @binding(1) var out : array<{{precision}}>; -@group(0) @binding(2) var params : Params; -struct Params { - N: u32, - C: u32, - Cp: u32, -}; -const NEG_INFINITY: f32 = -3.0e38; // WGSL has problem representing -3.4028235e+38 -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let N : u32 = params.N; - let C : u32 = params.C; - let Cp : u32 = params.Cp; - let i : u32 = global_id.x; - if (i < N) { - let inp_row_start : u32 = i * Cp; - var maxval : f32 = NEG_INFINITY; - // Find the maximum value in the row - for (var j : u32 = 0u; j < C; j++) { - let val : f32 = inp[inp_row_start + j]; - if (val > maxval) { - maxval = val; - } - } - var sum : f32 = 0.0; - // Compute the exponentials and sum them - for (var j : u32 = 0u; j < C; j++) { - let exp_val : f32 = exp(inp[inp_row_start + j] - maxval); - out[inp_row_start + j] = exp_val; - sum += exp_val; - } - // Normalize the row to get probabilities - let norm : f32 = 1.0f / sum; - for (var j : u32 = 0u; j < C; j++) { - out[inp_row_start + j] /= sum; - } - for (var j : u32 = C; j < Cp; j++) { - out[inp_row_start + j] = 0; - } - } -} -)"; - -// Encoder -static const char *kShaderEncoder = R"( -@group(0) @binding(0) var inp : array; -@group(0) @binding(1) var wte : array<{{precision}}>; -@group(0) @binding(2) var wpe : array<{{precision}}>; -@group(0) @binding(3) var out : array<{{precision}}>; -@group(0) @binding(4) var params : Params; -struct Params { - B: u32, - T: u32, - C: u32, -}; -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - if (b < B && t < T) { - let ix : u32 = u32(inp[b * T + t]); - let out_bt : u32 = b * T * C + t * C; - for (var i : u32 = 0u; i < C; i++) { - out[out_bt + i] = wte[ix * C + i] + wpe[t * C + i]; - } - } -} -)"; - -static const char *kShaderEncoderBackward = R"( -@group(0) @binding(0) var dwte : array<{{precision}}>; -@group(0) @binding(1) var dwpe : array<{{precision}}>; -@group(0) @binding(2) var dout : array<{{precision}}>; -@group(0) @binding(3) var inp : array; -@group(0) @binding(4) var params : Params; -struct Params { - B: u32, - T: u32, - C: u32, -}; -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - if (b < B && t < T) { - let ix : u32 = u32(inp[b * T + t]); - let dout_bt : u32 = b * T * C + t * C; - for (var i : u32 = 0u; i < C; i++) { - let d : {{precision}} = dout[dout_bt + i]; - atomicAdd(&dwte[ix * C + i], d); - atomicAdd(&dwpe[t * C + i], d); - } - } -} -)"; - - -// Matmul -static const char *kShaderMatmul = R"( -@group(0) @binding(0) var inp : array<{{precision}}>; -@group(0) @binding(1) var weight : array<{{precision}}>; -@group(0) @binding(2) var bias : array<{{precision}}>; -@group(0) @binding(3) var out : array<{{precision}}>; -@group(0) @binding(4) var params : Params; -struct Params { - B: u32, - T: u32, - C: u32, - OC: u32, -}; -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - let OC : u32 = params.OC; - // N == B*T == global_id.x - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - if (arrayLength(&bias) == 1) { - if (b < B && t < T) { - let bt : u32 = global_id.x; - for (var o : u32 = 0u; o < OC; o++) { - var val : {{precision}} = 0; - for (var i : u32 = 0u; i < C; i++) { - val += inp[bt * C + i] * weight[o * C + i]; - } - out[bt * OC + o] = val; - } - } - } else { - if (b < B && t < T) { - let bt : u32 = global_id.x; - for (var o : u32 = 0u; o < OC; o++) { - var val : {{precision}} = bias[o]; - for (var i : u32 = 0u; i < C; i++) { - val += inp[bt * C + i] * weight[o * C + i]; - } - out[bt * OC + o] = val; - } - } - } -} - -)"; - - -static const char *kShaderMatmul2DTiling = R"( -@group(0) @binding(0) var inp : array<{{precision}}>; -@group(0) @binding(1) var weight : array<{{precision}}>; -@group(0) @binding(2) var bias : array<{{precision}}>; -@group(0) @binding(3) var out : array<{{precision}}>; -@group(0) @binding(4) var params : Params; -struct Params { - B: u32, - T: u32, - C: u32, - OC: u32, -}; -var tileInp: array<{{precision}}, {{BT}} * {{BC}}>; -var tileWeight: array<{{precision}}, {{BOC}} * {{BC}}>; - -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(local_invocation_id) localID : vec3, - @builtin(workgroup_id) groupid : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - let OC : u32 = params.OC; - - var localT: array<{{precision}}, {{TT}}>; - var localOC: array<{{precision}}, {{TOC}}>; - - let outB: u32 = groupid.x; - let outT: u32 = groupid.y; - let outOC: u32 = groupid.z; - let numThread: u32 = ({{BT}} * {{BOC}}) / ({{TT}} * {{TOC}}); - - // position of the first c element computed by the thread - let threadRow: u32 = (localID.x / ({{BOC}} / {{TOC}})) * {{TT}}; - let threadCol: u32 = (localID.x % ({{BOC}} / {{TOC}})) * {{TOC}}; - - // inpPtr and weightPtr are the starting positions of the tiles in a and b, - // incremented in the bkidx loop. - // outPtr is the starting position of the tile in c which is fixed. - - var inpPtr = (outB * T + outT * {{BT}}) * C; // BTC - var weightPtr = outOC * {{BOC}} * C; //OCC - var threadResults: array<{{precision}}, {{TT}} * {{TOC}}>; - let outPtr = (outB * T + outT * {{BT}}) * OC + outOC * {{BOC}}; //BTOC - let biasPtr = outOC * {{BOC}}; - - for (var bkidx: u32 = 0; bkidx < C; bkidx += {{BC}}) { - // Load BC x BOC by numThread(BT * BOC / (TT * TOC)) - // The number of iteration == BC * BOC / (BT * BOC / (TT * TOC)) - for (var idx: u32 = 0; idx < {{NUM_TILEW}}; idx++) { - tileWeight[localID.x + idx * numThread] = weight[weightPtr + ((localID.x + idx * numThread) / {{BC}}) * C + ((localID.x + idx * numThread) % {{BC}})]; - } - weightPtr += {{BC}}; - - // Load tile - // Load BT x BC by numThread(BT * BOC / (TT * TOC)) - // The number of iteration == BT * BC / (BT * BOC / (TT * TOC)) - for (var idx: u32 = 0; idx < {{NUM_TILEI}}; idx++) { - tileInp[localID.x + idx * numThread] = inp[inpPtr + ((localID.x + idx * numThread) / {{BC}}) * C + (localID.x + idx * numThread) % {{BC}}]; - } - inpPtr += {{BC}}; - - workgroupBarrier(); - // Compute tile - for (var dotIdx: u32 = 0; dotIdx < {{BC}}; dotIdx = dotIdx + 1) { - for (var idx: u32 = 0; idx < {{TT}}; idx++) { - localT[idx] = tileInp[(threadRow + idx) * {{BC}} + dotIdx]; - } - for (var idx: u32 = 0; idx < {{TOC}}; idx++) { - localOC[idx] = tileWeight[(threadCol + idx) * {{BC}} + dotIdx]; - } - for (var resIdxT: u32 = 0; resIdxT < {{TT}}; resIdxT++) { - for (var resIdxOC: u32 = 0; resIdxOC < {{TOC}}; resIdxOC++) { - threadResults[resIdxT * {{TOC}} + resIdxOC] += localT[resIdxT] * localOC[resIdxOC]; - } - } - } - workgroupBarrier(); - } - - if (arrayLength(&bias) == 1) { - for (var resIdxT: u32 = 0; resIdxT < {{TT}}; resIdxT++) { - for (var resIdxOC: u32 = 0; resIdxOC < {{TOC}}; resIdxOC++) { - out[outPtr + (threadRow + resIdxT) * OC + threadCol + resIdxOC] = threadResults[resIdxT * {{TOC}} + resIdxOC]; - } - } - } else { - for (var resIdxT: u32 = 0; resIdxT < {{TT}}; resIdxT++) { - for (var resIdxOC: u32 = 0; resIdxOC < {{TOC}}; resIdxOC++) { - out[outPtr + (threadRow + resIdxT) * OC + threadCol + resIdxOC] = threadResults[resIdxT * {{TOC}} + resIdxOC] + bias[biasPtr + threadCol + resIdxOC]; - } - } - } -} -)"; - -static const char *kShaderMatmulBackward = R"( -@group(0) @binding(0) var dinp : array<{{precision}}>; -@group(0) @binding(1) var dweight : array<{{precision}}>; -@group(0) @binding(2) var dbias : array<{{precision}}>; -@group(0) @binding(3) var dout : array<{{precision}}>; -@group(0) @binding(4) var inp : array<{{precision}}>; -@group(0) @binding(5) var weight : array<{{precision}}>; -@group(0) @binding(6) var params : Params; -struct Params { - B: u32, - T: u32, - C: u32, - OC: u32, -}; -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - let OC : u32 = params.OC; - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - if (b < B && t < T) { - let bt : u32 = b * T + t; - for (var o : u32 = 0u; o < OC; o++) { - let d : {{precision}} = dout[bt * OC + o]; - atomicAdd(&dbias[o], d); - for (var i : u32 = 0u; i < C; i++) { - atomicAdd(&dinp[bt * C + i], weight[o * C + i] * d); - atomicAdd(&dweight[o * C + i], inp[bt * C + i] * d); - } - } - } -} -)"; - -// Attention -static const char *kShaderAttention = R"( -@group(0) @binding(0) var inp : array<{{precision}}>; -@group(0) @binding(1) var preatt : array<{{precision}}>; -@group(0) @binding(2) var att : array<{{precision}}>; -@group(0) @binding(3) var out : array<{{precision}}>; -@group(0) @binding(4) var params : Params; -struct Params { - B: u32, - T: u32, - C: u32, - NH: u32, -}; -const NEG_INFINITY: f32 = -3.0e38; // WGSL has problem representing -3.4028235e+38 -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - let NH : u32 = params.NH; - let C3 : u32 = C * 3u; - let hs : u32 = C / NH; - let scale : {{precision}} = 1.0 / sqrt({{precision}}(hs)); - - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - - if (b < B && t < T) { - for (var h : u32 = 0u; h < NH; h++) { - let query_t : u32 = b * T * C3 + t * C3 + h * hs; - let preatt_bth : u32 = b * NH * T * T + h * T * T + t * T; - let att_bth : u32 = b * NH * T * T + h * T * T + t * T; - - // pass 1: calculate query dot key and maxval - var maxval : {{precision}} = NEG_INFINITY; - for (var t2 : u32 = 0u; t2 <= t; t2++) { - let key_t2 : u32 = b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key - - // (query_t) dot (key_t2) - var val : {{precision}} = 0.0; - for (var i : u32 = 0u; i < hs; i++) { - val += inp[query_t + i] * inp[key_t2 + i]; - } - val *= scale; - if (val > maxval) { - maxval = val; - } - - preatt[preatt_bth + t2] = val; - } - - // pass 2: calculate the exp and keep track of sum - // maxval is being calculated and subtracted only for numerical stability - var expsum : {{precision}} = 0.0; - for (var t2 : u32 = 0u; t2 <= t; t2++) { - let expv : {{precision}} = exp(preatt[preatt_bth + t2] - maxval); - expsum += expv; - att[att_bth + t2] = expv; - } - let expsum_inv : {{precision}} = select(0.0, 1.0 / expsum, expsum != 0.0); - - // pass 3: normalize to get the softmax - for (var t2 : u32 = 0u; t2 < T; t2++) { - if (t2 <= t) { - att[att_bth + t2] *= expsum_inv; - } else { - // causal attention mask. not strictly necessary to set to zero here - // only doing this explicitly for debugging and checking to PyTorch - att[att_bth + t2] = 0.0; - } - } - - // pass 4: accumulate weighted values into the output of attention - let out_bth : u32 = b * T * C + t * C + h * hs; - for (var i : u32 = 0u; i < hs; i++) { out[out_bth + i] = 0.0; } - for (var t2 : u32 = 0u; t2 <= t; t2++) { - let value_t2 : u32 = b * T * C3 + t2 * C3 + h * hs + C * 2u; // +C*2 because it's value - let att_btht2 : {{precision}} = att[att_bth + t2]; - for (var i : u32 = 0u; i < hs; i++) { - out[out_bth + i] += att_btht2 * inp[value_t2 + i]; - } - } - } - } -} -)"; - -static const char *kShaderAttentionBackward = R"( -@group(0) @binding(0) var dinp : array<{{precision}}>; -@group(0) @binding(1) var dpreatt : array<{{precision}}>; -@group(0) @binding(2) var datt : array<{{precision}}>; -@group(0) @binding(3) var dout : array<{{precision}}>; -@group(0) @binding(4) var inp : array<{{precision}}>; -@group(0) @binding(5) var att : array<{{precision}}>; -@group(0) @binding(6) var params : Params; -struct Params { - B: u32, - T: u32, - C: u32, - NH: u32, -}; -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - let NH : u32 = params.NH; - let C3 : u32 = C * 3u; - let hs : u32 = C / NH; - let scale : {{precision}} = 1.0 / sqrt({{precision}}(hs)); - - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - - if (b < B && t < T) { - for (var h : u32 = 0u; h < NH; h++) { - let att_bth : u32 = b * NH * T * T + h * T * T + t * T; - let datt_bth : u32 = b * NH * T * T + h * T * T + t * T; - let dpreatt_bth : u32 = b * NH * T * T + h * T * T + t * T; - let dquery_t : u32 = b * T * C3 + t * C3 + h * hs; - let query_t : u32 = b * T * C3 + t * C3 + h * hs; - - // backward pass 4, through the value accumulation - let dout_bth : u32 = b * T * C + t * C + h * hs; - for (var t2 : u32 = 0u; t2 <= t; t2++) { - let value_t2 : u32 = b * T * C3 + t2 * C3 + h * hs + C * 2u; // +C*2 because it's value - let dvalue_t2 : u32 = b * T * C3 + t2 * C3 + h * hs + C * 2u; // +C*2 because it's value - for (var i : u32 = 0u; i < hs; i++) { - // in the forward pass this was: - // out_bth[i] += att_bth[t2] * value_t2[i]; - // so now we have: - atomicAdd(&datt[datt_bth + t2], inp[value_t2 + i] * dout[dout_bth + i]); - atomicAdd(&dinp[dvalue_t2 + i], att[att_bth + t2] * dout[dout_bth + i]); - } - } - - // backward pass 2 & 3, the softmax - // note that softmax (like e.g. tanh) doesn't need the input (preatt) to backward - for (var t2 : u32 = 0u; t2 <= t; t2++) { - for (var t3 : u32 = 0u; t3 <= t; t3++) { - let indicator : {{precision}} = select(0.0, 1.0, t2 == t3); - let local_derivative : {{precision}} = att[att_bth + t2] * (indicator - att[att_bth + t3]); - atomicAdd(&dpreatt[dpreatt_bth + t3], local_derivative * datt[datt_bth + t2]); - } - } - - // backward pass 1, the query @ key matmul - for (var t2 : u32 = 0u; t2 <= t; t2++) { - let key_t2 : u32 = b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key - let dkey_t2 : u32 = b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key - for (var i : u32 = 0u; i < hs; i++) { - // in the forward pass this was: - // preatt_bth[t2] += (query_t[i] * key_t2[i]) * scale; - // so now we have: - atomicAdd(&dinp[dquery_t + i], inp[key_t2 + i] * dpreatt[dpreatt_bth + t2] * scale); - atomicAdd(&dinp[dkey_t2 + i], inp[query_t + i] * dpreatt[dpreatt_bth + t2] * scale); - } - } - } - } -} -)"; - -// LayerNorm -static const char *kShaderLayerNorm = R"( -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var weight: array<{{precision}}>; -@group(0) @binding(2) var bias: array<{{precision}}>; -@group(0) @binding(3) var out: array<{{precision}}>; -@group(0) @binding(4) var mean: array<{{precision}}>; -@group(0) @binding(5) var rstd: array<{{precision}}>; -@group(0) @binding(6) var params: Params; - -struct Params { - B: u32, - T: u32, - C: u32, -}; - -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) GlobalInvocationID: vec3) { - let idx: u32 = GlobalInvocationID.x; - - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - - if (idx >= B * T) { return; } - - let b : u32 = idx / T; - let t : u32 = idx % T; - - // Calculate mean - var sum: f32 = 0.0; - for (var i: u32 = 0; i < C; i = i + 1) { - sum += inp[b * T * C + t * C + i]; - } - let mean_val: f32 = sum / f32(C); - mean[b * T + t] = mean_val; - - // Calculate rstd - sum = 0.0; - for (var i: u32 = 0; i < C; i = i + 1) { - let diff: f32 = inp[b * T * C + t * C + i] - mean_val; - sum += diff * diff; - } - let rstd_val: f32 = 1.0 / sqrt(sum / f32(C) + 1e-5); - rstd[b * T + t] = rstd_val; - - for (var i: u32 = 0; i < C; i = i + 1) { - let n: f32 = rstd_val * (inp[b * T * C + t * C + i] - mean_val); - out[b * T * C + t * C + i] = n * weight[i] + bias[i]; - } -} -)"; - -static const char *kShaderLayerNormBackward = R"( -@group(0) @binding(0) var dinp: array<{{precision}}>; -@group(0) @binding(1) var dweight: array<{{precision}}>; -@group(0) @binding(2) var dbias: array<{{precision}}>; -@group(0) @binding(3) var dout: array<{{precision}}>; -@group(0) @binding(4) var inp: array<{{precision}}>; -@group(0) @binding(5) var weight: array<{{precision}}>; -@group(0) @binding(6) var mean: array<{{precision}}>; -@group(0) @binding(7) var rstd: array<{{precision}}>; -@group(0) @binding(8) var params: Params; - -struct Params { - B: u32, - T: u32, - C: u32, -}; - -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) GlobalInvocationID: vec3) { - let idx: u32 = GlobalInvocationID.x; - - let B : u32 = params.B; - let T : u32 = params.T; - let C : u32 = params.C; - - if (idx >= B * T) { return; } - - let b : u32 = idx / T; - let t : u32 = idx % T; - - // first: two reduce operations - var dnorm_mean: f32 = 0.0f; - var dnorm_norm_mean: f32 = 0.0f; - for (var i: u32 = 0; i < C; i = i + 1) { - let norm_bti: f32 = (inp[b * T * C + t * C + i] - mean[b * T + t]) * rstd[b * T + t]; - let dnorm_i: f32 = weight[i] * dout[b * T * C + t * C + i]; - dnorm_mean += dnorm_i; - dnorm_norm_mean += dnorm_i * norm_bti; - } - dnorm_mean = dnorm_mean / f32(C); - dnorm_norm_mean = dnorm_norm_mean / f32(C); - - // now iterate again and accumulate all the gradients - for (var i: u32 = 0; i < C; i = i + 1) { - let norm_bti: f32 = (inp[b * T * C + t * C + i] - mean[b * T + t]) * rstd[b * T + t]; - let dnorm_i: f32 = weight[i] * dout[b * T * C + t * C + i]; - // gradient contribution to bias - atomicAdd(&dbias[i], dout[b * T * C + t * C + i]); - // gradient contribution to weight - atomicAdd(&dweight[i], norm_bti * dout[b * T * C + t * C + i]); - // gradient contribution to input - var dval: f32 = 0.0f; - dval += dnorm_i; // term 1 - dval -= dnorm_mean; // term 2 - dval -= norm_bti * dnorm_norm_mean; // term 3 - dval *= rstd[b * T + t]; // final scale - atomicAdd(&dinp[b * T * C + t * C + i], dval); - } -} -)"; - -static const char *kShaderCrossEntropyForward = R"( -@group(0) @binding(0) var losses : array<{{precision}}>; -@group(0) @binding(1) var probs : array<{{precision}}>; -@group(0) @binding(2) var targets : array; -@group(0) @binding(3) var params : Params; -struct Params { - B: u32, - T: u32, - Vp: u32, -}; -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let Vp : u32 = params.Vp; - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - if (b < B && t < T) { - let probs_bt : u32 = b * T * Vp + t * Vp; - let ix : u32 = u32(targets[b * T + t]); - losses[b * T + t] = -log(probs[probs_bt + ix]); - } -} -)"; - -static const char *kShaderCrossEntropySoftmaxBackward = R"( -@group(0) @binding(0) var dlogits : array<{{precision}}>; -@group(0) @binding(1) var dlosses : array<{{precision}}>; -@group(0) @binding(2) var probs : array<{{precision}}>; -@group(0) @binding(3) var targets : array; -@group(0) @binding(4) var params : Params; -struct Params { - B: u32, - T: u32, - V: u32, - Vp: u32, -}; -@compute @workgroup_size({{workgroupSize}}) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let B : u32 = params.B; - let T : u32 = params.T; - let V : u32 = params.V; - let Vp : u32 = params.Vp; - let b : u32 = global_id.x / T; - let t : u32 = global_id.x % T; - if (b < B && t < T) { - let dlogits_bt : u32 = b * T * Vp + t * Vp; - let probs_bt : u32 = b * T * Vp + t * Vp; - let dloss : {{precision}} = dlosses[b * T + t]; - let ix : u32 = u32(targets[b * T + t]); - for (var i : u32 = 0u; i < V; i++) { - let p : {{precision}} = probs[probs_bt + i]; - let indicator : {{precision}} = select(0.0, 1.0, i == ix); - dlogits[dlogits_bt + i] += (p - indicator) * dloss; - } - } -} -)"; - -static const char *kSum = R"( -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -var buffer: array<{{precision}}, 1024>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) globalID : vec3, - @builtin(local_invocation_id) localID : vec3, - @builtin(workgroup_id) groupid : vec3, - @builtin(num_workgroups) numGroups : vec3) { - let blockSize3d: vec3 = vec3({{workgroupSize}}); - let blockSize: u32 = blockSize3d.x; - let threadId: u32 = localID.x; - let blockId: u32 = groupid.x + groupid.y * numGroups.x; - let blockStart = blockId * blockSize * 2 + threadId; - - buffer[threadId] = inp[blockStart] + inp[blockStart + blockSize]; - workgroupBarrier(); - var stride: u32 = blockSize / 2; - - if (blockSize >= 1024 && threadId < 512) { - buffer[threadId] += buffer[threadId + 512]; - } - workgroupBarrier(); - - if (blockSize >= 512 && threadId < 256) { - buffer[threadId] += buffer[threadId + 256]; - } - workgroupBarrier(); - - if (blockSize >= 256 && threadId < 128) { - buffer[threadId] += buffer[threadId + 128]; - } - workgroupBarrier(); - - if (threadId < 64) { - buffer[threadId] += buffer[threadId + 64]; - } - workgroupBarrier(); - - if (threadId < 32) { - buffer[threadId] += buffer[threadId + 32]; - } - workgroupBarrier(); - - if (threadId < 16) { - buffer[threadId] += buffer[threadId + 16]; - } - workgroupBarrier(); - - if (threadId < 8) { - buffer[threadId] += buffer[threadId + 8]; - } - workgroupBarrier(); - - if (threadId < 4) { - buffer[threadId] += buffer[threadId + 4]; - } - workgroupBarrier(); - - if (threadId < 2) { - buffer[threadId] += buffer[threadId + 2]; - } - workgroupBarrier(); - - if (threadId == 0) { - buffer[0] += buffer[1]; - out[blockId] = buffer[0]; - } -} -)"; - -} // namespace gpu - -#endif // KERNELS_H diff --git a/experimental/kernels/ops.cpp b/experimental/kernels/ops.cpp deleted file mode 100644 index 0e9c076..0000000 --- a/experimental/kernels/ops.cpp +++ /dev/null @@ -1,799 +0,0 @@ -#include "gpu.hpp" -#include -#include -#include -#include - -#include "kernels.h" -#include "ops.hpp" -#include "experimental/wgsl.h" // loopUnrolling - -using namespace gpu; - -void encoder_forward(Context& ctx, float* out, - int* inp, float* wte, float* wpe, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long v = VOCAB_SIZE; - struct EncoderParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "encoder_forward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor input = createTensor(ctx, Shape{b * t}, ki32); - Tensor wte_t = createTensor(ctx, Shape{v, c}, kf32); - Tensor wpe_t = createTensor(ctx, Shape{t, c}, kf32); - Tensor output = createTensor(ctx, Shape{b * t * c}, kf32); - op = createKernel(ctx, {kShaderEncoder, 256, kf32}, - Bindings{input, wte_t, wpe_t, output}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - EncoderParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& input = ctx.pool.data[op->buffers[0]]; - Tensor& wte_t = ctx.pool.data[op->buffers[1]]; - Tensor& wpe_t = ctx.pool.data[op->buffers[2]]; - Tensor& output = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, inp, input); - toGPU(ctx, wte, wte_t); - toGPU(ctx, wpe, wpe_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, out, b * t * c * sizeof(float)); -} - -void encoder_backward(Context& ctx, float* dwte, float* dwpe, - float* dout, int* inp, - int B, int T, int C) { - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long v = VOCAB_SIZE; - struct EncoderParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "encoder_backward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dwte_t = createTensor(ctx, Shape{v, c}, kf32); - Tensor dwpe_t = createTensor(ctx, Shape{t, c}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor input = createTensor(ctx, Shape{b * t}, ki32); - op = createKernel(ctx, {kShaderEncoderBackward, 256, kf32}, - Bindings{dwte_t, dwpe_t, dout_t, input}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - EncoderParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dwte_t = ctx.pool.data[op->buffers[0]]; - Tensor& dwpe_t = ctx.pool.data[op->buffers[1]]; - Tensor& dout_t = ctx.pool.data[op->buffers[2]]; - Tensor& input = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, dwte, dwte_t); - toGPU(ctx, dwpe, dwpe_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, input); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dwte_t, dwte, v * c * sizeof(float)); - toCPU(ctx, dwpe_t, dwpe, t * c * sizeof(float)); -} - -void layernorm_forward(Context& ctx, float* out, float* mean, float* rstd, - float* inp, float* weight, float* bias, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - struct LayerNormParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "layernorm_forward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_t = createTensor(ctx, Shape{c}, kf32); - Tensor bias_t = createTensor(ctx, Shape{c}, kf32); - Tensor out_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor mean_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor rstd_t = createTensor(ctx, Shape{b * t}, kf32); - op = createKernel(ctx, {kShaderLayerNorm, 256, kf32}, - Bindings{inp_t, weight_t, bias_t, out_t, mean_t, rstd_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - LayerNormParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_t = ctx.pool.data[op->buffers[0]]; - Tensor& weight_t = ctx.pool.data[op->buffers[1]]; - Tensor& bias_t = ctx.pool.data[op->buffers[2]]; - Tensor& out_t = ctx.pool.data[op->buffers[3]]; - Tensor& mean_t = ctx.pool.data[op->buffers[4]]; - Tensor& rstd_t = ctx.pool.data[op->buffers[5]]; - - toGPU(ctx, inp, inp_t); - toGPU(ctx, weight, weight_t); - toGPU(ctx, bias, bias_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, out_t, out, b * t * c * sizeof(float)); - toCPU(ctx, mean_t, mean, b * t * sizeof(float)); - toCPU(ctx, rstd_t, rstd, b * t * sizeof(float)); -} - -void layernorm_backward(Context& ctx, float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* mean, float* rstd, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - struct LayerNormParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "layernorm_backward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dinp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor dweight_t = createTensor(ctx, Shape{c}, kf32); - Tensor dbias_t = createTensor(ctx, Shape{c}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_t = createTensor(ctx, Shape{c}, kf32); - Tensor mean_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor rstd_t = createTensor(ctx, Shape{b * t}, kf32); - op = createKernel(ctx, {kShaderLayerNormBackward, 256, kf32}, - Bindings{dinp_t, dweight_t, dbias_t, dout_t, inp_t, weight_t, mean_t, rstd_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - LayerNormParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dinp_t = ctx.pool.data[op->buffers[0]]; - Tensor& dweight_t = ctx.pool.data[op->buffers[1]]; - Tensor& dbias_t = ctx.pool.data[op->buffers[2]]; - Tensor& dout_t = ctx.pool.data[op->buffers[3]]; - Tensor& inp_t = ctx.pool.data[op->buffers[4]]; - Tensor& weight_t = ctx.pool.data[op->buffers[5]]; - Tensor& mean_t = ctx.pool.data[op->buffers[6]]; - Tensor& rstd_t = ctx.pool.data[op->buffers[7]]; - - toGPU(ctx, dinp, dinp_t); - toGPU(ctx, dweight, dweight_t); - toGPU(ctx, dbias, dbias_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, inp_t); - toGPU(ctx, weight, weight_t); - toGPU(ctx, mean, mean_t); - toGPU(ctx, rstd, rstd_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_t, dinp, b * t * c * sizeof(float)); - toCPU(ctx, dweight_t, dweight, c * sizeof(float)); - toCPU(ctx, dbias_t, dbias, c * sizeof(float)); -} - -struct DurationTime { - std::chrono::high_resolution_clock::time_point start; - std::chrono::high_resolution_clock::time_point end; - std::chrono::microseconds duration; - std::string src; - bool verbose; - - inline DurationTime(const std::string& src, bool verbose = true) { - this->src = src; - this->verbose = verbose; - start = std::chrono::high_resolution_clock::now(); - } - - inline ~DurationTime() { - end = std::chrono::high_resolution_clock::now(); - duration = std::chrono::duration_cast(end - start); - if (this->verbose) { - printf("Duration(%s): %.1f microseconds\n", src.c_str(), static_cast(duration.count())); - } - } -}; - - -void matmul_forward(Context& ctx, float* out, - const float* inp, const float* weight, const float* bias, - int B, int T, int C, int OC){ - bool verbose = false; - DurationTime duration("matmul_forward_gpu", verbose); - struct MatmulParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t OC; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long oc = static_cast(OC); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "matmul_forward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(OC); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - constexpr size_t BT = 64; - constexpr size_t BC = 8; - constexpr size_t BOC = 64; - constexpr size_t TT = BT / BC; - constexpr size_t TOC = BOC / BC; - size_t num_threads = BT * BOC / (TT * TOC); - Shape wgSize = {num_threads, 1, 1}; - Shape nWorkgroups = {b, cdiv(T, BT), cdiv(OC, BOC)}; - - std::string kShaderMatmul2DTiling_(kShaderMatmul2DTiling); - std::string kShaderMatmul2D(loopUnrolling( - replaceAll(kShaderMatmul2DTiling_, - {{"{{precision}}", toString(kf32)}, - {"{{BT}}", toString(BT)}, - {"{{BC}}", toString(BC)}, - {"{{BOC}}", toString(BOC)}, - {"{{TT}}", toString(TT)}, - {"{{TOC}}", toString(TOC)}, - {"{{NUM_TILEI}}", toString(BT * BC / num_threads)}, - {"{{NUM_TILEW}}", toString(BOC * BC / num_threads)} - }) - ) - ); - - Tensor inp_i = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_i = createTensor(ctx, Shape{oc * c}, kf32); - Tensor bias_i = bias == NULL ? createTensor(ctx, Shape{1}, kf32) : createTensor(ctx, Shape{oc}, kf32); - Tensor out_o = createTensor(ctx, Shape{b * t * oc}, kf32); - - op = createKernel(ctx, {kShaderMatmul2D, wgSize, kf32}, - Bindings{inp_i, weight_i, bias_i, out_o}, - nWorkgroups, - /* params */ - MatmulParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(oc) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_i = ctx.pool.data[op->buffers[0]]; - Tensor& weight_i = ctx.pool.data[op->buffers[1]]; - Tensor& bias_i = ctx.pool.data[op->buffers[2]]; - Tensor& out_o = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, inp, inp_i); - toGPU(ctx, weight, weight_i); - if (bias != NULL) { - toGPU(ctx, bias, bias_i); - } - - std::promise promise; - std::future future = promise.get_future(); - { - DurationTime duration("matmul_forward_gpu without creating tensors", verbose); - { - DurationTime duration("matmul_forward_gpu without creating kernel", verbose); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, out_o, out, b * t * oc * sizeof(float)); - } - } - toCPU(ctx, out_o, out, b * t * oc * sizeof(float)); -} - -void matmul_backward(Context& ctx, float* dinp, float* dweight, float* dbias, - const float* dout, const float* inp, const float* weight, - int B, int T, int C, int OC){ - struct MatmulParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t OC; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long oc = static_cast(OC); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "matmul_backward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(OC); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dinp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor dweight_t = createTensor(ctx, Shape{oc * c}, kf32); - Tensor dbias_t = createTensor(ctx, Shape{oc}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * oc}, kf32); - Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_t = createTensor(ctx, Shape{oc * c}, kf32); - op = createKernel(ctx, {kShaderMatmulBackward, 256, kf32}, - Bindings{dinp_t, dweight_t, dbias_t, dout_t, inp_t, weight_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - MatmulParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(oc) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dinp_t = ctx.pool.data[op->buffers[0]]; - Tensor& dweight_t = ctx.pool.data[op->buffers[1]]; - Tensor& dbias_t = ctx.pool.data[op->buffers[2]]; - Tensor& dout_t = ctx.pool.data[op->buffers[3]]; - Tensor& inp_t = ctx.pool.data[op->buffers[4]]; - Tensor& weight_t = ctx.pool.data[op->buffers[5]]; - - toGPU(ctx, dinp, dinp_t); - toGPU(ctx, dweight, dweight_t); - toGPU(ctx, dbias, dbias_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, inp_t); - toGPU(ctx, weight, weight_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_t, dinp, b * t * c * sizeof(float)); - toCPU(ctx, dweight_t, dweight, oc * c * sizeof(float)); - toCPU(ctx, dbias_t, dbias, oc * sizeof(float)); -} - -void attention_forward(Context& ctx, float* out, float* preatt, float* att, - float* inp, - int B, int T, int C, int NH){ - struct AttentionParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t NH; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long nh = static_cast(NH); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "attention_forward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(NH); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32); - Tensor preatt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor att_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor out_t = createTensor(ctx, Shape{b * t * c}, kf32); - op = createKernel(ctx, {kShaderAttention, 256, kf32}, - Bindings{inp_t, preatt_t, att_t, out_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - AttentionParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(nh) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_t = ctx.pool.data[op->buffers[0]]; - Tensor& preatt_t = ctx.pool.data[op->buffers[1]]; - Tensor& att_t = ctx.pool.data[op->buffers[2]]; - Tensor& out_t = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, inp, inp_t); - toGPU(ctx, preatt, preatt_t); - toGPU(ctx, att, att_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, preatt_t, preatt, b * nh * t * t * sizeof(float)); - toCPU(ctx, att_t, att, b * nh * t * t * sizeof(float)); - toCPU(ctx, out_t, out, b * t * c * sizeof(float)); -} - -void attention_backward(Context& ctx, float* dinp, float* dpreatt, float* datt, - float* dout, float* inp, float* att, - int B, int T, int C, int NH){ - struct AttentionParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t NH; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long nh = static_cast(NH); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "attention_backward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(NH); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dinp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32); - Tensor dpreatt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor datt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor inp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32); - Tensor att_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - op = createKernel(ctx, {kShaderAttentionBackward, 256, kf32}, - Bindings{dinp_t, dpreatt_t, datt_t, dout_t, inp_t, att_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - AttentionParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(nh) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dinp_t = ctx.pool.data[op->buffers[0]]; - Tensor& dpreatt_t = ctx.pool.data[op->buffers[1]]; - Tensor& datt_t = ctx.pool.data[op->buffers[2]]; - Tensor& dout_t = ctx.pool.data[op->buffers[3]]; - Tensor& inp_t = ctx.pool.data[op->buffers[4]]; - Tensor& att_t = ctx.pool.data[op->buffers[5]]; - - toGPU(ctx, dinp, dinp_t); - toGPU(ctx, dpreatt, dpreatt_t); - toGPU(ctx, datt, datt_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, inp_t); - toGPU(ctx, att, att_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_t, dinp, b * t * c * 3 * sizeof(float)); - toCPU(ctx, dpreatt_t, dpreatt, b * nh * t * t * sizeof(float)); - toCPU(ctx, datt_t, datt, b * nh * t * t * sizeof(float)); -} - -void gelu_forward(Context& ctx, float* out, float* inp, int n) { - unsigned long N = static_cast(n); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "gelu_forward_" + std::to_string(n); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor input = createTensor(ctx, Shape{N}, kf32); - Tensor output = createTensor(ctx, Shape{N}, kf32); - op = createKernel(ctx, {kShaderGelu, 256, kf32}, - Bindings{input, output}, - /* nWorkgroups */ {cdiv(N, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& input = ctx.pool.data[op->buffers[0]]; - Tensor& output = ctx.pool.data[op->buffers[1]]; - - toGPU(ctx, inp, input); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, out, N * sizeof(float)); -} - -void gelu_backward(Context& ctx, float* dinp, float* inp, float* dout, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "gelu_backward_" + std::to_string(N); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp_i = createTensor(ctx, Shape{n}, kf32); - Tensor dout_i = createTensor(ctx, Shape{n}, kf32); - Tensor dinp_o = createTensor(ctx, Shape{n}, kf32); - op = createKernel(ctx, {kShaderGeluBackward, 256, kf32}, - Bindings{inp_i, dout_i, dinp_o}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_i = ctx.pool.data[op->buffers[0]]; - Tensor& dout_i = ctx.pool.data[op->buffers[1]]; - Tensor& dinp_o = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, inp, inp_i); - toGPU(ctx, dout, dout_i); - toGPU(ctx, dinp, dinp_o); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_o, dinp, n * sizeof(float)); -} - -void residual_forward(Context& ctx, float* out, float* inp1, float* inp2, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "residual_forward_" + std::to_string(N); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp1_i = createTensor(ctx, Shape{n}, kf32); - Tensor inp2_i = createTensor(ctx, Shape{n}, kf32); - Tensor out_o = createTensor(ctx, Shape{n}, kf32); - op = createKernel(ctx, {kShaderResidual, 256, kf32}, - Bindings{inp1_i, inp2_i, out_o}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp1_i = ctx.pool.data[op->buffers[0]]; - Tensor& inp2_i = ctx.pool.data[op->buffers[1]]; - Tensor& out_o = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, inp1, inp1_i); - toGPU(ctx, inp2, inp2_i); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, out_o, out, n * sizeof(float)); -} - -void residual_backward(Context& ctx, float* dinp1, float* dinp2, float* dout, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "residual_backward_" + std::to_string(N); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dout_i = createTensor(ctx, Shape{n}, kf32); - Tensor dinp1_o = createTensor(ctx, Shape{n}, kf32); - Tensor dinp2_o = createTensor(ctx, Shape{n}, kf32); - op = createKernel(ctx, {kShaderResidualBackward, 256, kf32}, - Bindings{dout_i, dinp1_o, dinp2_o}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dout_i = ctx.pool.data[op->buffers[0]]; - Tensor& dinp1_o = ctx.pool.data[op->buffers[1]]; - Tensor& dinp2_o = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, dout, dout_i); - toGPU(ctx, dinp1, dinp1_o); - toGPU(ctx, dinp2, dinp2_o); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp1_o, dinp1, n * sizeof(float)); - toCPU(ctx, dinp2_o, dinp2, n * sizeof(float)); -} - -void softmax_forward(Context& ctx, float* probs, float* logits, int B, int T, int V, int Vp) { - struct SoftmaxParam { - uint32_t N; - uint32_t C; - uint32_t Cp; - }; - uint32_t b = static_cast(B); - uint32_t t = static_cast(T); - uint32_t c = static_cast(V); - uint32_t cp = static_cast(Vp); - // Generate the key of the cache by arguments. - std::string key = "softmax_forward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(V) + "_" + std::to_string(Vp); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor input = createTensor(ctx, {b * t, cp}, kf32); - Tensor output = createTensor(ctx, {b * t, cp}, kf32); - assert( (B*T) % 256 == 0); - op = createKernel( - ctx, {kShaderSoftmax1, 256, kf32}, Bindings{input, output}, - Shape{cdiv(B * T, 256), 1, 1}, SoftmaxParam{b * t, c, cp}, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& input = ctx.pool.data[op->buffers[0]]; - Tensor& output = ctx.pool.data[op->buffers[1]]; - - toGPU(ctx, logits, input); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, probs, sizeof(float)*b*t*cp); -} - -void crossentropy_forward(Context& ctx, float* losses, - float* probs, int* targets, - int B, int T, int Vp){ - struct CrossEntropyParams { - uint32_t B; - uint32_t T; - uint32_t VP; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long vp = static_cast(Vp); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "crossentropy_forward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(Vp); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor losses_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor probs_t = createTensor(ctx, Shape{b * t * vp}, kf32); - Tensor targets_t = createTensor(ctx, Shape{b * t}, ki32); - op = createKernel(ctx, {kShaderCrossEntropyForward, 256, kf32}, - Bindings{losses_t, probs_t, targets_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - CrossEntropyParams{ - static_cast(b), - static_cast(t), - static_cast(vp) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& losses_t = ctx.pool.data[op->buffers[0]]; - Tensor& probs_t = ctx.pool.data[op->buffers[1]]; - Tensor& targets_t = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, losses, losses_t); - toGPU(ctx, probs, probs_t); - toGPU(ctx, targets, targets_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, losses_t, losses, b * t * sizeof(float)); -} - -void crossentropy_softmax_backward(Context& ctx, float* dlogits, - float* dlosses, float* probs, int* targets, - int B, int T, int V, int Vp){ - struct CrossEntropySoftmaxBackwardParams { - uint32_t B; - uint32_t T; - uint32_t V; - uint32_t VP; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long v = static_cast(V); - unsigned long vp = static_cast(Vp); - setLogLevel(kError); - // Generate the key of the cache by arguments. - std::string key = "crossentropy_softmax_backward_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(V) + "_" + std::to_string(Vp); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dlogits_t = createTensor(ctx, Shape{b * t * vp}, kf32); - Tensor dlosses_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor probs_t = createTensor(ctx, Shape{b * t * vp}, kf32); - Tensor targets_t = createTensor(ctx, Shape{b * t}, ki32); - op = createKernel(ctx, {kShaderCrossEntropySoftmaxBackward, 256, kf32}, - Bindings{dlogits_t, dlosses_t, probs_t, targets_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - CrossEntropySoftmaxBackwardParams{ - static_cast(b), - static_cast(t), - static_cast(v), - static_cast(vp) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dlogits_t = ctx.pool.data[op->buffers[0]]; - Tensor& dlosses_t = ctx.pool.data[op->buffers[1]]; - Tensor& probs_t = ctx.pool.data[op->buffers[2]]; - Tensor& targets_t = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, dlogits, dlogits_t); - toGPU(ctx, dlosses, dlosses_t); - toGPU(ctx, probs, probs_t); - toGPU(ctx, targets, targets_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dlogits_t, dlogits, b * t * vp * sizeof(float)); -} diff --git a/experimental/kernels/ops.hpp b/experimental/kernels/ops.hpp deleted file mode 100644 index 2514329..0000000 --- a/experimental/kernels/ops.hpp +++ /dev/null @@ -1,108 +0,0 @@ -#ifndef OPS_H -#define OPS_H - -#include "gpu.hpp" - -using namespace gpu; - -#ifdef __cplusplus -extern "C" { -#endif - -#define VOCAB_SIZE 50257 - -// See https://github.com/google/dawn/blob/a8fbe981a86cb59536e2de423d2013a82d9b54a0/src/dawn/native/Limits.cpp -#define LIMITS_BUFFER_SIZE_1GB { \ - .nextInChain = nullptr, \ - .limits = { \ - .maxTextureDimension1D=8192, \ - .maxTextureDimension2D=8192, \ - .maxTextureDimension3D=2048, \ - .maxTextureArrayLayers=256, \ - .maxBindGroups=4, \ - .maxBindGroupsPlusVertexBuffers=24, \ - .maxBindingsPerBindGroup=1000, \ - .maxDynamicUniformBuffersPerPipelineLayout=8, \ - .maxDynamicStorageBuffersPerPipelineLayout=4, \ - .maxSampledTexturesPerShaderStage=16, \ - .maxSamplersPerShaderStage=16, \ - .maxStorageBuffersPerShaderStage=8, \ - .maxStorageTexturesPerShaderStage=4, \ - .maxUniformBuffersPerShaderStage=12, \ - .maxUniformBufferBindingSize=65536, \ - .maxStorageBufferBindingSize=1073741824, \ - .minUniformBufferOffsetAlignment=256, \ - .minStorageBufferOffsetAlignment=256, \ - .maxVertexBuffers=8, \ - .maxBufferSize=0x80000000, \ - .maxVertexAttributes=16, \ - .maxVertexBufferArrayStride=2048, \ - .maxInterStageShaderComponents=64, \ - .maxInterStageShaderVariables=16, \ - .maxColorAttachments=8, \ - .maxColorAttachmentBytesPerSample=32, \ - .maxComputeWorkgroupStorageSize=16384, \ - .maxComputeInvocationsPerWorkgroup=256, \ - .maxComputeWorkgroupSizeX=256, \ - .maxComputeWorkgroupSizeY=256, \ - .maxComputeWorkgroupSizeZ=64, \ - .maxComputeWorkgroupsPerDimension=65535 \ - } \ - } - - -void encoder_forward(Context& ctx, float* out, - int* inp, float* wte, float* wpe, - int B, int T, int C); - -void encoder_backward(Context& ctx, float* dwte, float* dwpe, - float* dout, int* inp, - int B, int T, int C); - -void layernorm_forward(Context& ctx, float* out, float* mean, float* rstd, - float* inp, float* weight, float* bias, - int B, int T, int C); - -void layernorm_backward(Context& ctx, float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* mean, float* rstd, - int B, int T, int C); - -void matmul_forward(Context& ctx, float* out, - const float* inp, const float* weight, const float* bias, - int B, int T, int C, int OC); - -void matmul_backward(Context& ctx, float* dinp, float* dweight, float* dbias, - const float* dout, const float* inp, const float* weight, - int B, int T, int C, int OC); - -void attention_forward(Context& ctx, float* out, float* preatt, float* att, - float* inp, - int B, int T, int C, int NH); - -void attention_backward(Context& ctx, float* dinp, float* dpreatt, float* datt, - float* dout, float* inp, float* att, - int B, int T, int C, int NH); - -void gelu_forward(Context& ctx, float* out, float* inp, int N); - -void gelu_backward(Context& ctx, float* dinp, float* inp, float* dout, int N); - -void residual_forward(Context& ctx, float* out, float* inp1, float* inp2, int N); - -void residual_backward(Context& ctx, float* dinp1, float* dinp2, float* dout, int N); - -void softmax_forward(Context& ctx, float* probs, float* logits, int B, int T, int V, int Vp); - -void crossentropy_forward(Context& ctx, float* losses, - float* probs, int* targets, - int B, int T, int Vp); - -void crossentropy_softmax_backward(Context& ctx, float* dlogits, - float* dlosses, float* probs, int* targets, - int B, int T, int V, int Vp); - -#ifdef __cplusplus -} -#endif - -#endif // OPS_H diff --git a/experimental/kernels/ops_aot.cpp b/experimental/kernels/ops_aot.cpp deleted file mode 100644 index f4ce9c0..0000000 --- a/experimental/kernels/ops_aot.cpp +++ /dev/null @@ -1,356 +0,0 @@ -#include "gpu.hpp" -#include -#include -#include -#include - -#include "kernels.h" -#include "ops_aot.hpp" -#include "experimental/wgsl.h" // loopUnrolling - -using namespace gpu; - -Kernel encoder_forward(Context& ctx, Tensor& out, - Tensor& inp, Tensor& wte, Tensor& wpe, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long v = VOCAB_SIZE; - struct EncoderParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - return createKernel(ctx, {kShaderEncoder, 256, kf32}, - Bindings{inp, wte, wpe, out}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - EncoderParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }); -} - -Kernel encoder_backward(Context& ctx, Tensor& dwte, Tensor& dwpe, - Tensor& dout, Tensor& inp, - int B, int T, int C) { - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long v = VOCAB_SIZE; - struct EncoderParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - return createKernel(ctx, {kShaderEncoderBackward, 256, kf32}, - Bindings{dwte, dwpe, dout, inp}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - EncoderParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }); -} - -Kernel layernorm_forward(Context& ctx, Tensor& out, Tensor& mean, Tensor& rstd, - Tensor& inp, Tensor& weight, Tensor& bias, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - struct LayerNormParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - return createKernel(ctx, {kShaderLayerNorm, 256, kf32}, - Bindings{inp, weight, bias, out, mean, rstd}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - LayerNormParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }); -} - -Kernel layernorm_backward(Context& ctx, Tensor& dinp, Tensor& dweight, Tensor& dbias, - Tensor& dout, Tensor& inp, Tensor& weight, Tensor& mean, Tensor& rstd, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - struct LayerNormParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - return createKernel(ctx, {kShaderLayerNormBackward, 256, kf32}, - Bindings{dinp, dweight, dbias, dout, inp, weight, mean, rstd}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - LayerNormParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }); -} - -struct DurationTime { - std::chrono::high_resolution_clock::time_point start; - std::chrono::high_resolution_clock::time_point end; - std::chrono::microseconds duration; - std::string src; - bool verbose; - - inline DurationTime(const std::string& src, bool verbose = true) { - this->src = src; - this->verbose = verbose; - start = std::chrono::high_resolution_clock::now(); - } - - inline ~DurationTime() { - end = std::chrono::high_resolution_clock::now(); - duration = std::chrono::duration_cast(end - start); - if (this->verbose) { - printf("Duration(%s): %.1f microseconds\n", src.c_str(), static_cast(duration.count())); - } - } -}; - - -Kernel matmul_forward(Context& ctx, Tensor& out, - const Tensor& inp, const Tensor& weight, const Tensor& bias, - int B, int T, int C, int OC){ - bool verbose = false; - DurationTime duration("matmul_forward_gpu", verbose); - struct MatmulParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t OC; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long oc = static_cast(OC); - setLogLevel(kError); - - constexpr size_t BT = 64; - constexpr size_t BC = 8; - constexpr size_t BOC = 64; - constexpr size_t TT = BT / BC; - constexpr size_t TOC = BOC / BC; - size_t num_threads = BT * BOC / (TT * TOC); - Shape wgSize = {num_threads, 1, 1}; - Shape nWorkgroups = {b, cdiv(T, BT), cdiv(OC, BOC)}; - - std::string kShaderMatmul2DTiling_(kShaderMatmul2DTiling); - std::string kShaderMatmul2D(loopUnrolling( - replaceAll(kShaderMatmul2DTiling_, - {{"{{precision}}", toString(kf32)}, - {"{{BT}}", toString(BT)}, - {"{{BC}}", toString(BC)}, - {"{{BOC}}", toString(BOC)}, - {"{{TT}}", toString(TT)}, - {"{{TOC}}", toString(TOC)}, - {"{{NUM_TILEI}}", toString(BT * BC / num_threads)}, - {"{{NUM_TILEW}}", toString(BOC * BC / num_threads)} - }) - ) - ); - - return createKernel(ctx, {kShaderMatmul2D, wgSize, kf32}, - Bindings{inp, weight, bias, out}, - nWorkgroups, - /* params */ - MatmulParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(oc) - }); -} - -Kernel matmul_backward(Context& ctx, Tensor& dinp, Tensor& dweight, Tensor& dbias, - const Tensor& dout, const Tensor& inp, const Tensor& weight, - int B, int T, int C, int OC){ - struct MatmulParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t OC; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long oc = static_cast(OC); - setLogLevel(kError); - return createKernel(ctx, {kShaderMatmulBackward, 256, kf32}, - Bindings{dinp, dweight, dbias, dout, inp, weight}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - MatmulParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(oc) - }); -} - -Kernel attention_forward(Context& ctx, Tensor& out, Tensor& preatt, Tensor& att, - Tensor& inp, - int B, int T, int C, int NH){ - struct AttentionParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t NH; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long nh = static_cast(NH); - setLogLevel(kError); - return createKernel(ctx, {kShaderAttention, 256, kf32}, - Bindings{inp, preatt, att, out}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - AttentionParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(nh) - }); -} - -Kernel attention_backward(Context& ctx, Tensor& dinp, Tensor& dpreatt, Tensor& datt, - Tensor& dout, Tensor& inp, Tensor& att, - int B, int T, int C, int NH){ - struct AttentionParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t NH; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long nh = static_cast(NH); - setLogLevel(kError); - return createKernel(ctx, {kShaderAttentionBackward, 256, kf32}, - Bindings{dinp, dpreatt, datt, dout, inp, att}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - AttentionParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(nh) - }); -} - -Kernel gelu_forward(Context& ctx, Tensor& out, Tensor& inp, int n) { - unsigned long N = static_cast(n); - setLogLevel(kError); - return createKernel(ctx, {kShaderGelu, 256, kf32}, - Bindings{inp, out}, - /* nWorkgroups */ {cdiv(N, 256), 1, 1}); -} - -Kernel gelu_backward(Context& ctx, Tensor& dinp, Tensor& inp, Tensor& dout, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - return createKernel(ctx, {kShaderGeluBackward, 256, kf32}, - Bindings{inp, dout, dinp}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}); -} - -Kernel residual_forward(Context& ctx, Tensor& out, Tensor& inp1, Tensor& inp2, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - return createKernel(ctx, {kShaderResidual, 256, kf32}, - Bindings{inp1, inp2, out}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}); -} - -Kernel residual_backward(Context& ctx, Tensor& dinp1, Tensor& dinp2, Tensor& dout, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - return createKernel(ctx, {kShaderResidualBackward, 256, kf32}, - Bindings{dout, dinp1, dinp2}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}); -} - -Kernel softmax_forward(Context& ctx, Tensor& probs, Tensor& logits, int B, int T, int V, int Vp) { - struct SoftmaxParam { - uint32_t N; - uint32_t C; - uint32_t Cp; - }; - uint32_t b = static_cast(B); - uint32_t t = static_cast(T); - uint32_t c = static_cast(V); - uint32_t cp = static_cast(Vp); - assert( (B*T) % 256 == 0); - return createKernel( - ctx, {kShaderSoftmax1, 256, kf32}, Bindings{logits, probs}, - Shape{cdiv(B * T, 256), 1, 1}, SoftmaxParam{b * t, c, cp}); -} - -Kernel crossentropy_forward(Context& ctx, Tensor& losses, - Tensor& probs, Tensor& targets, - int B, int T, int Vp){ - struct CrossEntropyParams { - uint32_t B; - uint32_t T; - uint32_t VP; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long vp = static_cast(Vp); - setLogLevel(kError); - return createKernel(ctx, {kShaderCrossEntropyForward, 256, kf32}, - Bindings{losses, probs, targets}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - CrossEntropyParams{ - static_cast(b), - static_cast(t), - static_cast(vp) - }); -} - -Kernel crossentropy_softmax_backward(Context& ctx, Tensor& dlogits, - Tensor& dlosses, Tensor& probs, Tensor& targets, - int B, int T, int V, int Vp){ - struct CrossEntropySoftmaxBackwardParams { - uint32_t B; - uint32_t T; - uint32_t V; - uint32_t VP; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long v = static_cast(V); - unsigned long vp = static_cast(Vp); - setLogLevel(kError); - return createKernel(ctx, {kShaderCrossEntropySoftmaxBackward, 256, kf32}, - Bindings{dlogits, dlosses, probs, targets}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - CrossEntropySoftmaxBackwardParams{ - static_cast(b), - static_cast(t), - static_cast(v), - static_cast(vp) - }); -} diff --git a/experimental/kernels/ops_aot.hpp b/experimental/kernels/ops_aot.hpp deleted file mode 100644 index 5db9ff7..0000000 --- a/experimental/kernels/ops_aot.hpp +++ /dev/null @@ -1,100 +0,0 @@ -#ifndef OPS_H -#define OPS_H - -#include "gpu.hpp" - -using namespace gpu; - -#define VOCAB_SIZE 50257 - -// See https://github.com/google/dawn/blob/a8fbe981a86cb59536e2de423d2013a82d9b54a0/src/dawn/native/Limits.cpp -#define LIMITS_BUFFER_SIZE_1GB { \ - .nextInChain = nullptr, \ - .limits = { \ - .maxTextureDimension1D=8192, \ - .maxTextureDimension2D=8192, \ - .maxTextureDimension3D=2048, \ - .maxTextureArrayLayers=256, \ - .maxBindGroups=4, \ - .maxBindGroupsPlusVertexBuffers=24, \ - .maxBindingsPerBindGroup=1000, \ - .maxDynamicUniformBuffersPerPipelineLayout=8, \ - .maxDynamicStorageBuffersPerPipelineLayout=4, \ - .maxSampledTexturesPerShaderStage=16, \ - .maxSamplersPerShaderStage=16, \ - .maxStorageBuffersPerShaderStage=8, \ - .maxStorageTexturesPerShaderStage=4, \ - .maxUniformBuffersPerShaderStage=12, \ - .maxUniformBufferBindingSize=65536, \ - .maxStorageBufferBindingSize=1073741824, \ - .minUniformBufferOffsetAlignment=256, \ - .minStorageBufferOffsetAlignment=256, \ - .maxVertexBuffers=8, \ - .maxBufferSize=0x80000000, \ - .maxVertexAttributes=16, \ - .maxVertexBufferArrayStride=2048, \ - .maxInterStageShaderComponents=64, \ - .maxInterStageShaderVariables=16, \ - .maxColorAttachments=8, \ - .maxColorAttachmentBytesPerSample=32, \ - .maxComputeWorkgroupStorageSize=16384, \ - .maxComputeInvocationsPerWorkgroup=256, \ - .maxComputeWorkgroupSizeX=256, \ - .maxComputeWorkgroupSizeY=256, \ - .maxComputeWorkgroupSizeZ=64, \ - .maxComputeWorkgroupsPerDimension=65535 \ - } \ - } - - -Kernel encoder_forward(Context& ctx, Tensor& out, - Tensor& inp, Tensor& wte, Tensor& wpe, - int B, int T, int C); - -Kernel encoder_backward(Context& ctx, Tensor& dwte, Tensor& dwpe, - Tensor& dout, Tensor& inp, - int B, int T, int C); - -Kernel layernorm_forward(Context& ctx, Tensor& out, Tensor& mean, Tensor& rstd, - Tensor& inp, Tensor& weight, Tensor& bias, - int B, int T, int C); - -Kernel layernorm_backward(Context& ctx, Tensor& dinp, Tensor& dweight, Tensor& dbias, - Tensor& dout, Tensor& inp, Tensor& weight, Tensor& mean, Tensor& rstd, - int B, int T, int C); - -Kernel matmul_forward(Context& ctx, Tensor& out, - const Tensor& inp, const Tensor& weight, const Tensor& bias, - int B, int T, int C, int OC); - -Kernel matmul_backward(Context& ctx, Tensor& dinp, Tensor& dweight, Tensor& dbias, - const Tensor& dout, const Tensor& inp, const Tensor& weight, - int B, int T, int C, int OC); - -Kernel attention_forward(Context& ctx, Tensor& out, Tensor& preatt, Tensor& att, - Tensor& inp, - int B, int T, int C, int NH); - -Kernel attention_backward(Context& ctx, Tensor& dinp, Tensor& dpreatt, Tensor& datt, - Tensor& dout, Tensor& inp, Tensor& att, - int B, int T, int C, int NH); - -Kernel gelu_forward(Context& ctx, Tensor& out, Tensor& inp, int N); - -Kernel gelu_backward(Context& ctx, Tensor& dinp, Tensor& inp, Tensor& dout, int N); - -Kernel residual_forward(Context& ctx, Tensor& out, Tensor& inp1, Tensor& inp2, int N); - -Kernel residual_backward(Context& ctx, Tensor& dinp1, Tensor& dinp2, Tensor& dout, int N); - -Kernel softmax_forward(Context& ctx, Tensor& probs, Tensor& logits, int B, int T, int V, int Vp); - -Kernel crossentropy_forward(Context& ctx, Tensor& losses, - Tensor& probs, Tensor& targets, - int B, int T, int Vp); - -Kernel crossentropy_softmax_backward(Context& ctx, Tensor& dlogits, - Tensor& dlosses, Tensor& probs, Tensor& targets, - int B, int T, int V, int Vp); - -#endif // OPS_H diff --git a/experimental/kernels/reduce.cpp b/experimental/kernels/reduce.cpp deleted file mode 100644 index 46460df..0000000 --- a/experimental/kernels/reduce.cpp +++ /dev/null @@ -1,514 +0,0 @@ -#include "gpu.hpp" -#include -#include -#include -#include -#include -#include "utils/array_utils.hpp" // show, isclose, randn, randint -#include "kernels.h" - -using namespace gpu; - -#define LIMITS { \ - .nextInChain = nullptr, \ - .limits = { \ - .maxTextureDimension1D=8192, \ - .maxTextureDimension2D=8192, \ - .maxTextureDimension3D=2048, \ - .maxTextureArrayLayers=256, \ - .maxBindGroups=4, \ - .maxBindGroupsPlusVertexBuffers=24, \ - .maxBindingsPerBindGroup=1000, \ - .maxDynamicUniformBuffersPerPipelineLayout=8, \ - .maxDynamicStorageBuffersPerPipelineLayout=4, \ - .maxSampledTexturesPerShaderStage=16, \ - .maxSamplersPerShaderStage=16, \ - .maxStorageBuffersPerShaderStage=8, \ - .maxStorageTexturesPerShaderStage=4, \ - .maxUniformBuffersPerShaderStage=12, \ - .maxUniformBufferBindingSize=65536, \ - .maxStorageBufferBindingSize=1073741824, \ - .minUniformBufferOffsetAlignment=256, \ - .minStorageBufferOffsetAlignment=256, \ - .maxVertexBuffers=8, \ - .maxBufferSize=0x80000000, \ - .maxVertexAttributes=16, \ - .maxVertexBufferArrayStride=2048, \ - .maxInterStageShaderComponents=64, \ - .maxInterStageShaderVariables=16, \ - .maxColorAttachments=8, \ - .maxColorAttachmentBytesPerSample=32, \ - .maxComputeWorkgroupStorageSize=16384, \ - .maxComputeInvocationsPerWorkgroup=1024, \ - .maxComputeWorkgroupSizeX=1024, \ - .maxComputeWorkgroupSizeY=1024, \ - .maxComputeWorkgroupSizeZ=64, \ - .maxComputeWorkgroupsPerDimension=65535 \ - } \ - } - - -struct DurationTime { - std::chrono::high_resolution_clock::time_point start; - std::chrono::high_resolution_clock::time_point end; - std::chrono::microseconds duration; - std::string src; - bool verbose; - int num; - - inline DurationTime(const std::string& src, bool verbose = true, int num = 1) { - this->src = src; - this->verbose = verbose; - this->num = num; - start = std::chrono::high_resolution_clock::now(); - } - - inline ~DurationTime() { - end = std::chrono::high_resolution_clock::now(); - duration = std::chrono::duration_cast(end - start); - if (this->verbose) { - printf("Duration(%s): %.1f microseconds\n", src.c_str(), static_cast(duration.count()) / static_cast(num)); - } - } -}; - -static const char *kSumVersion1 = R"( -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -var buffer: array<{{precision}}, 1024>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(local_invocation_id) localID : vec3, - @builtin(workgroup_id) groupid : vec3, - @builtin(num_workgroups) numGroups : vec3) { - let blockSize3d: vec3 = vec3({{workgroupSize}}); - let blockSize: u32 = blockSize3d.x; - let threadId: u32 = localID.x; - let blockId: u32 = groupid.x + groupid.y * numGroups.x; - let blockStart = blockId * blockSize * 2 + threadId; - - buffer[threadId] = inp[blockStart] + inp[blockStart + blockSize]; - workgroupBarrier(); - - for (var stride: u32 = blockSize / 2; stride > 0; stride /= 2) { - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - } - - if (threadId == 0) { - out[blockId] = buffer[0]; - } -} -)"; - -static const char *kSumVersion2 = R"( -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -var buffer: array<{{precision}}, 1024>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) globalID : vec3, - @builtin(local_invocation_id) localID : vec3, - @builtin(workgroup_id) groupid : vec3, - @builtin(num_workgroups) numGroups : vec3) { - let blockSize3d: vec3 = vec3({{workgroupSize}}); - let blockSize: u32 = blockSize3d.x; - let threadId: u32 = localID.x; - let blockId: u32 = groupid.x + groupid.y * numGroups.x; - let n: u32 = arrayLength(&inp); - let blockStart = blockId * blockSize * 2 + threadId; - - buffer[threadId] = inp[blockStart] + inp[blockStart + blockSize]; - workgroupBarrier(); - var stride: u32 = blockSize / 2; - - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/4 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/8 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/16 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/32 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/64 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/128 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/256 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/512 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - - stride /= 2; // 1/1024 - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - - if (threadId == 0) { - out[blockId] = buffer[0]; - } -} -)"; - -static const char *kSum2d = R"( -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -@group(0) @binding(2) var params : Params; -struct Params { - N: u32, - C: u32, -}; -var buffer: array<{{precision}}, 1024>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(local_invocation_id) localID : vec3, - @builtin(workgroup_id) groupid : vec3, - @builtin(num_workgroups) numGroups : vec3) { - let N : u32 = params.N; - let C : u32 = params.C; - let blockSize3d: vec3 = vec3({{workgroupSize}}); - let blockSize: u32 = blockSize3d.x; - let threadId: u32 = localID.x; - let blockId: u32 = groupid.x + groupid.y * numGroups.x; - - for (var i: u32 = 0; i= N) { - } else if(blockStart + blockSize >= N) { - buffer[threadId] = inp[blockStart * C + i]; - } else { - buffer[threadId] = inp[blockStart * C + i] + inp[(blockStart + blockSize) * C + i]; - } - workgroupBarrier(); - - for (var stride: u32 = blockSize / 2; stride > 0; stride /= 2) { - if (threadId < stride) { - buffer[threadId] += buffer[threadId + stride]; - } - workgroupBarrier(); - } - - if (threadId == 0) { - out[blockId * C + i] = buffer[0]; - } - workgroupBarrier(); - } -} -)"; - -float sum_cpu(const float* data, size_t size) { - float result = 0; - for (size_t i = 0; i < size; ++i) { - result += data[i]; - } - return result; -} - -void sum_cpu_2d(const float* data, float* out, size_t size0, size_t size1) { - float result = 0; - for (size_t j = 0; j < size1; ++j) { - out[j] = 0; - } - for (size_t i = 0; i < size0; ++i) { - for (size_t j = 0; j < size1; ++j) { - out[j] += data[(i * size1) + j]; - } - } -} - -Kernel createSumKernel(Context& ctx, Tensor& input, Tensor& output, size_t size, uint32_t num_threads = 1024) { - uint32_t num_blocks = ((size + num_threads -1) / num_threads); - uint32_t size_x = 32768u < num_blocks ? 32768u : num_blocks; - uint32_t size_y = size_x == 32768u ? num_blocks / 32768u : 1; - size_x /= 2; - size_x = size_x < 1 ? 1 : size_x; - // print size_x, size_y - printf("size_x: %u, size_y: %u, num_blocks: %u\n", size_x, size_y, num_blocks); - return createKernel(ctx, {kSum, num_threads, kf32}, Bindings{input, output}, {size_x, size_y, 1}); -} - -Kernel createSumKernel2d(Context& ctx, Tensor& input, Tensor& output, size_t size0, size_t size1, uint32_t num_threads = 1024) { - struct Params { - uint32_t N; - uint32_t C; - }; - uint32_t num_blocks = ((size0 + num_threads -1) / num_threads); - uint32_t size_x = num_blocks; - uint32_t size_y = size1; - size_x /= 2; - size_x = size_x < 1 ? 1 : size_x; - printf("size_x: %u, size_y: %u, num_blocks: %u\n", size_x, size_y, num_blocks); - return createKernel(ctx, - {kSum2d, num_threads, kf32}, - Bindings{input, output}, - {size_x, size_y, 1}, - Params{ - static_cast(size0), - static_cast(size1), - }); -} - -struct SumKernel { - std::vector outputs; - std::vector ops; - SumKernel(Context& ctx, size_t size, uint32_t num_threads = 1024) { - int input_size = size; - unsigned long output_size = size; - outputs.push_back(createTensor(ctx, Shape{std::max(size, static_cast(num_threads*2))}, kf32)); - for(int j=0;output_size>1;j++){ - output_size = (output_size + (num_threads * 2) - 1) / (num_threads * 2); - outputs.push_back(createTensor(ctx, Shape{std::max(output_size, static_cast(num_threads*2))}, kf32)); - ops.push_back(createSumKernel(ctx, outputs[j], outputs[j+1], input_size, num_threads)); - input_size = output_size; - } - } - void dispatchKernel(Context& ctx) { - for(int i=0;i promise; - std::future future = promise.get_future(); - gpu::dispatchKernel(ctx, ops[i], promise); - wait(ctx, future); - resetCommandBuffer(ctx.device, ops[i]); - } - } - void toGPU(Context& ctx, const float* data, size_t size) { - gpu::toGPU(ctx, data, outputs[0], size); - } - void toCPU(Context& ctx, float* data, size_t size) { - gpu::toCPU(ctx, outputs[outputs.size()-1], data, size); - } -}; - -struct SumKernel2d { - std::vector outputs; - std::vector ops; - bool debug; - SumKernel2d(Context& ctx, size_t size0, size_t size1, uint32_t num_threads = 1024) { - debug = false; - int input_size = size0; - unsigned long output_size = size0; - outputs.push_back(createTensor(ctx, Shape{std::max(size0, static_cast(num_threads*2)),size1}, kf32)); - for(int j=0;output_size>1;j++){ - output_size = (output_size + (num_threads * 2) - 1) / (num_threads * 2); - if (debug) - printf("size0: %zu, num_threads: %d, output_size: %lu\n", size0, num_threads, output_size); - outputs.push_back(createTensor(ctx, Shape{std::max(output_size, static_cast(num_threads*2)), size1}, kf32)); - ops.push_back(createSumKernel2d(ctx, outputs[j], outputs[j+1], input_size, size1, num_threads)); - input_size = output_size; - } - if (debug) - printf("ops.size(): %zu\n", ops.size()); - } - void dispatchKernel(Context& ctx) { - for(int i=0;i promise; - std::future future = promise.get_future(); - gpu::dispatchKernel(ctx, ops[i], promise); - wait(ctx, future); - resetCommandBuffer(ctx.device, ops[i]); - } - if (debug) { - std::unique_ptr buffer = std::make_unique(8); - for(int i=0;i inputArr = std::make_unique(M * N); - std::unique_ptr buffer = std::make_unique(BUF_SIZE); - std::mt19937 gen(314159); - printf("Initializing %zu values\n", M*N); - randn(inputArr.get(), M*N, gen); - // for(int i=0;i= 1e-0f) { - printf("Error: diff = %.6f\n", diff); - } else { - printf("Success: diff = %.6f\n", diff); - } - - printf("Computed %zu values of kSum(x)\n\n", M*N); - return 0; -} - -int main_2d(int argc, char **argv) { - static constexpr size_t M = 4096; - static constexpr size_t N = 4096; - std::unique_ptr inputArr = std::make_unique(M * N); - std::unique_ptr outputCpuArr = std::make_unique(N); - std::unique_ptr outputGpuArr = std::make_unique(N); - std::mt19937 gen(314159); - printf("Initializing %zu values\n", M*N); - randn(inputArr.get(), M*N, gen); - for(int i=0;i= 1e-0f) { - printf("Error: diff = %.6f\n", diff); - } else { - printf("Success: diff = %.6f\n", diff); - } - - return 0; -} - -int main(int argc, char **argv) { - printf("================================\n"); - printf("Start testing reduce-1d\n"); - main_1d(argc,argv); - printf("================================\n"); - printf("Start testing reduce-2d\n"); - main_2d(argc,argv); - return 0; -} diff --git a/experimental/kernels/term.html b/experimental/kernels/term.html deleted file mode 100644 index 0774ab8..0000000 --- a/experimental/kernels/term.html +++ /dev/null @@ -1,54 +0,0 @@ - - - - - gpu.cpp - - - - - - -
- - - - - {{{ SCRIPT }}} - - - diff --git a/experimental/kernels/unittest_llmc/unittest_kernels.cpp b/experimental/kernels/unittest_llmc/unittest_kernels.cpp deleted file mode 100644 index d037eac..0000000 --- a/experimental/kernels/unittest_llmc/unittest_kernels.cpp +++ /dev/null @@ -1,932 +0,0 @@ -#include "gpu.hpp" -#include -#include -#include -#include - -#include "kernels.h" -#include "unittest_llmc/unittest_kernels.h" -#include "experimental/wgsl.h" // loopUnrolling - -using namespace gpu; // createContext, createTensor, createKernel, - // createShader, dispatchKernel, wait, toCPU - // Tensor, Kernel, Context, Shape, kf32 - -#define VOCAB_SIZE 50257 - -// See https://github.com/google/dawn/blob/a8fbe981a86cb59536e2de423d2013a82d9b54a0/src/dawn/native/Limits.cpp -#define LIMITS_BUFFER_SIZE_1GB { \ - .nextInChain = nullptr, \ - .limits = { \ - .maxTextureDimension1D=8192, \ - .maxTextureDimension2D=8192, \ - .maxTextureDimension3D=2048, \ - .maxTextureArrayLayers=256, \ - .maxBindGroups=4, \ - .maxBindGroupsPlusVertexBuffers=24, \ - .maxBindingsPerBindGroup=1000, \ - .maxDynamicUniformBuffersPerPipelineLayout=8, \ - .maxDynamicStorageBuffersPerPipelineLayout=4, \ - .maxSampledTexturesPerShaderStage=16, \ - .maxSamplersPerShaderStage=16, \ - .maxStorageBuffersPerShaderStage=8, \ - .maxStorageTexturesPerShaderStage=4, \ - .maxUniformBuffersPerShaderStage=12, \ - .maxUniformBufferBindingSize=65536, \ - .maxStorageBufferBindingSize=1073741824, \ - .minUniformBufferOffsetAlignment=256, \ - .minStorageBufferOffsetAlignment=256, \ - .maxVertexBuffers=8, \ - .maxBufferSize=0x80000000, \ - .maxVertexAttributes=16, \ - .maxVertexBufferArrayStride=2048, \ - .maxInterStageShaderComponents=64, \ - .maxInterStageShaderVariables=16, \ - .maxColorAttachments=8, \ - .maxColorAttachmentBytesPerSample=32, \ - .maxComputeWorkgroupStorageSize=16384, \ - .maxComputeInvocationsPerWorkgroup=256, \ - .maxComputeWorkgroupSizeX=256, \ - .maxComputeWorkgroupSizeY=256, \ - .maxComputeWorkgroupSizeZ=64, \ - .maxComputeWorkgroupsPerDimension=65535 \ - } \ - } - -struct DurationTime { - std::chrono::high_resolution_clock::time_point start; - std::chrono::high_resolution_clock::time_point end; - std::chrono::microseconds duration; - std::string src; - bool verbose; - - inline DurationTime(const std::string& src, bool verbose = true) { - this->src = src; - this->verbose = verbose; - start = std::chrono::high_resolution_clock::now(); - } - - inline ~DurationTime() { - end = std::chrono::high_resolution_clock::now(); - duration = std::chrono::duration_cast(end - start); - if (this->verbose) { - printf("Duration(%s): %.1f microseconds\n", src.c_str(), static_cast(duration.count())); - } - } -}; - -static WGPURequiredLimits requiredLimits = LIMITS_BUFFER_SIZE_1GB; -static Context ctx = createContext({},{},{ - .requiredLimits = &requiredLimits - }); - -void ENCODER_FORWARD_GPU(float* out, - int* inp, float* wte, float* wpe, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long v = VOCAB_SIZE; - struct EncoderParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "ENCODER_FORWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor input = createTensor(ctx, Shape{b * t}, ki32); - Tensor wte_t = createTensor(ctx, Shape{v, c}, kf32); - Tensor wpe_t = createTensor(ctx, Shape{t, c}, kf32); - Tensor output = createTensor(ctx, Shape{b * t * c}, kf32); - op = createKernel(ctx, {kShaderEncoder, 256, kf32}, - Bindings{input, wte_t, wpe_t, output}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - EncoderParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& input = ctx.pool.data[op->buffers[0]]; - Tensor& wte_t = ctx.pool.data[op->buffers[1]]; - Tensor& wpe_t = ctx.pool.data[op->buffers[2]]; - Tensor& output = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, inp, input); - toGPU(ctx, wte, wte_t); - toGPU(ctx, wpe, wpe_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, out, b * t * c * sizeof(float)); -} - -void ENCODER_BACKWARD_GPU(float* dwte, float* dwpe, - float* dout, int* inp, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long v = VOCAB_SIZE; - struct EncoderParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "ENCODER_BACKWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dwte_t = createTensor(ctx, Shape{v, c}, kf32); - Tensor dwpe_t = createTensor(ctx, Shape{t, c}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor input = createTensor(ctx, Shape{b * t}, ki32); - op = createKernel(ctx, {kShaderEncoderBackward, 256, kf32}, - Bindings{dwte_t, dwpe_t, dout_t, input}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - EncoderParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dwte_t = ctx.pool.data[op->buffers[0]]; - Tensor& dwpe_t = ctx.pool.data[op->buffers[1]]; - Tensor& dout_t = ctx.pool.data[op->buffers[2]]; - Tensor& input = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, dwte, dwte_t); - toGPU(ctx, dwpe, dwpe_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, input); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dwte_t, dwte, v * c * sizeof(float)); - toCPU(ctx, dwpe_t, dwpe, t * c * sizeof(float)); -} - -void LAYERNORM_FORWARD_GPU(float* out, float* mean, float* rstd, - float* inp, float* weight, float* bias, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - struct LayerNormParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "LAYERNORM_FORWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_t = createTensor(ctx, Shape{c}, kf32); - Tensor bias_t = createTensor(ctx, Shape{c}, kf32); - Tensor out_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor mean_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor rstd_t = createTensor(ctx, Shape{b * t}, kf32); - op = createKernel(ctx, {kShaderLayerNorm, 256, kf32}, - Bindings{inp_t, weight_t, bias_t, out_t, mean_t, rstd_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - LayerNormParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_t = ctx.pool.data[op->buffers[0]]; - Tensor& weight_t = ctx.pool.data[op->buffers[1]]; - Tensor& bias_t = ctx.pool.data[op->buffers[2]]; - Tensor& out_t = ctx.pool.data[op->buffers[3]]; - Tensor& mean_t = ctx.pool.data[op->buffers[4]]; - Tensor& rstd_t = ctx.pool.data[op->buffers[5]]; - - toGPU(ctx, inp, inp_t); - toGPU(ctx, weight, weight_t); - toGPU(ctx, bias, bias_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, out_t, out, b * t * c * sizeof(float)); - toCPU(ctx, mean_t, mean, b * t * sizeof(float)); - toCPU(ctx, rstd_t, rstd, b * t * sizeof(float)); -} - -void LAYERNORM_BACKWARD_GPU(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* mean, float* rstd, - int B, int T, int C){ - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - struct LayerNormParams { - uint32_t B; - uint32_t T; - uint32_t C; - }; - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "LAYERNORM_BACKWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dinp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor dweight_t = createTensor(ctx, Shape{c}, kf32); - Tensor dbias_t = createTensor(ctx, Shape{c}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_t = createTensor(ctx, Shape{c}, kf32); - Tensor mean_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor rstd_t = createTensor(ctx, Shape{b * t}, kf32); - op = createKernel(ctx, {kShaderLayerNormBackward, 256, kf32}, - Bindings{dinp_t, dweight_t, dbias_t, dout_t, inp_t, weight_t, mean_t, rstd_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - LayerNormParams{ - static_cast(b), - static_cast(t), - static_cast(c) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dinp_t = ctx.pool.data[op->buffers[0]]; - Tensor& dweight_t = ctx.pool.data[op->buffers[1]]; - Tensor& dbias_t = ctx.pool.data[op->buffers[2]]; - Tensor& dout_t = ctx.pool.data[op->buffers[3]]; - Tensor& inp_t = ctx.pool.data[op->buffers[4]]; - Tensor& weight_t = ctx.pool.data[op->buffers[5]]; - Tensor& mean_t = ctx.pool.data[op->buffers[6]]; - Tensor& rstd_t = ctx.pool.data[op->buffers[7]]; - - toGPU(ctx, dinp, dinp_t); - toGPU(ctx, dweight, dweight_t); - toGPU(ctx, dbias, dbias_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, inp_t); - toGPU(ctx, weight, weight_t); - toGPU(ctx, mean, mean_t); - toGPU(ctx, rstd, rstd_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_t, dinp, b * t * c * sizeof(float)); - toCPU(ctx, dweight_t, dweight, c * sizeof(float)); - toCPU(ctx, dbias_t, dbias, c * sizeof(float)); -} - -void matmul_forward_dummy(float* out, - const float* inp, const float* weight, const float* bias, - int B, int T, int C, int OC); - - -void MATMUL_FORWARD_GPU(float* out, - const float* inp, const float* weight, const float* bias, - int B, int T, int C, int OC){ - int version = 2; - bool verbose = false; - bool debug = false; - float *out_exp; - DurationTime duration("matmul_forward_gpu with preparing a kernel", verbose); - if (verbose) { - printf("matmul forward: B=%d, T=%d, C=%d, OC=%d, bias=%d\n", B, T, C, OC, bias != NULL); - } - if (debug) { - out_exp = new float[B*T*OC]; - { - DurationTime duration("matmul_forward_cpu", verbose); - matmul_forward_dummy(out_exp, inp, weight, bias, B, T, C, OC); - } - } - struct MatmulParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t OC; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long oc = static_cast(OC); - setLogLevel(kError); - - if (version == 2 || version == 1) { - // Generate the key of the cache by arguments. - std::string key = "MATMUL_FORWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(OC); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp_i = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_i = createTensor(ctx, Shape{oc * c}, kf32); - Tensor bias_i = bias == NULL ? createTensor(ctx, Shape{1}, kf32) : createTensor(ctx, Shape{oc}, kf32); - Tensor out_o = createTensor(ctx, Shape{b * t * oc}, kf32); - - if (version == 2) { - constexpr size_t BT = 64; - constexpr size_t BC = 16; - constexpr size_t BOC = 64; - constexpr size_t TT = BT / BC; - constexpr size_t TOC = BOC / BC; - constexpr size_t num_threads = BT * BOC / (TT * TOC); - Shape wgSize = {num_threads, 1, 1}; - - std::string codeString(kShaderMatmul2DTiling); - std::string unrolledCode = loopUnrolling(replaceAll(codeString, {{"{{precision}}", toString(kf32)}, - {"{{BT}}", toString(BT)}, - {"{{BC}}", toString(BC)}, - {"{{BOC}}", toString(BOC)}, - {"{{TT}}", toString(TT)}, - {"{{TOC}}", toString(TOC)}, - {"{{NUM_TILEI}}", toString(BT * BC / num_threads)}, - {"{{NUM_TILEW}}", toString(BOC * BC / num_threads)} - })); - - Shape nWorkgroups = {b, cdiv(T, BT), cdiv(OC, BOC)}; - op = createKernel(ctx, {unrolledCode, wgSize, kf32}, - Bindings{inp_i, weight_i, bias_i, out_o}, - nWorkgroups, - /* params */ - MatmulParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(oc) - }, - nullptr, - key.c_str() - ); - } else { - op = createKernel(ctx, {kShaderMatmul, 256, kf32}, - Bindings{inp_i, weight_i, bias_i, out_o}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - MatmulParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(oc) - }, - nullptr, - key.c_str() - ); - } - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_i = ctx.pool.data[op->buffers[0]]; - Tensor& weight_i = ctx.pool.data[op->buffers[1]]; - Tensor& bias_i = ctx.pool.data[op->buffers[2]]; - Tensor& out_o = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, inp, inp_i); - toGPU(ctx, weight, weight_i); - if (bias != NULL) { - toGPU(ctx, bias, bias_i); - } - - std::promise promise; - std::future future = promise.get_future(); - - { - DurationTime duration("matmul_forward_gpu", verbose); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - } - toCPU(ctx, out_o, out, b * t * oc * sizeof(float)); - } else { - DurationTime duration("matmul_forward_cpu", verbose); - matmul_forward_dummy(out, inp, weight, bias, B, T, C, OC); - } - - if (debug) { // compare out with out_exp. - for (int i = 0; i < B*T*OC; i++) { - if (fabs(out[i] - out_exp[i]) > 1e-2) { - printf("matmul forward: out[%d] = %f, out_exp[%d] = %f\n", i, out[i], i, out_exp[i]); - //Dump the first 4 x 4 elements by table, at first output out, then output out_exp - printf("inp:\n"); - for (int j = 0; j < 4; j++) { - for (int k = 0; k < 4; k++) { - printf("%f ", inp[j * C + k]); - } - printf("\n"); - } - printf("weight:\n"); - for (int j = 0; j < 4; j++) { - for (int k = 0; k < 4; k++) { - printf("%f ", weight[j * OC + k]); - } - printf("\n"); - } - printf("out:\n"); - for (int j = 0; j < 4; j++) { - for (int k = 0; k < 4; k++) { - printf("%f ", out[j * OC + k]); - } - printf("\n"); - } - printf("out_exp:\n"); - for (int j = 0; j < 4; j++) { - for (int k = 0; k < 4; k++) { - printf("%f ", out_exp[j * OC + k]); - } - printf("\n"); - } - exit(1); - } - } - delete[] out_exp; - } -} - -void MATMUL_BACKWARD_GPU(float* dinp, float* dweight, float* dbias, - const float* dout, const float* inp, const float* weight, - int B, int T, int C, int OC){ - struct MatmulParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t OC; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long oc = static_cast(OC); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "MATMUL_BACKWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(OC); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dinp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor dweight_t = createTensor(ctx, Shape{oc * c}, kf32); - Tensor dbias_t = createTensor(ctx, Shape{oc}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * oc}, kf32); - Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor weight_t = createTensor(ctx, Shape{oc * c}, kf32); - op = createKernel(ctx, {kShaderMatmulBackward, 256, kf32}, - Bindings{dinp_t, dweight_t, dbias_t, dout_t, inp_t, weight_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - MatmulParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(oc) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dinp_t = ctx.pool.data[op->buffers[0]]; - Tensor& dweight_t = ctx.pool.data[op->buffers[1]]; - Tensor& dbias_t = ctx.pool.data[op->buffers[2]]; - Tensor& dout_t = ctx.pool.data[op->buffers[3]]; - Tensor& inp_t = ctx.pool.data[op->buffers[4]]; - Tensor& weight_t = ctx.pool.data[op->buffers[5]]; - - toGPU(ctx, dinp, dinp_t); - toGPU(ctx, dweight, dweight_t); - toGPU(ctx, dbias, dbias_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, inp_t); - toGPU(ctx, weight, weight_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_t, dinp, b * t * c * sizeof(float)); - toCPU(ctx, dweight_t, dweight, oc * c * sizeof(float)); - toCPU(ctx, dbias_t, dbias, oc * sizeof(float)); -} - -void ATTENTION_FORWARD_GPU(float* out, float* preatt, float* att, - float* inp, - int B, int T, int C, int NH){ - struct AttentionParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t NH; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long nh = static_cast(NH); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "ATTENTION_FORWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(NH); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32); - Tensor preatt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor att_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor out_t = createTensor(ctx, Shape{b * t * c}, kf32); - op = createKernel(ctx, {kShaderAttention, 256, kf32}, - Bindings{inp_t, preatt_t, att_t, out_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - AttentionParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(nh) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_t = ctx.pool.data[op->buffers[0]]; - Tensor& preatt_t = ctx.pool.data[op->buffers[1]]; - Tensor& att_t = ctx.pool.data[op->buffers[2]]; - Tensor& out_t = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, inp, inp_t); - toGPU(ctx, preatt, preatt_t); - toGPU(ctx, att, att_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, preatt_t, preatt, b * nh * t * t * sizeof(float)); - toCPU(ctx, att_t, att, b * nh * t * t * sizeof(float)); - toCPU(ctx, out_t, out, b * t * c * sizeof(float)); -} - -void ATTENTION_BACKWARD_GPU(float* dinp, float* dpreatt, float* datt, - float* dout, float* inp, float* att, - int B, int T, int C, int NH){ - struct AttentionParams { - uint32_t B; - uint32_t T; - uint32_t C; - uint32_t NH; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long c = static_cast(C); - unsigned long nh = static_cast(NH); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "ATTENTION_BACKWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(C) + "_" + std::to_string(NH); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dinp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32); - Tensor dpreatt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor datt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32); - Tensor inp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32); - Tensor att_t = createTensor(ctx, Shape{b * nh * t * t}, kf32); - op = createKernel(ctx, {kShaderAttentionBackward, 256, kf32}, - Bindings{dinp_t, dpreatt_t, datt_t, dout_t, inp_t, att_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - AttentionParams{ - static_cast(b), - static_cast(t), - static_cast(c), - static_cast(nh) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dinp_t = ctx.pool.data[op->buffers[0]]; - Tensor& dpreatt_t = ctx.pool.data[op->buffers[1]]; - Tensor& datt_t = ctx.pool.data[op->buffers[2]]; - Tensor& dout_t = ctx.pool.data[op->buffers[3]]; - Tensor& inp_t = ctx.pool.data[op->buffers[4]]; - Tensor& att_t = ctx.pool.data[op->buffers[5]]; - - toGPU(ctx, dinp, dinp_t); - toGPU(ctx, dpreatt, dpreatt_t); - toGPU(ctx, datt, datt_t); - toGPU(ctx, dout, dout_t); - toGPU(ctx, inp, inp_t); - toGPU(ctx, att, att_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_t, dinp, b * t * c * 3 * sizeof(float)); - toCPU(ctx, dpreatt_t, dpreatt, b * nh * t * t * sizeof(float)); - toCPU(ctx, datt_t, datt, b * nh * t * t * sizeof(float)); -} - -void GELU_FORWARD_GPU(float* out, float* inp, int n) { - unsigned long N = static_cast(n); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "GELU_FORWARD_GPU_" + std::to_string(n); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor input = createTensor(ctx, Shape{N}, kf32); - Tensor output = createTensor(ctx, Shape{N}, kf32); - op = createKernel(ctx, {kShaderGelu, 256, kf32}, - Bindings{input, output}, - /* nWorkgroups */ {cdiv(N, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& input = ctx.pool.data[op->buffers[0]]; - Tensor& output = ctx.pool.data[op->buffers[1]]; - - toGPU(ctx, inp, input); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, out, N * sizeof(float)); -} - -void GELU_BACKWARD_GPU(float* dinp, float* inp, float* dout, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "GELU_BACKWARD_GPU_" + std::to_string(N); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp_i = createTensor(ctx, Shape{n}, kf32); - Tensor dout_i = createTensor(ctx, Shape{n}, kf32); - Tensor dinp_o = createTensor(ctx, Shape{n}, kf32); - op = createKernel(ctx, {kShaderGeluBackward, 256, kf32}, - Bindings{inp_i, dout_i, dinp_o}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp_i = ctx.pool.data[op->buffers[0]]; - Tensor& dout_i = ctx.pool.data[op->buffers[1]]; - Tensor& dinp_o = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, inp, inp_i); - toGPU(ctx, dout, dout_i); - toGPU(ctx, dinp, dinp_o); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp_o, dinp, n * sizeof(float)); -} - -void RESIDUAL_FORWARD_GPU(float* out, float* inp1, float* inp2, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "RESIDUAL_FORWARD_GPU_" + std::to_string(N); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor inp1_i = createTensor(ctx, Shape{n}, kf32); - Tensor inp2_i = createTensor(ctx, Shape{n}, kf32); - Tensor out_o = createTensor(ctx, Shape{n}, kf32); - op = createKernel(ctx, {kShaderResidual, 256, kf32}, - Bindings{inp1_i, inp2_i, out_o}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& inp1_i = ctx.pool.data[op->buffers[0]]; - Tensor& inp2_i = ctx.pool.data[op->buffers[1]]; - Tensor& out_o = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, inp1, inp1_i); - toGPU(ctx, inp2, inp2_i); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, out_o, out, n * sizeof(float)); -} - -void RESIDUAL_BACKWARD_GPU(float* dinp1, float* dinp2, float* dout, int N){ - unsigned long n = static_cast(N); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "RESIDUAL_BACKWARD_GPU_" + std::to_string(N); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dout_i = createTensor(ctx, Shape{n}, kf32); - Tensor dinp1_o = createTensor(ctx, Shape{n}, kf32); - Tensor dinp2_o = createTensor(ctx, Shape{n}, kf32); - op = createKernel(ctx, {kShaderResidualBackward, 256, kf32}, - Bindings{dout_i, dinp1_o, dinp2_o}, - /* nWorkgroups */ {cdiv(n, 256), 1, 1}, - nullptr, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dout_i = ctx.pool.data[op->buffers[0]]; - Tensor& dinp1_o = ctx.pool.data[op->buffers[1]]; - Tensor& dinp2_o = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, dout, dout_i); - toGPU(ctx, dinp1, dinp1_o); - toGPU(ctx, dinp2, dinp2_o); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dinp1_o, dinp1, n * sizeof(float)); - toCPU(ctx, dinp2_o, dinp2, n * sizeof(float)); -} - -void SOFTMAX_FORWARD_GPU(float* probs, float* logits, int B, int T, int V, int Vp) { - struct SoftmaxParam { - uint32_t N; - uint32_t C; - uint32_t Cp; - }; - uint32_t b = static_cast(B); - uint32_t t = static_cast(T); - uint32_t c = static_cast(V); - uint32_t cp = static_cast(Vp); - - // Generate the key of the cache by arguments. - std::string key = "SOFTMAX_FORWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(V) + "_" + std::to_string(Vp); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor input = createTensor(ctx, {b * t, cp}, kf32); - Tensor output = createTensor(ctx, {b * t, cp}, kf32); - assert( (B*T) % 256 == 0); - op = createKernel( - ctx, {kShaderSoftmax1, 256, kf32}, Bindings{input, output}, - Shape{cdiv(B * T, 256), 1, 1}, SoftmaxParam{b * t, c, cp}, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& input = ctx.pool.data[op->buffers[0]]; - Tensor& output = ctx.pool.data[op->buffers[1]]; - - toGPU(ctx, logits, input); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, probs, sizeof(float)*b*t*cp); -} - -void CROSSENTROPY_FORWARD_GPU(float* losses, - float* probs, int* targets, - int B, int T, int Vp){ - struct CrossEntropyParams { - uint32_t B; - uint32_t T; - uint32_t VP; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long vp = static_cast(Vp); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "CROSSENTROPY_FORWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(Vp); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor losses_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor probs_t = createTensor(ctx, Shape{b * t * vp}, kf32); - Tensor targets_t = createTensor(ctx, Shape{b * t}, ki32); - op = createKernel(ctx, {kShaderCrossEntropyForward, 256, kf32}, - Bindings{losses_t, probs_t, targets_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - CrossEntropyParams{ - static_cast(b), - static_cast(t), - static_cast(vp) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& losses_t = ctx.pool.data[op->buffers[0]]; - Tensor& probs_t = ctx.pool.data[op->buffers[1]]; - Tensor& targets_t = ctx.pool.data[op->buffers[2]]; - - toGPU(ctx, losses, losses_t); - toGPU(ctx, probs, probs_t); - toGPU(ctx, targets, targets_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, losses_t, losses, b * t * sizeof(float)); -} - -void CROSSENTROPY_SOFTMAX_BACKWARD_GPU(float* dlogits, - float* dlosses, float* probs, int* targets, - int B, int T, int V, int Vp){ - struct CrossEntropySoftmaxBackwardParams { - uint32_t B; - uint32_t T; - uint32_t V; - uint32_t VP; - }; - unsigned long b = static_cast(B); - unsigned long t = static_cast(T); - unsigned long v = static_cast(V); - unsigned long vp = static_cast(Vp); - setLogLevel(kError); - - // Generate the key of the cache by arguments. - std::string key = "CROSSENTROPY_SOFTMAX_BACKWARD_GPU_" + std::to_string(B) + "_" + std::to_string(T) + "_" + std::to_string(V) + "_" + std::to_string(Vp); - Kernel op; - if (ctx.kernelPool.data.find(key) == ctx.kernelPool.data.end()) { - Tensor dlogits_t = createTensor(ctx, Shape{b * t * vp}, kf32); - Tensor dlosses_t = createTensor(ctx, Shape{b * t}, kf32); - Tensor probs_t = createTensor(ctx, Shape{b * t * vp}, kf32); - Tensor targets_t = createTensor(ctx, Shape{b * t}, ki32); - op = createKernel(ctx, {kShaderCrossEntropySoftmaxBackward, 256, kf32}, - Bindings{dlogits_t, dlosses_t, probs_t, targets_t}, - /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, - /* params */ - CrossEntropySoftmaxBackwardParams{ - static_cast(b), - static_cast(t), - static_cast(v), - static_cast(vp) - }, - nullptr, - key.c_str()); - } else { - op = ctx.kernelPool.data[key]; - } - Tensor& dlogits_t = ctx.pool.data[op->buffers[0]]; - Tensor& dlosses_t = ctx.pool.data[op->buffers[1]]; - Tensor& probs_t = ctx.pool.data[op->buffers[2]]; - Tensor& targets_t = ctx.pool.data[op->buffers[3]]; - - toGPU(ctx, dlogits, dlogits_t); - toGPU(ctx, dlosses, dlosses_t); - toGPU(ctx, probs, probs_t); - toGPU(ctx, targets, targets_t); - - std::promise promise; - std::future future = promise.get_future(); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, dlogits_t, dlogits, b * t * vp * sizeof(float)); -} diff --git a/experimental/kernels/unittest_llmc/unittest_kernels.h b/experimental/kernels/unittest_llmc/unittest_kernels.h deleted file mode 100644 index 95370d9..0000000 --- a/experimental/kernels/unittest_llmc/unittest_kernels.h +++ /dev/null @@ -1,216 +0,0 @@ -#ifdef __cplusplus -extern "C" { -#endif - -#ifdef METAL_PROFILER -#include "experimental/profiler/metal.hpp" - -#define MAIN main_wrapper -static int main_wrapper(int argc, char *argv[]); - -int main(int argc, char *argv[]) { - startCapture(); - int ret = main_wrapper(argc, argv); - stopCapture(); - return ret; -} - -#else -#define MAIN main -#endif - -// -- USE_GPU_FOR_* are the GPU/CPU switching flags for the kernels in llm.c. -- - -#define USE_GPU_FOR_ENCODER_FORWARD 1 -// #define USE_GPU_FOR_ENCODER_BACKWARD 1 -#define USE_GPU_FOR_LAYERNORM_FORWARD 1 -// -- Note: atomicAdd should be used with i32 or u32 not f32. -// #define USE_GPU_FOR_LAYERNORM_BACKWARD 1 -// -- Note: matmul_forward kernel works, but it is too slow. -#define USE_GPU_FOR_MATMUL_FORWARD 1 -// #define USE_GPU_FOR_MATMUL_BACKWARD 1 -#define USE_GPU_FOR_ATTENTION_FORWARD 1 -// #define USE_GPU_FOR_ATTENTION_BACKWARD 1 -#define USE_GPU_FOR_GELU_FORWARD 1 -#define USE_GPU_FOR_GELU_BACKWARD 1 -#define USE_GPU_FOR_RESIDUAL_FORWARD 1 -#define USE_GPU_FOR_RESIDUAL_BACKWARD 1 -#define USE_GPU_FOR_SOFTMAX_FORWARD 1 -#define USE_GPU_FOR_CROSSENTROPY_FORWARD 1 -#define USE_GPU_FOR_CROSSENTROPY_SOFTMAX_BACKWARD 1 - - -#ifdef USE_GPU_FOR_ENCODER_FORWARD -#define ENCODER_FORWARD_CPU encoder_forward_dummy -#define ENCODER_FORWARD_GPU encoder_forward -#else -#define ENCODER_FORWARD_CPU encoder_forward -#define ENCODER_FORWARD_GPU encoder_forward_dummy -#endif - -#ifdef USE_GPU_FOR_ENCODER_BACKWARD -#define ENCODER_BACKWARD_CPU encoder_backward_dummy -#define ENCODER_BACKWARD_GPU encoder_backward -#else -#define ENCODER_BACKWARD_CPU encoder_backward -#define ENCODER_BACKWARD_GPU encoder_backward_dummy -#endif - -#ifdef USE_GPU_FOR_LAYERNORM_FORWARD -#define LAYERNORM_FORWARD_CPU layernorm_forward_dummy -#define LAYERNORM_FORWARD_GPU layernorm_forward -#else -#define LAYERNORM_FORWARD_CPU layernorm_forward -#define LAYERNORM_FORWARD_GPU layernorm_forward_dummy -#endif - -#ifdef USE_GPU_FOR_LAYERNORM_BACKWARD -#define LAYERNORM_BACKWARD_CPU layernorm_backward_dummy -#define LAYERNORM_BACKWARD_GPU layernorm_backward -#else -#define LAYERNORM_BACKWARD_CPU layernorm_backward -#define LAYERNORM_BACKWARD_GPU layernorm_backward_dummy -#endif - -#ifdef USE_GPU_FOR_MATMUL_FORWARD -#define MATMUL_FORWARD_CPU matmul_forward_dummy -#define MATMUL_FORWARD_GPU matmul_forward -#else -#define MATMUL_FORWARD_CPU matmul_forward -#define MATMUL_FORWARD_GPU matmul_forward_dummy -#endif - -#ifdef USE_GPU_FOR_MATMUL_BACKWARD -#define MATMUL_BACKWARD_CPU matmul_backward_dummy -#define MATMUL_BACKWARD_GPU matmul_backward -#else -#define MATMUL_BACKWARD_CPU matmul_backward -#define MATMUL_BACKWARD_GPU matmul_backward_dummy -#endif - -#ifdef USE_GPU_FOR_ATTENTION_FORWARD -#define ATTENTION_FORWARD_CPU attention_forward_dummy -#define ATTENTION_FORWARD_GPU attention_forward -#else -#define ATTENTION_FORWARD_CPU attention_forward -#define ATTENTION_FORWARD_GPU attention_forward_dummy -#endif - -#ifdef USE_GPU_FOR_ATTENTION_BACKWARD -#define ATTENTION_BACKWARD_CPU attention_backward_dummy -#define ATTENTION_BACKWARD_GPU attention_backward -#else -#define ATTENTION_BACKWARD_CPU attention_backward -#define ATTENTION_BACKWARD_GPU attention_backward_dummy -#endif - -#ifdef USE_GPU_FOR_GELU_FORWARD -#define GELU_FORWARD_CPU gelu_forward_dummy -#define GELU_FORWARD_GPU gelu_forward -#else -#define GELU_FORWARD_CPU gelu_forward -#define GELU_FORWARD_GPU gelu_forward_dummy -#endif - -#ifdef USE_GPU_FOR_GELU_BACKWARD -#define GELU_BACKWARD_CPU gelu_backward_dummy -#define GELU_BACKWARD_GPU gelu_backward -#else -#define GELU_BACKWARD_CPU gelu_backward -#define GELU_BACKWARD_GPU gelu_backward_dummy -#endif - -#ifdef USE_GPU_FOR_RESIDUAL_FORWARD -#define RESIDUAL_FORWARD_CPU residual_forward_dummy -#define RESIDUAL_FORWARD_GPU residual_forward -#else -#define RESIDUAL_FORWARD_CPU residual_forward -#define RESIDUAL_FORWARD_GPU residual_forward_dummy -#endif - -#ifdef USE_GPU_FOR_RESIDUAL_BACKWARD -#define RESIDUAL_BACKWARD_CPU residual_backward_dummy -#define RESIDUAL_BACKWARD_GPU residual_backward -#else -#define RESIDUAL_BACKWARD_CPU residual_backward -#define RESIDUAL_BACKWARD_GPU residual_backward_dummy -#endif - -#ifdef USE_GPU_FOR_SOFTMAX_FORWARD -#define SOFTMAX_FORWARD_CPU softmax_forward_dummy -#define SOFTMAX_FORWARD_GPU softmax_forward -#else -#define SOFTMAX_FORWARD_CPU softmax_forward -#define SOFTMAX_FORWARD_GPU softmax_forward_dummy -#endif - -#ifdef USE_GPU_FOR_CROSSENTROPY_FORWARD -#define CROSSENTROPY_FORWARD_CPU crossentropy_forward_dummy -#define CROSSENTROPY_FORWARD_GPU crossentropy_forward -#else -#define CROSSENTROPY_FORWARD_CPU crossentropy_forward -#define CROSSENTROPY_FORWARD_GPU crossentropy_forward_dummy -#endif - -#ifdef USE_GPU_FOR_CROSSENTROPY_SOFTMAX_BACKWARD -#define CROSSENTROPY_SOFTMAX_BACKWARD_CPU crossentropy_softmax_backward_dummy -#define CROSSENTROPY_SOFTMAX_BACKWARD_GPU crossentropy_softmax_backward -#else -#define CROSSENTROPY_SOFTMAX_BACKWARD_CPU crossentropy_softmax_backward -#define CROSSENTROPY_SOFTMAX_BACKWARD_GPU crossentropy_softmax_backward_dummy -#endif - -void encoder_forward(float* out, - int* inp, float* wte, float* wpe, - int B, int T, int C); - -void encoder_backward(float* dwte, float* dwpe, - float* dout, int* inp, - int B, int T, int C); - -void layernorm_forward(float* out, float* mean, float* rstd, - float* inp, float* weight, float* bias, - int B, int T, int C); - -void layernorm_backward(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* mean, float* rstd, - int B, int T, int C); - -void matmul_forward(float* out, - const float* inp, const float* weight, const float* bias, - int B, int T, int C, int OC); - -void matmul_backward(float* dinp, float* dweight, float* dbias, - const float* dout, const float* inp, const float* weight, - int B, int T, int C, int OC); - -void attention_forward(float* out, float* preatt, float* att, - float* inp, - int B, int T, int C, int NH); - -void attention_backward(float* dinp, float* dpreatt, float* datt, - float* dout, float* inp, float* att, - int B, int T, int C, int NH); - -void gelu_forward(float* out, float* inp, int N); - -void gelu_backward(float* dinp, float* inp, float* dout, int N); - -void residual_forward(float* out, float* inp1, float* inp2, int N); - -void residual_backward(float* dinp1, float* dinp2, float* dout, int N); - -void softmax_forward(float* probs, float* logits, int B, int T, int V, int Vp); - -void crossentropy_forward(float* losses, - float* probs, int* targets, - int B, int T, int Vp); - -void crossentropy_softmax_backward(float* dlogits, - float* dlosses, float* probs, int* targets, - int B, int T, int V, int Vp); - -#ifdef __cplusplus -} -#endif - diff --git a/experimental/profiler/metal.hpp b/experimental/profiler/metal.hpp deleted file mode 100644 index 3aaaad4..0000000 --- a/experimental/profiler/metal.hpp +++ /dev/null @@ -1,6 +0,0 @@ -#ifdef __APPLE__ -extern "C" { - void startCapture(); - void stopCapture(); -} -#endif diff --git a/experimental/profiler/metal.mm b/experimental/profiler/metal.mm deleted file mode 100644 index f4207b6..0000000 --- a/experimental/profiler/metal.mm +++ /dev/null @@ -1,46 +0,0 @@ -#import -#import -#import - - -extern "C" { - void startCapture() { - if (![[NSProcessInfo processInfo].environment[@"METAL_CAPTURE_ENABLED"] boolValue]) { - NSLog(@"METAL_CAPTURE_ENABLED is not set. Please set it to 1 to enable Metal capture."); - return; - } - - MTLCaptureDescriptor *descriptor = [[MTLCaptureDescriptor alloc] init]; - descriptor.destination = MTLCaptureDestinationGPUTraceDocument; - descriptor.outputURL = [NSURL fileURLWithPath:@"gpu.cpp.gputrace"]; - - NSFileManager *fileManager = [NSFileManager defaultManager]; - if ([fileManager fileExistsAtPath:@"gpu.cpp.gputrace"]) { - NSError *error = nil; - [fileManager removeItemAtPath:@"gpu.cpp.gputrace" error:&error]; - if (error) { - NSLog(@"Error deleting existing gpu.cpp.gputrace directory: %@", error); - return; - } else { - NSLog(@"Deleted existing gpu.cpp.gputrace directory."); - } - } - - NSError *error = nil; - id device = MTLCreateSystemDefaultDevice(); - if (!device) { - NSLog(@"MTLCreateSystemDefaultDevice returned nil. Metal may not be supported on this system."); - return; - } - descriptor.captureObject = device; - - BOOL success = [MTLCaptureManager.sharedCaptureManager startCaptureWithDescriptor:descriptor error:&error]; - if (!success) { - NSLog(@" error capturing mtl => %@ ", [error localizedDescription] ); - } - } - - void stopCapture() { - [MTLCaptureManager.sharedCaptureManager stopCapture]; - } -} diff --git a/experimental/web/CMakeLists.txt b/experimental/web/CMakeLists.txt deleted file mode 100644 index db4c281..0000000 --- a/experimental/web/CMakeLists.txt +++ /dev/null @@ -1,16 +0,0 @@ -cmake_minimum_required(VERSION 3.13) -project(run) -set(CMAKE_CXX_STANDARD 20) - -add_executable(run "run.cpp") - -set_target_properties(run PROPERTIES SUFFIX ".html") -set(SHELL_FILE_PATH "../custom_shell.html") - -# target_link_options(run PRIVATE "-sUSE_WEBGPU=1" "-sUSE_GLFW=3") - -# asyncify for emscripten_sleep() -# see https://emscripten.org/docs/api_reference/emscripten.h.html#pseudo-synchronous-functions -target_link_options(run PRIVATE "-sUSE_WEBGPU=1" "-sASYNCIFY=1" "--shell-file=${SHELL_FILE_PATH}") - -target_include_directories(run PRIVATE "../../") diff --git a/experimental/web/Makefile b/experimental/web/Makefile deleted file mode 100644 index 69a0000..0000000 --- a/experimental/web/Makefile +++ /dev/null @@ -1,38 +0,0 @@ -GPUCPP=../.. -FLAGS=-std=c++17 -s USE_WEBGPU=1 -s ASYNCIFY=1 -I$(GPUCPP) - -.PHONY: default cmake check-emsdk browser clean - -default: server - -build/run.html: check-emsdk run.cpp custom_shell.html - em++ run.cpp -o build/run.html \ - $(FLAGS) \ - --shell-file custom_shell.html \ - -build/run.wasm: check-emsdk run.cpp custom_shell.html - em++ run.cpp -o build/run.wasm \ - $(FLAGS) \ - --shell-file custom_shell.html \ - -# make clean explicit here because custom_shell.html changes don't trigger a rebuild -cmake: check-emsdk clean - emcmake cmake -B build && cmake --build build -j4 - python3 -m http.server --directory . - -server: build/run.html - @echo "\n┌───────────────────────────────────────────────────────────────────────────────────┐" - @echo "│ Open http://localhost:8000/build/run.html in your browser to see the output. │" - @echo "│ │" - @echo "│ Press Ctrl+C to stop the server. │" - @echo "└───────────────────────────────────────────────────────────────────────────────────┘\n\n" - python3 -m http.server --directory . - -open-browser: - open http://127.0.0.1:8000/build/run.html - -clean: - rm -rf build/* - -check-emsdk: - @which em++ > /dev/null || (echo "emsdk not found. Please install emsdk and run 'source emsdk_env.sh' in the emsdk directory." && exit 1) diff --git a/experimental/web/README.md b/experimental/web/README.md deleted file mode 100644 index 6b4e11b..0000000 --- a/experimental/web/README.md +++ /dev/null @@ -1,3 +0,0 @@ -Warning: web targets are not supported for now. - -We'll enable them and move this to examples/ once emscripten's WebGPU implementation catches up with the Dawn commit we're using. diff --git a/experimental/web/build/.gitkeep b/experimental/web/build/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/experimental/web/custom_shell.html b/experimental/web/custom_shell.html deleted file mode 100644 index 516d9b9..0000000 --- a/experimental/web/custom_shell.html +++ /dev/null @@ -1,53 +0,0 @@ - - - - - gpu.cpp - - - - - - -
- - - - - {{{ SCRIPT }}} - - diff --git a/experimental/web/run.cpp b/experimental/web/run.cpp deleted file mode 100644 index 224c24d..0000000 --- a/experimental/web/run.cpp +++ /dev/null @@ -1,57 +0,0 @@ -#include -#include -#include -#include - -#include "gpu.hpp" -#include "emscripten/emscripten.h" - -using namespace gpu; // createContext, createTensor, createKernel, - // createShader, dispatchKernel, wait, toCPU - // Tensor, Kernel, Context, Shape, kf32 - -static const char *kGelu = R"( -const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI) -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -@group(0) @binding(1) var dummy: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&inp)) { - let x: f32 = inp[i]; - out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR - * (x + .044715 * x * x * x))), x, x > 10.0); - } -} -)"; - -int main(int argc, char **argv) { - printf("\033[2J\033[1;1H"); - printf("\nHello gpu.cpp!\n"); - printf("--------------\n\n"); - - Context ctx = createContext({}); - static constexpr size_t N = 5000; - std::array inputArr, outputArr; - for (int i = 0; i < N; ++i) { - inputArr[i] = static_cast(i) / 10.0; // dummy input data - } - Tensor input = createTensor(ctx, Shape{N}, kf32, inputArr.data()); - Tensor output = createTensor(ctx, Shape{N}, kf32); - std::promise promise; - std::future future = promise.get_future(); - Kernel op = createKernel(ctx, {kGelu, 256, kf32}, - Bindings{input, output}, - {cdiv(N, 256), 1, 1}); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, outputArr.data(), sizeof(outputArr)); - for (int i = 0; i < 12; ++i) { - printf(" gelu(%.2f) = %.2f\n", inputArr[i], outputArr[i]); - } - printf(" ...\n\n"); - printf("Computed %zu values of GELU(x)\n\n", N); - return 0; -} diff --git a/experimental/wgsl.h b/experimental/wgsl.h deleted file mode 100644 index 25f3c93..0000000 --- a/experimental/wgsl.h +++ /dev/null @@ -1,84 +0,0 @@ -#ifndef GPU_CPP_WGSL_H -#define GPU_CPP_WGSL_H - -#include -#include -#include "utils/logging.hpp" // LOG - -namespace gpu { - -// Loop-unrolling optimization with regex -// -// Note: Be cautious, as it does not correctly recognize comments or lexical tokens. -std::string loopUnrolling(const std::string& code, int threshold = 32) { - // This regex pattern matches a for loop with the following structure: - // for (var : u32 = ; < ; ++) { } - std::regex forLoopPattern(R"(for\s*\(\s*var\s+(\w+):\s*u32\s*=\s*(\d+)\s*;\s*\1\s*<\s*(\d+)\s*;\s*\1\+\+\s*\)\s*\{\s*([^{}]*)\})"); - // Explanation of the regex: - // for\s*\( : Matches 'for (' with optional whitespace - // \s*var\s+ : Matches 'var ' with optional whitespace - // (\w+) : Captures the variable name (alphanumeric characters and underscores) - // :\s*u32\s*=\s* : Matches ': u32 = ' with optional whitespace - // (\d+) : Captures the start value (one or more digits) - // \s*;\s* : Matches ';' with optional whitespace - // \1\s*<\s* : Matches the captured variable name followed by '<' with optional whitespace - // (\d+) : Captures the end value (one or more digits) - // \s*;\s* : Matches ';' with optional whitespace - // \1\+\+\s* : Matches the captured variable name followed by '++' with optional whitespace - // \)\s*\{\s* : Matches ')' followed by '{' with optional whitespace - // ([^{}]*) : Captures the loop body (anything except '{' or '}') - // \} : Matches the closing '}' - - // Example: - // - // Input code: - // for (var i: u32 = 0; i < 3; i++) { std::cout << i << std::endl; } - // - // Matches: - // varName = "i" - // start = "0" - // end = "3" - // loopBody = "std::cout << i << std::endl;" - // - // Unrolled: - // std::cout << 0 << std::endl; - // std::cout << 1 << std::endl; - // std::cout << 2 << std::endl; - // - std::smatch match; - std::string unrolledCode = code; - while (std::regex_search(unrolledCode, match, forLoopPattern)) { - std::string varName = match[1]; - int start = std::stoi(match[2]); - int end = std::stoi(match[3]); - std::string loopBody = match[4]; - - if (end - start > threshold ) { - std::string skippedLoop = - "for (var " + - std::string(match[1]) + ": u32 = " + std::string(match[2]) + ";"+ - std::string(match[1]) + " < " + std::string(match[3]) + ";"+ - std::string(match[1]) + "++) /* Skipped */ {"+ - std::string(match[4]) + - "}"; - // LOG(kDefLog, kInfo, "Roll loop:%s", skippedLoop.c_str()); - unrolledCode = unrolledCode.substr(0, match.position()) + skippedLoop + unrolledCode.substr(match.position() + match.length()); - } else { - // LOG(kDefLog, kInfo, "Unroll loop(var: %s, start:%d, end:%d, body:%s)", varName.c_str(), start, end, loopBody.c_str()); - std::string unrolledLoop; - for (int i = start; i < end; ++i) { - std::string unrolledIteration = loopBody; - std::regex varPattern(varName); - unrolledIteration = std::regex_replace(unrolledIteration, varPattern, std::to_string(i)); - unrolledLoop += unrolledIteration; - } - unrolledCode = unrolledCode.substr(0, match.position()) + unrolledLoop + unrolledCode.substr(match.position() + match.length()); - } - } - - return unrolledCode; -} - -} // namespace gpu - -#endif diff --git a/gpu.hpp b/gpu.hpp index 5327fe7..f4b318b 100644 --- a/gpu.hpp +++ b/gpu.hpp @@ -1,1666 +1,738 @@ #ifndef GPU_HPP #define GPU_HPP +#include #include -#include +#include +#include +#include +#include #include #include #include -#include -#include +#include +#include +#include +#include #include -#include +#include +#include #include -#include -#include // std::pair +#include +#include #include +#if defined(__EMSCRIPTEN__) #include "webgpu/webgpu.h" - -#include "numeric_types/half.hpp" -#include "utils/logging.hpp" - -#ifdef __EMSCRIPTEN__ -#include "emscripten/emscripten.h" +#include "webgpu/webgpu_cpp.h" +#else +// Include Dawn's extended C header before the generated C++ facade. This makes +// the opt-in SPIR-V source descriptor available without dropping down to C. +#include "dawn/webgpu.h" +#include "dawn/webgpu_cpp.h" #endif -#ifdef USE_DAWN_API -#include "dawn/native/DawnNative.h" -#endif +#include "utils/logging.hpp" namespace gpu { -/** - * @brief Represents a buffer of values on the GPU. - */ -struct Array { - WGPUBuffer buffer; - WGPUBufferUsage usage; - size_t size; // in bytes -}; - -/** - * @brief Represents the shape of a tensor. - * - * The rank of the tensor is the - * number of dimensions in the shape. The data array stores the size of each - * dimension. For now, we limit the rank to 8 to avoid dynamic allocation. - * - * @code - * Shape shape = {256, 256}; - * @endcode - */ struct Shape { - static constexpr size_t kMaxRank = 8; // Maximum rank of a tensor, avoids - // dynamic allocation for shape data - std::array data = {0}; + static constexpr size_t kMaxRank = 8; + std::array data{}; size_t rank = 0; - inline Shape() = default; - inline Shape(std::initializer_list dims) { - assert(dims.size() <= kMaxRank); - std::copy(dims.begin(), dims.end(), data.begin()); - rank = dims.size(); + + Shape() = default; + Shape(std::initializer_list dimensions) : rank(dimensions.size()) { + if (rank > kMaxRank) throw std::invalid_argument("tensor rank exceeds 8"); + std::copy(dimensions.begin(), dimensions.end(), data.begin()); } - inline size_t &operator[](size_t index) { - assert(index < rank); + + size_t &operator[](size_t index) { + if (index >= rank) throw std::out_of_range("shape index"); return data[index]; } - inline const size_t &operator[](size_t index) const { - assert(index < rank); + const size_t &operator[](size_t index) const { + if (index >= rank) throw std::out_of_range("shape index"); return data[index]; } }; -/** - * @brief Returns the number of elements in a tensor with the given shape, - * which is equal to the product of the dimensions. - * @param[in] shape Shape of the tensor - * @return Number of elements in the tensor - * - * @code - * size({256, 256}) -> 65536 - * @endcode - */ inline size_t size(const Shape &shape) { - size_t numels = 1; - for (size_t i = 0; i < shape.rank; i++) { - numels *= shape.data[i]; - } - return numels; + size_t result = 1; + for (size_t i = 0; i < shape.rank; ++i) result *= shape[i]; + return result; } -/** - * @brief Represents a tensor on the GPU, which is a buffer of values with a - * shape. - * - * @code - * Tensor tensor = createTensor(ctx, {256, 256}, kf32); - * @endcode - */ -struct Tensor { - Array data; - Shape shape; -}; - -/** - * @brief Represents a non-owning view into a tensor specifying an offset and a - * subspan. This is useful for specifying a slice of a tensor on the GPU - * without copying the data. - * - * @code - * TensorView view = {tensor, 0, 256}; - * @endcode - */ -struct TensorView { - Tensor data; // non-owning view - size_t offset = 0; - size_t span = 0; -}; - -/** - * @brief Represents an ordered collection of WGPUBuffers (wrapped as tensors, - * non-overlapping views, or arrays) for the purpose of binding them to a - * kernel operation to make them accessible to the GPU kernel. - * - * The ordering of the bindings should match the binding indices in the WGSL - * code. - */ -template struct Bindings { - std::array data; - std::array viewOffsets; - std::array viewSpans; - Bindings(const std::initializer_list &init) { - std::copy(begin(init), end(init), begin(data)); - std::fill(begin(viewOffsets), end(viewOffsets), 0); - for (size_t i = 0; i < N; ++i) { - viewSpans[i] = data[i].data.size; - } - } - - Bindings(const std::array &init) { - std::copy(begin(init), end(init), begin(data)); - std::fill(begin(viewOffsets), end(viewOffsets), 0); - for (size_t i = 0; i < N; ++i) { - viewSpans[i] = data[i].data.size; - } - } - - Bindings(const std::initializer_list &init) { - size_t i = 0; - for (const auto &tv : init) { - data[i] = tv.data; - viewOffsets[i] = tv.offset; - viewSpans[i] = tv.span; - ++i; - } - } - - Bindings(const std::initializer_list &init) { - std::copy(begin(init), end(init), begin(data)); - std::fill(begin(viewOffsets), end(viewOffsets), 0); - for (size_t i = 0; i < N; ++i) { - viewSpans[i] = data[i].size; - } - } - - Tensor &operator[](std::size_t index) { return data[index]; } - const Tensor &operator[](std::size_t index) const { return data[index]; } -}; +inline size_t ceilDiv(size_t value, size_t divisor) { + if (!divisor) throw std::invalid_argument("division by zero"); + return value / divisor + (value % divisor != 0); +} -/** - * @brief Deduction guide for Bindings - */ -template Bindings(std::array) -> Bindings; -template Bindings(Args...) -> Bindings; - -struct Context; // Forward declaration so that TensorPool can have a pointer to - // Context - -/** - * @brief Represents a pool of tensors to manage GPU resources. The pool is - * responsible for managing the lifetime of the tensors and freeing them when - * the pool is destroyed. - * - * Most users do not need to interact with the TensorPool type, as there is a - * member instance in the Context struct to simplify lifetime management of GPU - * resources. - */ -struct TensorPool { - inline TensorPool(Context *ctx) : ctx(ctx), data() {}; - Context *ctx; - std::unordered_map data; - ~TensorPool(); -}; +inline Shape ceilDiv(const Shape &value, const Shape &divisor) { + if (value.rank != divisor.rank) + throw std::invalid_argument("shape ranks differ"); + Shape result = value; + for (size_t i = 0; i < value.rank; ++i) + result[i] = ceilDiv(value[i], divisor[i]); + return result; +} -enum NumType { - kf16, // (experimental) - kf32, - ki32 -}; +enum NumType { kf16, kf32, ki32 }; -/** - * @brief Returns the number of bytes of a number type. - */ -inline size_t sizeBytes(const NumType &type) { +inline size_t sizeBytes(NumType type) { switch (type) { - case kf16: - return sizeof(uint16_t); - case kf32: - return sizeof(float); - case ki32: - return sizeof(int32_t); - default: - LOG(kDefLog, kError, "Invalid NumType in size calculation."); - return 0; + case kf16: return sizeof(uint16_t); + case kf32: return sizeof(float); + case ki32: return sizeof(int32_t); } + throw std::invalid_argument("unknown numeric type"); } -/** - * @brief Converts NumType to string. - */ inline std::string toString(NumType type) { switch (type) { - case kf16: - return "f16"; - case kf32: - return "f32"; - case ki32: - return "i32"; - default: - LOG(kDefLog, kError, "Invalid NumType in string conversion."); - return "unknown"; + case kf16: return "f16"; + case kf32: return "f32"; + case ki32: return "i32"; } + throw std::invalid_argument("unknown numeric type"); } -/** - * @brief Converts Shape to string. The string formatting is meant to be - * slotted into WGSL code (hence no additional parentheses or brackets). - */ inline std::string toString(const Shape &shape) { - std::string str; - for (size_t i = 0; i < shape.rank; i++) { - str += std::to_string(shape.data[i]); - if (i < shape.rank - 1) { - str += ", "; - } + std::string result; + for (size_t i = 0; i < shape.rank; ++i) { + if (i) result += ", "; + result += std::to_string(shape[i]); } - return str; + return result; } -/** - * @brief Converts size_t to string. Wraps std::to_string for consistency, - * instead of having to remember to switch between std::to_string and toString - * depending on the type. - */ -inline std::string toString(size_t value) { return std::to_string(value); } - -/** - * @brief simple in-place string replacement helper function for substituting - * placeholders in a WGSL string template. - * - * Note this is not meant to be used in performance-critical code paths and - * should be used ahead-of-time before any performance-critical codepath to - * preprocess WGSL code strings. - * - * @param[in] str String to mutate with substitution replacements. - * @param[in] from Substring to replace - * @param[in] to Substring to replace with - * - * @code - * replaceAll(str, "{{workgroupSize}}", "256"); - * @endcode - */ -inline void replaceAll(std::string &str, const std::string &from, - const std::string &to) { - size_t start_pos = 0; - while ((start_pos = str.find(from, start_pos)) != std::string::npos) { - str.replace(start_pos, from.length(), to); - start_pos += to.length(); - } +inline void replaceAll(std::string &text, std::string_view from, + std::string_view to) { + if (from.empty()) return; + for (size_t pos = 0; (pos = text.find(from, pos)) != std::string::npos; + pos += to.size()) + text.replace(pos, from.size(), to); } -/** - * @brief KernelCode is the representation of WGSL GPU code with template - * substitutions applied. It is a type around the code string with additional - * metadata for workgroup size and precision since they are specified in the - * WGSL code. Additionally, label and entryPoint are used by `createKernel()` - * to specify the label and entry point of the kernel. - */ -struct KernelCode { - /** - * @brief Constructor to create a code object from a template - * string and optional workgroup size and precision. - * - * @param[in] pData Shader template string with placeholders - * @param[in] workgroupSize Shape of the workgroup. Unlike tensor shapes which - * can be of arbitrary rank, workgroup size is always of rank 3 corresponding - * to x y and z. workgroupSize is stored as a field in the KernelCode instance - * that is returned by createShader(). - * @param[in] precision Data type precision to be substituted for - * {{precision}} in the WGSL code. As with workgroupSize, precision is stored - * as a field in the KernelCode instance that is returned by createShader(). - * @code - * KernelCode code = {kShaderText, {256, 1, 1}, kf32}; - * @endcode - */ - inline KernelCode(const std::string &pData = "", size_t workgroupSize = 256, - NumType precision = kf32) - : data(pData), workgroupSize({workgroupSize, 1, 1}), - precision(precision) { - if (precision == kf16) { - data = "enable f16;\n" + data; - } - replaceAll(data, "{{workgroupSize}}", toString({workgroupSize, 1, 1})); - replaceAll(data, "{{precision}}", toString(precision)); - LOG(kDefLog, kTrace, "Shader code:\n%s", data.c_str()); - } +inline void replaceAll( + std::string &text, + std::initializer_list> values) { + for (const auto &[from, to] : values) replaceAll(text, from, to); +} - /** - * @brief Overload of the constructor to create a code object from a template - * string and workgroup size. This overload takes a single size_t - * workgroupSize parameter instead of a 3D shape for the workgroup size and - * instantiates a 3D shape with the workgroupSize in the x dimension and 1 in - * the y and z dimensions. - * - * @param[in] pData Shader template string with placeholders @param[in] - * workgroupSize 3D Workgroup size - * @param[in] precision Data type precision for the shader - * - * @code KernelCode code = {kPuzzle1, 256, kf32}; @endcode - */ - inline KernelCode(const std::string &pData, - const Shape &workgroupSize = {256, 1, 1}, - NumType precision = kf32) - : data(pData), workgroupSize(workgroupSize), precision(precision) { - if (precision == kf16) { - data = "enable f16;\n" + data; - } - replaceAll(data, "{{workgroupSize}}", toString(workgroupSize)); - replaceAll(data, "{{precision}}", toString(precision)); - LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str()); - } +struct WGSL { + std::string code; + Shape workgroupSize{256, 1, 1}; + std::string label = "kernel"; + std::string entryPoint = "main"; - /** - * @brief Overload of the constructor, adding totalWorkgroups parameter to - * perform a string replacement for the total number of workgroups in the - * kernel code. - * - * @param[in] pData Shader template string with placeholders - * @param[in] workgroupSize 3D Workgroup size - * @param[in] precision Data type precision for the shader - * @param[in] totalWorkgroups Total number of workgroups in the kernel - * - * @code - * KernelCode code = {kPuzzle1, {256, 1, 1}, kf32, {2, 2, 1}}; - * @endcode - */ - inline KernelCode(const std::string &pData, const Shape &workgroupSize, - NumType precision, const Shape &totalWorkgroups) - : data(pData), workgroupSize(workgroupSize), precision(precision) { - if (precision == kf16) { - data = "enable f16;\n" + data; - } - replaceAll(data, "{{workgroupSize}}", toString(workgroupSize)); - replaceAll(data, "{{precision}}", toString(precision)); - replaceAll(data, "{{totalWorkgroups}}", toString(totalWorkgroups)); - LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str()); + WGSL(std::string code = {}, size_t workgroupSize = 256, + NumType precision = kf32) + : WGSL(std::move(code), Shape{workgroupSize, 1, 1}, precision) {} + + WGSL(std::string code, Shape workgroupSize, NumType precision = kf32) + : code(std::move(code)), workgroupSize(workgroupSize) { + if (this->workgroupSize.rank != 3) + throw std::invalid_argument("workgroup size must have three dimensions"); + if (precision == kf16) this->code = "enable f16;\n" + this->code; + replaceAll(this->code, "{{workgroupSize}}", toString(this->workgroupSize)); + replaceAll(this->code, "{{precision}}", toString(precision)); } - /** - * @brief Overload of the constructor, adding totalWorkgroups parameter as - * well as the size_t 1D workgroupSize parameter. - * - * @param[in] pData Shader template string with placeholders - * @param[in] workgroupSize Workgroup size in the x dimension - * @param[in] precision Data type precision for the shader - * @param[in] totalWorkgroups Total number of workgroups in the kernel - * - * @code - * KernelCode code = {kPuzzle1, {256, 1, 1}, kf32, {2, 2, 1}}; - * @endcode - */ - inline KernelCode(const std::string &pData, const size_t &workgroupSize, - NumType precision, const Shape &totalWorkgroups) - : data(pData), workgroupSize({workgroupSize, 1, 1}), - precision(precision) { - if (precision == kf16) { - data = "enable f16;\n" + data; - } - replaceAll(data, "{{workgroupSize}}", toString({workgroupSize, 1, 1})); - replaceAll(data, "{{precision}}", toString(precision)); - replaceAll(data, "{{totalWorkgroups}}", toString(totalWorkgroups)); - LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str()); + WGSL(std::string code, Shape workgroupSize, NumType precision, + const Shape &totalWorkgroups) + : WGSL(std::move(code), workgroupSize, precision) { + replaceAll(this->code, "{{totalWorkgroups}}", toString(totalWorkgroups)); } +}; - std::string data; - Shape workgroupSize; - NumType precision = kf32; +struct SPIRV { + std::vector code; std::string label = "kernel"; std::string entryPoint = "main"; }; -/** - * @brief Overload of the string replacement helper function to replace - * multiple substrings in a string with multiple replacements. - * - * @param[in] str String to mutate with substitution replacements. - * @param[in] reps Vector of pairs of substrings to replace and their - * replacements. - * - * @code - * replaceAll(str, {{"{{workgroupSize}}", "256"}, {"{{precision}}", - * @endcode - * "f32"}}); - */ -inline const std::string -replaceAll(std::string &str, - const std::vector> &reps) { - for (const auto &rep : reps) { - replaceAll(str, rep.first, rep.second); - } - - return str; -} - -/** - * @brief Used for on-done callback data for asynchronous operations sduch as - * kernel launching. - */ -struct CallbackData { - WGPUBuffer buffer; // managed by owning Kernel - size_t bufferSize; - void *output; // non-owning, only for target memory in toCPU, not used for - // kernel invocations - std::promise *promise; - std::future *future; -}; +using Shader = std::variant; -/** - * @brief Staging buffer and callback data for copying data between the GPU and - * CPU. - */ -struct CopyData { - WGPUCommandBuffer commandBuffer; - WGPUBuffer readbackBuffer; - std::promise promise; - std::future future; +struct Tensor { + wgpu::Buffer buffer; + Shape shape; + NumType type = kf32; + size_t bytes = 0; }; -/** - * @brief Represents handles + metadata for a reusable kernel on the GPU. - * The struct members can be divided into "consumed upon dispatch" - * (commandBuffer) and reusable ahead-of-time setup (all other members). - */ -struct RawKernel { - std::unique_ptr buffers; // non-owning - std::unique_ptr bufferSizes; - size_t numBindings; - Shape totalWorkgroups; - WGPUBindGroup bindGroup; // persists between submission - WGPUComputePipeline computePipeline; // persists between submission - WGPUCommandBuffer commandBuffer; // destroyed upon submission - bool used; +struct TensorView { + const Tensor &tensor; + size_t offset = 0; + size_t bytes = 0; }; -typedef std::shared_ptr Kernel; - -/** - * @brief A struct to package the result of a WGSL code compilation. - */ -struct CompilationInfo { - WGPUCompilationInfoRequestStatus status; - std::vector messages; - std::vector lineNums; - std::vector linePos; - bool finished; // true if the compilation is finished +struct Binding { + wgpu::Buffer buffer; + size_t offset; + size_t bytes; + wgpu::BufferBindingType type; }; -/** - * @brief Operator implementation to make the Kernel type hashable. - * @param[in] lhs First Kernel instance to compare - * @param[in] rhs Second Kernel instance to compare - * @return True if lhs < rhs, false otherwise - */ -inline bool operator<(const Kernel &lhs, const Kernel &rhs) { - return lhs->commandBuffer < rhs->commandBuffer; +inline Binding read(const Tensor &tensor) { + return {tensor.buffer, 0, tensor.bytes, + wgpu::BufferBindingType::ReadOnlyStorage}; } -/** - * @brief A pool of kernels to manage GPU resources. For simple use cases this - * is instantiated as a member in the Context struct although it's possible to - * have multiple resource pools of kernels in more complex scenarios. - */ -struct KernelPool { - inline KernelPool(Context *ctx) : ctx(ctx), data() {} - Context *ctx; - std::unordered_map data; - inline ~KernelPool() { - // Note : Some kernel resources such as commandBuffer are harvested by - // queue submission, explicitly destroying readback and callback buffers - // produces runtime errors. - data.clear(); - } -}; +inline Binding readWrite(const Tensor &tensor) { + return {tensor.buffer, 0, tensor.bytes, wgpu::BufferBindingType::Storage}; +} -inline void processEvents(const WGPUInstance &instance) { -#ifdef __EMSCRIPTEN__ - emscripten_sleep(0); -#else - wgpuInstanceProcessEvents(instance); -#endif +inline Binding read(const TensorView &view) { + if (view.offset > view.tensor.bytes) + throw std::out_of_range("tensor view starts beyond its buffer"); + const size_t bytes = view.bytes ? view.bytes : view.tensor.bytes - view.offset; + if (bytes > view.tensor.bytes - view.offset) + throw std::out_of_range("tensor view exceeds its buffer"); + return {view.tensor.buffer, view.offset, bytes, + wgpu::BufferBindingType::ReadOnlyStorage}; } -/** - * @brief Represents a GPU context, aggregates WebGPU API handles to interact - * with the GPU including the instance, adapter, device, and queue. - * - * Additionally contains a TensorPool and KernelPool for managing GPU resources - * to simplify lifetime management of GPU resources. - */ -struct Context { - WGPUInstance instance = nullptr; - WGPUAdapter adapter = nullptr; - WGPUDevice device = nullptr; - WGPUQueue queue = nullptr; - TensorPool pool = TensorPool(this); - KernelPool kernelPool = KernelPool(this); - WGPURequestAdapterStatus adapterStatus; - WGPURequestDeviceStatus deviceStatus; - - // Default constructor - Context() = default; +inline Binding readWrite(const TensorView &view) { + auto binding = read(view); + binding.type = wgpu::BufferBindingType::Storage; + return binding; +} - Context(Context&& other) noexcept - : instance(other.instance), - adapter(other.adapter), - device(other.device), - queue(other.queue), - // Re‐initialize pools to point to *this*: - pool(this), - kernelPool(this), - adapterStatus(other.adapterStatus), - deviceStatus(other.deviceStatus) - { - LOG(kDefLog, kTrace, "Moving Context ownership"); - // Move over the resources in the pools: - pool.data = std::move(other.pool.data); - kernelPool.data = std::move(other.kernelPool.data); - - // Null out handles in the source so its destructor won't release them. - other.instance = nullptr; - other.adapter = nullptr; - other.device = nullptr; - other.queue = nullptr; - // other.adapterStatus = 0; - // other.deviceStatus = 0; - } +template struct Bindings { + std::array data; - Context& operator=(Context&& other) noexcept { - if (this != &other) { - // Free any existing resources. In most cases, this should be a no-op - // since we typically shouldn't have two active initialized Context - // instances with resources acquired. - this->~Context(); - // Then placement‐new a move‐constructed copy in-place: - new (this) Context(std::move(other)); - } - return *this; - } + template + requires(sizeof...(T) == N && (std::same_as, Binding> && ...)) + explicit Bindings(T &&...bindings) + : data{std::forward(bindings)...} {} - ~Context() { - LOG(kDefLog, kTrace, "Destroying context"); - if (queue) { - wgpuQueueRelease(queue); - } else { - LOG(kDefLog, kTrace, "Queue is null"); - } - if (device) { - wgpuDeviceRelease(device); - processEvents(instance); - } else { - LOG(kDefLog, kTrace, "Device is null"); - } - if (adapter) { - wgpuAdapterRelease(adapter); - processEvents(instance); - } else { - LOG(kDefLog, kTrace, "Adapter is null"); - } - if (instance) { - wgpuInstanceRelease(instance); - } else { - LOG(kDefLog, kTrace, "Instance is null"); - } - LOG(kDefLog, kTrace, "Context destroyed"); - } + const Binding &operator[](size_t index) const { return data.at(index); } }; -/** - * @brief Tensor factory function to create a tensor (a Tensor type is simply - * an Array with an N-dimensional Shape specification) on the GPU. The tensor - * is created with the given shape, data type, and usage flags, added to the - * TensorPool, and returned. - * - * This is the core implementation which takes the minimal set of parameters in - * terms of the raw WebGPU API, and is used by the other createTensor overloads - * which provide more ergonomic interfaces. - * - * @param[in] pool TensorPool instance to manage the tensor - * @param[in] device WGPUDevice instance to create the tensor on - * @param[in] shape Shape of the tensor - * @param[in] dtype Data type of the tensor (e.g. kf32) - * @param[in] usage Usage flags for the tensor buffer - * @return Tensor instance representing the created tensor - * - * @code - * Tensor tensor = createTensor(pool, device, {256, 256}, kf32); - * @endcode - */ -inline Tensor createTensor(TensorPool &pool, WGPUDevice &device, - const Shape &shape, NumType dtype, - WGPUBufferUsage usage = WGPUBufferUsage_Storage | - WGPUBufferUsage_CopyDst | - WGPUBufferUsage_CopySrc) { - LOG(kDefLog, kTrace, "Creating tensor"); - size_t numElements = size(shape); - size_t size = sizeBytes(dtype) * numElements; - WGPUBufferDescriptor bufferDesc = { - .label = {.data = nullptr, .length = 0}, - .usage = usage, - .size = size, - }; - WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &bufferDesc); - pool.data[buffer] = Tensor{ - .data = Array{.buffer = buffer, .usage = usage, .size = size}, - .shape = shape, - }; - return pool.data[buffer]; -} +template Bindings(T &&...) -> Bindings; -/** - * @brief Overload of the tensor factory function to instantiate a tensor on - * the GPU with a given shape and data type. - * - * Instead of taking the TensoPool and raw WebGPU API WGPUDevice and - * WGPUBufferUsage arguments, this is a convenience wrapper around the - * core createTensor function which has default usage flags for a storage - * buffer, and also takes in the Context object. - * - * instance instead of the narrower TensorPool object. - * @param[in] ctx Context instance to manage the tensor - * @param[in] shape Shape of the tensor - * @param[in] dtype Data type of the tensor (e.g. kf32) - * @return Tensor instance representing the created tensor - * - * @code - * Tensor tensor = createTensor(ctx, {256, 256}, kf32); - * @endcode - */ -inline Tensor createTensor(Context &ctx, const Shape &shape, NumType dtype) { - return createTensor(ctx.pool, ctx.device, shape, dtype); -} +struct ContextOptions { + wgpu::PowerPreference powerPreference = wgpu::PowerPreference::HighPerformance; +#if defined(__APPLE__) + wgpu::BackendType backend = wgpu::BackendType::Metal; +#elif defined(__linux__) + wgpu::BackendType backend = wgpu::BackendType::Vulkan; +#else + wgpu::BackendType backend = wgpu::BackendType::Undefined; +#endif + wgpu::FeatureLevel featureLevel = wgpu::FeatureLevel::Core; + std::vector requiredFeatures; + std::optional requiredLimits; + bool enableSPIRV = false; +}; -/** - * @brief Overload of the tensor factory function to instantiate a tensor on - * the GPU with a given shape, data type. This overload also takes initial - * float* data to populate the tensor with. - * - * The data is assumed to be of size equal to the product of the dimensions in - * the shape, and is copied to the GPU buffer. - * - * @param[in] ctx Context instance to manage the tensor - * @param[in] shape Shape of the tensor - * @param[in] dtype Data type of the tensor (e.g. kf32) - * @param[in] data Initial data to populate the tensor with - * @return Tensor instance representing the created tensor - * - * @code - * Tensor tensor = createTensor(ctx, {256, 256}, kf32, data); - * @endcode - */ -inline Tensor createTensor(Context &ctx, const Shape &shape, NumType dtype, - const float *data) { - assert(dtype == kf32); - Tensor tensor = - createTensor(ctx.pool, ctx.device, shape, dtype, - WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | - WGPUBufferUsage_CopySrc); - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, - tensor.data.size); - return tensor; -} +namespace detail { -inline Tensor createTensor(Context &ctx, const Shape &shape, NumType dtype, - const int32_t *data) { - assert(dtype == ki32); - Tensor tensor = - createTensor(ctx.pool, ctx.device, shape, dtype, - WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | - WGPUBufferUsage_CopySrc); - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, - tensor.data.size); - return tensor; +inline std::string string(wgpu::StringView value) { + return std::string(static_cast(value)); } -/** - * @brief Overload of the tensor factory function to instantiate a tensor on - * the GPU with a given shape, data type. This overload also takes initial - * half* data to populate the tensor with. - * - * The data is assumed to be of size equal to the product of the dimensions in - * the shape, and is copied to the GPU buffer. - * - * @param[in] ctx Context instance to manage the tensor - * @param[in] shape Shape of the tensor - * @param[in] dtype Data type of the tensor (e.g. kf32) - * @param[in] data Initial data to populate the tensor with - * @return Tensor instance representing the created tensor - * - * @code - * Tensor tensor = createTensor(ctx, {256, 256}, kf32, data); - * @endcode - */ -inline Tensor createTensor(Context &ctx, const Shape &shape, NumType dtype, - const half *data) { - assert(dtype == kf16); - Tensor tensor = - createTensor(ctx.pool, ctx.device, shape, dtype, - WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | - WGPUBufferUsage_CopySrc); - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, - tensor.data.size); - return tensor; -} +struct ErrorState { + std::mutex mutex; + std::string message; -/** - * @brief Frees a tensor resource and updates the tensor pool. - * - * Only needed if the use case requires manually managing resource lifetimes of - * GPU tensors. For simple use cases, the TensorPool destructor will - * automatically free all tensors. - * - * @param[in] pool TensorPool instance to manage the tensor - * @param[in] tensor Tensor instance to free - * - * @code - * FreeTensor(pool, tensor); - * @endcode - */ -inline void FreeTensor(TensorPool &pool, Tensor tensor) { - if (tensor.data.buffer) { - wgpuBufferRelease(tensor.data.buffer); - } else { - LOG(kDefLog, kWarn, "Tried to free tensor with null buffer"); - } - if (pool.data.find(tensor.data.buffer) != pool.data.end()) { - pool.data.erase(tensor.data.buffer); - } else { - LOG(kDefLog, kWarn, "Tried to free tensor that was not in pool"); + void set(std::string value) { + std::lock_guard lock(mutex); + message = std::move(value); } -} -/** - * @brief Destructor for TensorPool which frees all tensors in the pool. - */ -inline TensorPool::~TensorPool() { - // Need to get keys in a separate iteration, otherwise iterator is getting - // invalidated during erase. - std::vector keys; - for (auto &pair : data) { - keys.push_back(pair.first); + std::string take() { + std::lock_guard lock(mutex); + return std::exchange(message, {}); } - for (auto &key : keys) { - FreeTensor(*this, data[key]); - LOG(kDefLog, kTrace, "Freed tensor"); - } -} - -/** - * @brief Checks a condition and logs an error message if the condition is - * false. - * @param[in] condition The condition to check. - * @param[in] message The error message to log if the condition is false. - * @param[in] file The source file where the check is performed. - * @param[in] line The line number in the source file where the check is - * performed. - */ -inline void check(bool condition, const char *message, - const char *file = "unkown", int line = -1) { - if (!condition) { - LOG(kDefLog, kError, "Error in file %s line %d:\n%s", file, line, message); - } else { - LOG(kDefLog, kTrace, "Success in file %s line %d:\n%s", file, line, - message); - } -} +}; -/** - * @brief Factory function to create a GPU context, which aggregates WebGPU API - * handles to interact with the GPU including the instance, adapter, device, and - * queue. - * - * The function takes optional descriptor parameters for the instance - * descriptor, adapter request options, and device descriptor, which are passed - * through to the WebGPU API calls to create the instance, adapter, and device. - * - * If dawn is used, it also sets up an error callback for device loss. - * - * @param[in] desc Instance descriptor for the WebGPU instance (optional) - * @param[in] adapterOpts Adapter request options for the WebGPU adapter - * (optional) - * @param[in] devDescriptor Device descriptor for the WebGPU device (optional) - * @return Context instance representing the created GPU context - * - */ -inline Context createContext( - const WGPUInstanceDescriptor &desc = {}, - const WGPURequestAdapterOptions &adapterOpts = {}, - const WGPUDeviceDescriptor &devDescriptor = {}) -{ - Context ctx; // stack-allocated - -#ifdef __EMSCRIPTEN__ - ctx.instance = wgpuCreateInstance(nullptr); +inline constexpr wgpu::CallbackMode completionMode() { +#if defined(__EMSCRIPTEN__) + return wgpu::CallbackMode::WaitAnyOnly; #else - ctx.instance = wgpuCreateInstance(&desc); + return wgpu::CallbackMode::AllowProcessEvents; #endif - check(ctx.instance, "Initialize WebGPU", __FILE__, __LINE__); - - LOG(kDefLog, kTrace, "Requesting adapter"); - { - struct AdapterData { - WGPUAdapter adapter = nullptr; - bool requestEnded = false; - WGPURequestAdapterStatus status; - }; - AdapterData adapterData; - - auto onAdapterRequestEnded = [](WGPURequestAdapterStatus status, - WGPUAdapter adapter, - WGPUStringView message, - void *pUserData, void *) { - auto &ad = *reinterpret_cast(pUserData); - ad.status = status; -#ifdef __EMSCRIPTEN__ - if (status != WGPURequestAdapterStatus_Success) { - LOG(kDefLog, kError, "Could not get WebGPU adapter: %.*s", - static_cast(message.length), message.data); - } -#endif - check(status == WGPURequestAdapterStatus_Success, - "Request WebGPU adapter", __FILE__, __LINE__); - ad.adapter = adapter; - ad.requestEnded = true; - }; - - WGPURequestAdapterCallbackInfo callbackInfo { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = onAdapterRequestEnded, - .userdata1 = &adapterData, - .userdata2 = nullptr - }; - wgpuInstanceRequestAdapter(ctx.instance, &adapterOpts, callbackInfo); - - while (!adapterData.requestEnded) { - processEvents(ctx.instance); - } - ctx.adapter = adapterData.adapter; - ctx.adapterStatus = adapterData.status; - } - - LOG(kDefLog, kTrace, "Requesting device"); - { - struct DeviceData { - WGPUDevice device = nullptr; - bool requestEnded = false; - WGPURequestDeviceStatus status; - }; - DeviceData devData; - - auto onDeviceRequestEnded = [](WGPURequestDeviceStatus status, - WGPUDevice device, - WGPUStringView message, - void *pUserData, void *) { - auto &dd = *reinterpret_cast(pUserData); - dd.status = status; - check(status == WGPURequestDeviceStatus_Success, - "Could not get WebGPU device.", __FILE__, __LINE__); - LOG(kDefLog, kTrace, "Device Request succeeded %p", - static_cast(device)); - dd.device = device; - dd.requestEnded= true; - }; - - WGPURequestDeviceCallbackInfo deviceCallbackInfo { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = onDeviceRequestEnded, - .userdata1= &devData, - .userdata2= nullptr - }; - wgpuAdapterRequestDevice(ctx.adapter, &devDescriptor, deviceCallbackInfo); - - LOG(kDefLog, kTrace, "Waiting for device request to end"); - while (!devData.requestEnded) { - processEvents(ctx.instance); - } - LOG(kDefLog, kTrace, "Device request ended"); - - ctx.device = devData.device; - ctx.deviceStatus = devData.status; - - // If the device was created, set up logging and fetch the queue - if (devData.status == WGPURequestDeviceStatus_Success) { - WGPULoggingCallbackInfo loggingCallbackInfo { - .nextInChain = nullptr, - .callback = - [](WGPULoggingType type, WGPUStringView message, - void *, void *) { - LOG(kDefLog, kError, "Device logging callback: %.*s", - static_cast(message.length), message.data); - if (type == WGPULoggingType_Error) { - throw std::runtime_error("Device error logged."); - } - }, - .userdata1 = nullptr, - .userdata2 = nullptr - }; - wgpuDeviceSetLoggingCallback(ctx.device, loggingCallbackInfo); - ctx.queue = wgpuDeviceGetQueue(ctx.device); - } - } - - return std::move(ctx); } - -#ifdef USE_DAWN_API -/** - * @brief Factory function to create a GPU context, which aggregates WebGPU API - * handles to interact with the GPU including the instance, adapter, device, and - * queue. - * - * The function takes gpu index to support for multi GPUs. - * To activate this function, it needs not only webgpu's headers but also DAWN's - * headers. - * - * If dawn is used, it also sets up an error callback for device loss. - * - * @param[in] gpuIdx GPU index - * @param[in] desc Instance descriptor for the WebGPU instance (optional) - * @param[in] devDescriptor Device descriptor for the WebGPU device (optional) - * @return Context instance representing the created GPU context - * - * @code - * Context ctx = createContextByGpuIdx(1); - * @endcode - */ -inline Context -createContextByGpuIdx(int gpuIdx, const WGPUInstanceDescriptor &desc = {}, - const WGPUDeviceDescriptor &devDescriptor = {}) { - Context context; - { -#ifdef __EMSCRIPTEN__ - // Emscripten does not support the instance descriptor - // and throws an assertion error if it is not nullptr. - context.instance = wgpuCreateInstance(nullptr); +inline constexpr wgpu::CallbackMode persistentMode() { +#if defined(__EMSCRIPTEN__) + return wgpu::CallbackMode::AllowSpontaneous; #else - context.instance = wgpuCreateInstance(&desc); + return wgpu::CallbackMode::AllowProcessEvents; #endif - // check status - check(context.instance, "Initialize WebGPU", __FILE__, __LINE__); - } - - LOG(kDefLog, kInfo, "Requesting adapter"); - { - std::vector adapters = - dawn::native::Instance( - reinterpret_cast(context.instance)) - .EnumerateAdapters(); - LOG(kDefLog, kInfo, "The number of GPUs=%d\n", adapters.size()); - // Note: Second gpu is not available on Macos, but the number of GPUs is 2 - // on Macos. - // Calling wgpuAdapterGetInfo function for the second gpu becomes - // segfault. When you check all GPUs on linux, uncomment out following - // codes. - // - // for (size_t i = 0; i < adapters.size(); i++) { - // WGPUAdapterInfo info {}; - // auto ptr = adapters[i].Get(); - // if (ptr && adapters[i]) { - // wgpuAdapterGetInfo(ptr, &info); - // LOG(kDefLog, kInfo, "GPU(Adapter)[%d] = %s\n", i, info.description); - // wgpuAdapterInfoFreeMembers(info); - // } - // } - - { - LOG(kDefLog, kInfo, "Use GPU(Adapter)[%d]\n", gpuIdx); - auto ptr = adapters[gpuIdx].Get(); - if (ptr) { - WGPUAdapterInfo info{}; - wgpuAdapterGetInfo(ptr, &info); - LOG(kDefLog, kInfo, "GPU(Adapter)[%d] = %s\n", gpuIdx, - info.description); - wgpuAdapterInfoFreeMembers(info); - } - context.adapter = adapters[gpuIdx].Get(); - dawn::native::GetProcs().adapterAddRef(context.adapter); - } - } - - LOG(kDefLog, kInfo, "Requesting device"); - { - struct DeviceData { - WGPUDevice device = nullptr; - bool requestEnded = false; - }; - DeviceData devData; - - auto onDeviceRequestEnded = [](WGPURequestDeviceStatus status, - WGPUDevice device, WGPUStringView message, - void *pUserData, void *) { - DeviceData &devData = *reinterpret_cast(pUserData); - check(status == WGPURequestDeviceStatus_Success, - "Could not get WebGPU device.", __FILE__, __LINE__); - LOG(kDefLog, kTrace, "Device Request succeeded %x", - static_cast(device)); - devData.device = device; - devData.requestEnded = true; - }; - - WGPURequestDeviceCallbackInfo deviceCallbackInfo = { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = onDeviceRequestEnded, - .userdata1 = &devData, - .userdata2 = nullptr}; - wgpuAdapterRequestDevice(context.adapter, &devDescriptor, - deviceCallbackInfo); - - LOG(kDefLog, kInfo, "Waiting for device request to end"); - while (!devData.requestEnded) { - processEvents(context.instance); - } - LOG(kDefLog, kInfo, "Device request ended"); - assert(devData.requestEnded); - context.device = devData.device; - - WGPULoggingCallbackInfo loggingCallbackInfo = { - .nextInChain = nullptr, - .callback = - [](WGPULoggingType type, WGPUStringView message, void *userdata1, - void *userdata2) { - LOG(kDefLog, kError, "Device logging callback: %.*s", - static_cast(message.length), message.data); - if (type == WGPULoggingType_Error) { - throw std::runtime_error("Device error logged."); - } - }, - .userdata1 = nullptr, - .userdata2 = nullptr}; - wgpuDeviceSetLoggingCallback(context.device, loggingCallbackInfo); - } - context.queue = wgpuDeviceGetQueue(context.device); - return context; } -#endif -inline void wait(Context &ctx, std::future &future) { - while (future.wait_for(std::chrono::seconds(0)) != +template +T await(const wgpu::Instance &instance, std::future &result, + wgpu::Future event) { +#if defined(__EMSCRIPTEN__) + if (instance.WaitAny(event, UINT64_MAX) != wgpu::WaitStatus::Success) + throw std::runtime_error("could not wait for WebGPU operation"); +#else + (void)event; + while (result.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) { - processEvents(ctx.instance); + instance.ProcessEvents(); + std::this_thread::sleep_for(std::chrono::microseconds(100)); } +#endif + return result.get(); } -/** - * @brief Copies data from a GPU buffer to CPU memory. - * @param[in] ctx Context instance to manage the operation - * @param[in] tensor Tensor instance representing the GPU buffer to copy from - * @param[out] data Pointer to the CPU memory to copy the data to - * @param[in] bufferSize Size of the data buffer in bytes - * @param[in] op StagingBuffer instance to manage the operation - * - * @code - * toCPU(ctx, tensor, data, bufferSize); - * @endcode - */ -inline void toCPU(Context &ctx, Tensor &tensor, void *data, size_t bufferSize, - CopyData &op) { - wgpuQueueSubmit(ctx.queue, 1, &op.commandBuffer); - wgpuCommandBufferRelease(op.commandBuffer); - CallbackData callbackData = {op.readbackBuffer, bufferSize, data, &op.promise, - &op.future}; - - WGPUQueueWorkDoneCallbackInfo workDoneCallbackInfo = { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = - [](WGPUQueueWorkDoneStatus status, void *userdata1, void *userdata2) { - check(status == WGPUQueueWorkDoneStatus_Success, "Queue work done", - __FILE__, __LINE__); - const auto *data = static_cast(userdata1); - WGPUBufferMapCallbackInfo mapCallbackInfo = { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = - [](WGPUMapAsyncStatus status, WGPUStringView message, - void *userdata1, void *userdata2) { - const auto *data = static_cast(userdata1); - check(status == WGPUMapAsyncStatus_Success, - "Map readbackBuffer", __FILE__, __LINE__); - const void *mappedData = wgpuBufferGetConstMappedRange( - data->buffer, /*offset=*/0, data->bufferSize); - check(mappedData, "Get mapped range", __FILE__, __LINE__); - memcpy(data->output, mappedData, data->bufferSize); - wgpuBufferUnmap(data->buffer); - data->promise->set_value(); - }, - .userdata1 = const_cast(data), - .userdata2 = nullptr}; - wgpuBufferMapAsync(data->buffer, WGPUMapMode_Read, 0, - data->bufferSize, mapCallbackInfo); - }, - .userdata1 = &callbackData, - .userdata2 = nullptr}; - wgpuQueueOnSubmittedWorkDone(ctx.queue, workDoneCallbackInfo); - - wait(ctx, op.future); -} +} // namespace detail -/** - * @brief Overload of the toCPU function to copy data from a GPU buffer to CPU - * but initializes a staging buffer and promise/future for the operation for - * you. - * - * For simple use cases, this overload is recommended as it abstracts away the - * staging buffer and promise/future management. For more custom use cases - * where the staging buffer is initialized ahead of time, use the other - * overload. - * - * @param[in] ctx Context instance to manage the operation - * @param[in] tensor Tensor instance representing the GPU buffer to copy from - * @param[in] bufferSize Size of the data buffer in bytes - * @param[out] data Pointer to the CPU memory to copy the data to - */ -inline void toCPU(Context &ctx, Tensor &tensor, void *data, size_t bufferSize) { - CopyData op; - op.future = op.promise.get_future(); - { - WGPUBufferDescriptor readbackBufferDescriptor = { - .label = {.data = nullptr, .length = 0}, - .usage = WGPUBufferUsage_CopyDst | WGPUBufferUsage_MapRead, - .size = bufferSize, - }; - op.readbackBuffer = - wgpuDeviceCreateBuffer(ctx.device, &readbackBufferDescriptor); - } - { - WGPUCommandEncoder commandEncoder; - commandEncoder = wgpuDeviceCreateCommandEncoder(ctx.device, nullptr); - wgpuCommandEncoderCopyBufferToBuffer(commandEncoder, tensor.data.buffer, 0, - op.readbackBuffer, 0, bufferSize); - op.commandBuffer = wgpuCommandEncoderFinish(commandEncoder, nullptr); - wgpuCommandEncoderRelease(commandEncoder); - check(op.commandBuffer, "Create command buffer", __FILE__, __LINE__); - } - toCPU(ctx, tensor, data, bufferSize, op); - if (op.readbackBuffer) { - wgpuBufferRelease(op.readbackBuffer); +struct Context { + std::shared_ptr errors = + std::make_shared(); + wgpu::Instance instance; + wgpu::Adapter adapter; + wgpu::Device device; + wgpu::Queue queue; + + Context() = default; + Context(const Context &) = delete; + Context &operator=(const Context &) = delete; + Context(Context &&) = default; + Context &operator=(Context &&) = default; + + void check() { +#if !defined(__EMSCRIPTEN__) + instance.ProcessEvents(); +#endif + if (auto message = errors->take(); !message.empty()) + throw std::runtime_error(message); } -} +}; -/** - * @brief Overload of the toCPU function to copy data from a GPU buffer to CPU - * memory for an array of floats instead of a pointer to a float buffer. - * @param[in] ctx Context instance to manage the operation - * @param[in] tensor Tensor instance representing the GPU buffer to copy from - * @param[out] data Array of floats to copy the data to - * - * @code - * toCPU(ctx, tensor, data); - * @endcode - */ -template -void toCPU(Context &ctx, Tensor &tensor, std::array &data) { - toCPU(ctx, tensor, data.data(), sizeof(data)); -} +class Future { + std::future result; + wgpu::Future event; -inline void toCPU(Context &ctx, WGPUBuffer buffer, void *data, size_t size) { - uint64_t bufferSize = size; - CopyData op; - op.future = op.promise.get_future(); - { - WGPUBufferDescriptor readbackBufferDescriptor = { - .label = {.data = nullptr, .length = 0}, - .usage = WGPUBufferUsage_CopyDst | WGPUBufferUsage_MapRead, - .size = bufferSize, - }; - op.readbackBuffer = - wgpuDeviceCreateBuffer(ctx.device, &readbackBufferDescriptor); - } - { - WGPUCommandEncoder commandEncoder; - commandEncoder = wgpuDeviceCreateCommandEncoder(ctx.device, nullptr); - wgpuCommandEncoderCopyBufferToBuffer(commandEncoder, buffer, 0, - op.readbackBuffer, 0, bufferSize); - op.commandBuffer = wgpuCommandEncoderFinish(commandEncoder, nullptr); - wgpuCommandEncoderRelease(commandEncoder); - check(op.commandBuffer, "Create command buffer", __FILE__, __LINE__); - } - wgpuQueueSubmit(ctx.queue, 1, &op.commandBuffer); - wgpuCommandBufferRelease(op.commandBuffer); - CallbackData callbackData = {op.readbackBuffer, bufferSize, data, &op.promise, - &op.future}; - - WGPUQueueWorkDoneCallbackInfo workDoneCallbackInfo = { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = - [](WGPUQueueWorkDoneStatus status, void *userdata1, void *userdata2) { - check(status == WGPUQueueWorkDoneStatus_Success, "Queue work done", - __FILE__, __LINE__); - const auto *data = static_cast(userdata1); - WGPUBufferMapCallbackInfo mapCallbackInfo = { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = - [](WGPUMapAsyncStatus status, WGPUStringView message, - void *userdata1, void *userdata2) { - const auto *data = static_cast(userdata1); - check(status == WGPUMapAsyncStatus_Success, - "Map readbackBuffer", __FILE__, __LINE__); - const void *mappedData = wgpuBufferGetConstMappedRange( - data->buffer, /*offset=*/0, data->bufferSize); - check(mappedData, "Get mapped range", __FILE__, __LINE__); - memcpy(data->output, mappedData, data->bufferSize); - wgpuBufferUnmap(data->buffer); - data->promise->set_value(); - }, - .userdata1 = const_cast(data), - .userdata2 = nullptr}; - wgpuBufferMapAsync(data->buffer, WGPUMapMode_Read, 0, - data->bufferSize, mapCallbackInfo); - }, - .userdata1 = &callbackData, - .userdata2 = nullptr}; - wgpuQueueOnSubmittedWorkDone(ctx.queue, workDoneCallbackInfo); - - wait(ctx, op.future); - if (op.readbackBuffer) { - wgpuBufferRelease(op.readbackBuffer); - } -} +public: + Future() = default; + Future(std::future result, wgpu::Future event) + : result(std::move(result)), event(event) {} + Future(const Future &) = delete; + Future &operator=(const Future &) = delete; + Future(Future &&) = default; + Future &operator=(Future &&) = default; -/** - * @brief Copies data from CPU memory to a GPU buffer. The toGPU overloads are - * effectively a convenience wrapper around the WebGPU API call - * wgpuQueueWriteBuffer. - * - * @param[in] ctx Context instance to manage the operation - * @param[in] data Pointer to the CPU memory to copy from - * @param[in] buffer WGPUBuffer instance representing the GPU buffer to copy - * to - * @param[in] size Size of the data buffer in bytes - * - * @code - * toGPU(ctx, data, buffer, size); - * @endcode - */ -inline void toGPU(Context &ctx, const void *data, WGPUBuffer buffer, - size_t size) { - wgpuQueueWriteBuffer(ctx.queue, buffer, 0, data, size); -} + friend void wait(Context &, Future &); +}; -/** - * @brief Overload of the toGPU function to copy data from CPU memory to a GPU - * taking a Tensor instance instead of a WGPUBuffer instance. - * @param[in] ctx Context instance to manage the operation - * @param[in] data Pointer to the CPU memory to copy from - * @param[in] tensor Tensor instance representing the GPU buffer to copy to - * - * @code - * toGPU(ctx, data, tensor); - * @endcode - */ -inline void toGPU(Context &ctx, const float *data, Tensor &tensor) { - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, - tensor.data.size); -} +inline Context createContext(const ContextOptions &options = {}) { + Context context; + std::vector instanceFeatures; +#if defined(__EMSCRIPTEN__) + if (options.enableSPIRV) + throw std::invalid_argument("SPIR-V input is unavailable in browsers"); + instanceFeatures.push_back(wgpu::InstanceFeatureName::TimedWaitAny); +#else + if (options.enableSPIRV) + instanceFeatures.push_back(wgpu::InstanceFeatureName::ShaderSourceSPIRV); +#endif -inline void toGPU(Context &ctx, const half *data, Tensor &tensor) { - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, - tensor.data.size); + wgpu::InstanceDescriptor instanceDescriptor{}; + instanceDescriptor.requiredFeatureCount = instanceFeatures.size(); + instanceDescriptor.requiredFeatures = instanceFeatures.data(); + context.instance = wgpu::CreateInstance(&instanceDescriptor); + if (!context.instance) throw std::runtime_error("could not create WebGPU instance"); + + wgpu::RequestAdapterOptions adapterOptions{}; + adapterOptions.featureLevel = options.featureLevel; + adapterOptions.powerPreference = options.powerPreference; + adapterOptions.backendType = options.backend; + auto adapterPromise = std::make_shared>(); + auto adapterFuture = adapterPromise->get_future(); + auto adapterEvent = context.instance.RequestAdapter( + &adapterOptions, detail::completionMode(), + [adapterPromise](wgpu::RequestAdapterStatus status, wgpu::Adapter adapter, + wgpu::StringView message) { + if (status == wgpu::RequestAdapterStatus::Success) + adapterPromise->set_value(std::move(adapter)); + else + adapterPromise->set_exception(std::make_exception_ptr( + std::runtime_error("could not request WebGPU adapter: " + + detail::string(message)))); + }); + context.adapter = + detail::await(context.instance, adapterFuture, adapterEvent); + + wgpu::DeviceDescriptor deviceDescriptor{}; + deviceDescriptor.requiredFeatureCount = options.requiredFeatures.size(); + deviceDescriptor.requiredFeatures = options.requiredFeatures.data(); + deviceDescriptor.requiredLimits = options.requiredLimits + ? &*options.requiredLimits + : nullptr; + deviceDescriptor.SetUncapturedErrorCallback( + [](const wgpu::Device &, wgpu::ErrorType, wgpu::StringView message, + detail::ErrorState *errors) { + errors->set(detail::string(message)); + }, context.errors.get()); + deviceDescriptor.SetDeviceLostCallback( + detail::persistentMode(), + [](const wgpu::Device &, wgpu::DeviceLostReason reason, + wgpu::StringView message, detail::ErrorState *errors) { + if (reason != wgpu::DeviceLostReason::Destroyed) + errors->set("WebGPU device lost: " + detail::string(message)); + }, context.errors.get()); + + auto devicePromise = std::make_shared>(); + auto deviceFuture = devicePromise->get_future(); + auto deviceEvent = context.adapter.RequestDevice( + &deviceDescriptor, detail::completionMode(), + [devicePromise](wgpu::RequestDeviceStatus status, wgpu::Device device, + wgpu::StringView message) { + if (status == wgpu::RequestDeviceStatus::Success) + devicePromise->set_value(std::move(device)); + else + devicePromise->set_exception(std::make_exception_ptr( + std::runtime_error("could not request WebGPU device: " + + detail::string(message)))); + }); + context.device = detail::await(context.instance, deviceFuture, deviceEvent); + context.queue = context.device.GetQueue(); + return context; } -inline void toGPU(Context &ctx, const int *data, Tensor &tensor) { - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, - tensor.data.size); +inline Tensor createTensor(Context &context, const Shape &shape, NumType type) { + Tensor tensor{.shape = shape, .type = type, + .bytes = size(shape) * sizeBytes(type)}; + if (!tensor.bytes) throw std::invalid_argument("tensor cannot be empty"); + wgpu::BufferDescriptor descriptor{}; + descriptor.label = "gpu.cpp tensor"; + descriptor.size = tensor.bytes; + descriptor.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | + wgpu::BufferUsage::CopyDst; + tensor.buffer = context.device.CreateBuffer(&descriptor); + context.check(); + return tensor; } -inline void toGPU(Context &ctx, const float *data, Tensor &tensor, - size_t size) { - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, size); +template +Tensor createTensor(Context &context, const Shape &shape, NumType type, + std::span values) { + if (sizeof(T) != sizeBytes(type)) + throw std::invalid_argument("host and tensor element sizes differ"); + Tensor tensor{.shape = shape, .type = type, + .bytes = size(shape) * sizeBytes(type)}; + if (!tensor.bytes) throw std::invalid_argument("tensor cannot be empty"); + if (values.size_bytes() != tensor.bytes) + throw std::invalid_argument("host data size does not match tensor shape"); + wgpu::BufferDescriptor descriptor{}; + descriptor.label = "gpu.cpp initialized tensor"; + descriptor.size = tensor.bytes; + descriptor.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | + wgpu::BufferUsage::CopyDst; + descriptor.mappedAtCreation = true; + tensor.buffer = context.device.CreateBuffer(&descriptor); + auto *mapped = tensor.buffer.GetMappedRange(0, tensor.bytes); + if (!mapped) throw std::runtime_error("could not map initialized tensor"); + std::memcpy(mapped, values.data(), tensor.bytes); + tensor.buffer.Unmap(); + context.check(); + return tensor; } -inline void toGPU(Context &ctx, const half *data, Tensor &tensor, size_t size) { - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, size); +template +Tensor createTensor(Context &context, const Shape &shape, NumType type, + const std::vector &values) { + return createTensor(context, shape, type, std::span(values)); } -inline void toGPU(Context &ctx, const int *data, Tensor &tensor, size_t size) { - wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data, size); +template +Tensor createTensor(Context &context, const Shape &shape, NumType type, + const std::array &values) { + return createTensor(context, shape, type, std::span(values)); } -template -inline void toGPU(Context &ctx, Params ¶ms, Kernel &op) { - // TODO(avh): Maintain params metadata in Kernel and check for consistency. - // If a kernel does not have parameters this will quietly overwrite - // the last buffer in the bind group with the parameters buffer. - if (op->numBindings > 0) { - wgpuQueueWriteBuffer(ctx.queue, op->buffers[op->numBindings - 1], 0, - static_cast(¶ms), sizeof(params)); - } +template +void toGPU(Context &context, const Tensor &tensor, std::span values) { + if (sizeof(T) != sizeBytes(tensor.type)) + throw std::invalid_argument("host and tensor element sizes differ"); + if (values.size_bytes() != tensor.bytes) + throw std::invalid_argument("host data size does not match tensor"); + context.queue.WriteBuffer(tensor.buffer, 0, values.data(), values.size_bytes()); + context.check(); } -/** - * @brief Resets the command buffer in preparation for a kernel dispatch. - * Since command buffers are consumed upon submission, this function is used - * both in the initial kernel creation and every time the kernel is to be - * reused for a dispatch. - * @param[in] device WGPUDevice instance to manage the operation - * @param[in] op Kernel instance representing the kernel to reset - * - * @code - * resetCommandBuffer(device, op); - * @endcode - */ -inline void resetCommandBuffer(WGPUDevice &device, Kernel &op) { - { - WGPUCommandEncoder commandEncoder = - wgpuDeviceCreateCommandEncoder(device, nullptr); - WGPUComputePassEncoder computePassEncoder = - wgpuCommandEncoderBeginComputePass(commandEncoder, nullptr); - wgpuComputePassEncoderSetPipeline(computePassEncoder, op->computePipeline); - wgpuComputePassEncoderSetBindGroup(computePassEncoder, 0, op->bindGroup, 0, - nullptr); - wgpuComputePassEncoderDispatchWorkgroups( - computePassEncoder, op->totalWorkgroups[0], op->totalWorkgroups[1], - op->totalWorkgroups[2]); - wgpuComputePassEncoderEnd(computePassEncoder); - wgpuComputePassEncoderRelease(computePassEncoder); - op->commandBuffer = wgpuCommandEncoderFinish(commandEncoder, nullptr); - wgpuCommandEncoderRelease(commandEncoder); - op->used = false; - } +template +void toGPU(Context &context, const Tensor &tensor, const std::vector &values) { + toGPU(context, tensor, std::span(values)); } -/** - * @brief NoParam is a no-op type used to indicate that a kernel does not have - * any parameters. - */ -struct NoParam {}; -template constexpr bool IsNoParam = std::is_same_v; - -/** - * @brief Ceiling division. - */ -inline size_t cdiv(size_t n, size_t d) { return (n + d - 1) / d; } - -/** - * @brief cdiv for shape specification. Mostly useful for evenly dividing - * total # threads by workgroup size dimensions. - */ -inline Shape cdiv(Shape total, Shape group) { - assert(total.rank == group.rank); - Shape result; - result.rank = total.rank; - for (size_t dim = 0; dim < total.rank; ++dim) { - result[dim] = cdiv(total[dim], group[dim]); - } - return result; +template +void toGPU(Context &context, const Tensor &tensor, + const std::array &values) { + toGPU(context, tensor, std::span(values)); } -/** - * @brief A factory function to create a kernel on the GPU. The kernel is - * created with the given WGSL code, input tensors, output tensor, and - * optional parameters. - * - * Note that the values of the input tensors are not used here, only the - * reference handles to the underlying buffers as well as the size of the - * buffers. - * - * @param[in] ctx Context instance to manage the kernel - * @param[in] code WGSL code for the kernel - * @param[in] dataBindings Pointer to a span of tensors bound to the kernel - * @param[in] numTensors Number of tensors in the dataBindings span - * @param[in] viewOffsets Pointer to an array of view offsets for the input - * tensors - * @param[in] totalWorkgroups Shape of the workgroup - * @param[in] params Optional parameters for the kernel. If the kernel does - * not have any parameters, use NoParam. This is cast as void* to allow for - * arbitrary types to be passed as parameters. - * @param[in] paramsSize Size of the parameters buffer in bytes. - * @return Kernel instance representing the created kernel - * - * @code - * Kernel kernel = createKernel(ctx, code, dataBindings, numInputs, - * @endcode - * output, nThreads, params, paramsSize); - */ -inline Kernel createKernel(Context& ctx, const KernelCode &code, - const Tensor *dataBindings, size_t numTensors, - const size_t *viewOffsets, - const Shape &totalWorkgroups, - const void *params = nullptr, size_t paramsSize = 0, - CompilationInfo *compilationInfo = nullptr, - const char *cacheKey = nullptr) { - // Create a cache key by the pointer values of the data bindings and the - // kernel code - if (cacheKey != nullptr && - ctx.kernelPool.data.find(cacheKey) != ctx.kernelPool.data.end()) { - LOG(kDefLog, kInfo, "Kernel cache hit"); - return ctx.kernelPool.data[cacheKey]; - } - - assert(totalWorkgroups.rank == 3); - WGPUDevice device = ctx.device; - WGPUQueue queue = ctx.queue; - Kernel op(new RawKernel()); - - // paramIndex is the index into bgLayoutEntries for the parameters buffer If - // there are no parameters for the kernel, paramsSize == 0 and paramIndex is - // effectively undefined (== -1) - size_t paramIndex = -1; - // Note: paramIndex is undefined unless paramsSize > 0 - size_t numBindings = numTensors; - if (paramsSize > 0) { - numBindings++; // parameters buffer - paramIndex = numBindings - 1; // index of the parameters buffer within - // op.buffers, op.bufferSizes and - // bgLayoutEntries - } - op->buffers = std::make_unique(numBindings); - op->bufferSizes = std::make_unique(numBindings); - op->numBindings = numBindings; - std::vector bgLayoutEntries(numBindings); - // Create layout entries for input buffers - for (size_t i = 0; i < numTensors; ++i) { - bgLayoutEntries[i] = WGPUBindGroupLayoutEntry{ - .binding = static_cast(i), - .visibility = WGPUShaderStage_Compute, - .buffer = - WGPUBufferBindingLayout{ - .type = WGPUBufferBindingType_Storage, - .minBindingSize = dataBindings[i].data.size, - }, - }; - } - if (paramsSize > 0) { - LOG(kDefLog, kInfo, "Create layout entry for the params buffer"); - // Create layout entry for the params buffer - bgLayoutEntries[paramIndex] = WGPUBindGroupLayoutEntry{ - .binding = static_cast(paramIndex), - .visibility = WGPUShaderStage_Compute, - .buffer = - WGPUBufferBindingLayout{ - .type = WGPUBufferBindingType_Uniform, - .minBindingSize = paramsSize, - }, - }; - } - WGPUBindGroupLayoutDescriptor bgLayoutDesc = { - .entryCount = static_cast(bgLayoutEntries.size()), - .entries = bgLayoutEntries.data(), - }; - WGPUBindGroupLayout bgLayout = - wgpuDeviceCreateBindGroupLayout(device, &bgLayoutDesc); - for (size_t i = 0; i < numTensors; ++i) { - op->buffers[i] = dataBindings[i].data.buffer; - op->bufferSizes[i] = dataBindings[i].data.size; - } - // Create a buffer for the Params struct - if (paramsSize > 0) { - WGPUBufferDescriptor paramsBufferDesc = { - .label = {.data = nullptr, .length = 0}, - .usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst, - .size = paramsSize, - .mappedAtCreation = false, - }; - op->buffers[paramIndex] = wgpuDeviceCreateBuffer(device, ¶msBufferDesc); - op->bufferSizes[paramIndex] = paramsSize; - wgpuQueueWriteBuffer(queue, op->buffers[paramIndex], 0, params, paramsSize); - LOG(kDefLog, kTrace, "Params buffer written"); - } else { - LOG(kDefLog, kTrace, "No params buffer needed"); - } - std::vector bindGroupEntries(numBindings); - for (size_t i = 0; i < numTensors; ++i) { - bindGroupEntries[i] = WGPUBindGroupEntry{ - .binding = static_cast(i), - .buffer = op->buffers[i], - .offset = viewOffsets[i], - .size = op->bufferSizes[i], - }; - } - if (paramsSize > 0) { - LOG(kDefLog, kInfo, "Create bind group entry for the params buffer"); - LOG(kDefLog, kInfo, "paramIndex: %d", paramIndex); - bindGroupEntries[paramIndex] = WGPUBindGroupEntry{ - .binding = static_cast(paramIndex), - .buffer = op->buffers[paramIndex], - .offset = 0, - .size = paramsSize, - }; - } - LOG(kDefLog, kTrace, "BG Entries Size: %d", numBindings); - WGPUBindGroupDescriptor bindGroupDesc = { - .layout = bgLayout, - .entryCount = static_cast(numBindings), - .entries = bindGroupEntries.data(), - }; - op->bindGroup = wgpuDeviceCreateBindGroup(device, &bindGroupDesc); - - WGPUPipelineLayoutDescriptor pipelineLayoutDesc = { - .bindGroupLayoutCount = 1, - .bindGroupLayouts = &bgLayout, - }; - WGPUPipelineLayout pipelineLayout = - wgpuDeviceCreatePipelineLayout(device, &pipelineLayoutDesc); - - WGPUShaderSourceWGSL wgslDesc = { - .chain = {.sType = WGPUSType_ShaderSourceWGSL}, - .code = {.data = code.data.c_str(), .length = code.data.length()}}; - - WGPUShaderModuleDescriptor shaderModuleDesc = {}; - shaderModuleDesc.nextInChain = &wgslDesc.chain; - shaderModuleDesc.label = {code.label.c_str(), code.label.length()}; - - WGPUComputePipelineDescriptor computePipelineDesc = {}; - computePipelineDesc.layout = pipelineLayout; - computePipelineDesc.compute.module = - wgpuDeviceCreateShaderModule(device, &shaderModuleDesc); - - computePipelineDesc.compute.entryPoint = {code.entryPoint.c_str(), - code.entryPoint.length()}; - computePipelineDesc.label = {code.label.c_str(), code.label.length()}; - - op->computePipeline = - wgpuDeviceCreateComputePipeline(device, &computePipelineDesc); - op->totalWorkgroups = {totalWorkgroups[0], totalWorkgroups[1], - totalWorkgroups[2]}; - resetCommandBuffer(device, op); - if (cacheKey != nullptr) - ctx.kernelPool.data[cacheKey] = op; - - auto compilationInfoCallback = [](WGPUCompilationInfoRequestStatus status, - WGPUCompilationInfo const *compilationInfo, - void *userdata1, void *userdata2) { - CompilationInfo *result = static_cast(userdata1); - if (compilationInfo && result) { - result->status = status; - for (uint32_t i = 0; i < compilationInfo->messageCount; ++i) { - printf("Message %d: %.*s\n", i, - static_cast(compilationInfo->messages[i].message.length), - compilationInfo->messages[i].message.data); - result->messages.push_back( - std::string(compilationInfo->messages[i].message.data, - compilationInfo->messages[i].message.length)); - result->lineNums.push_back(compilationInfo->messages[i].lineNum); - result->linePos.push_back(compilationInfo->messages[i].linePos); - } - result->finished = true; - } else { - LOG(kDefLog, kTrace, "No compilation info or result"); - } - }; - - WGPUCompilationInfoCallbackInfo compilationCallbackInfo = { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = compilationInfoCallback, - .userdata1 = static_cast(compilationInfo), - .userdata2 = nullptr}; - - while (compilationInfo && !compilationInfo->finished) { - processEvents(ctx.instance); - } - return op; -} +struct Kernel { + wgpu::ComputePipeline pipeline; + wgpu::BindGroup bindGroup; + Shape workgroups{1, 1, 1}; + wgpu::Buffer parameters; + size_t parameterBytes = 0; +}; -/** - * @brief Overload which wraps the createKernel factory function to create a - * kernel on the GPU. This overload uses takes a static collection of input - * tensors instead of a pointer and a statically determined ParamsType instead - * of casting params to a void pointer. - * - * @param[in] ctx Context instance to manage the kernel - * @param[in] code WGSL code for the kernel - * @param[in] dataBindings A Bindings of tensors whose GPU buffers are bound - * to the kernel as inputs and outputs. - * @param[in] totalWorkgroups Number of workgroups in the x, y, z grid, must be - * a Shape of rank == 3. - * @param[in] params Optional parameters for the kernel. If the kernel does - * not have any parameters, use NoParam. - * @return Kernel instance representing the created kernel - * - * @code - * Kernel kernel = createKernel(ctx, code, tensorData, output, - * @endcode - * totalWorkgroups, params); - */ -template -Kernel createKernel(Context &ctx, const KernelCode &code, - const Bindings &dataBindings, - const Shape &totalWorkgroups, - const ParamsType ¶ms = ParamsType{}, - CompilationInfo *compilationInfo = nullptr, - const char *cacheKey = nullptr) { - if constexpr (!IsNoParam) { - return createKernel(ctx, code, dataBindings.data.data(), numInputs, - dataBindings.viewOffsets.data(), totalWorkgroups, - reinterpret_cast(¶ms), - sizeof(ParamsType), compilationInfo, cacheKey); - } else { - return createKernel(ctx, code, dataBindings.data.data(), numInputs, - dataBindings.viewOffsets.data(), totalWorkgroups, - nullptr, 0, compilationInfo, cacheKey); - } +namespace detail { + +inline wgpu::ShaderModule createShaderModule(Context &context, + const Shader &shader) { + return std::visit( + [&](const auto &source) -> wgpu::ShaderModule { + using T = std::decay_t; + wgpu::ShaderModuleDescriptor descriptor{}; + descriptor.label = std::string_view(source.label); + if constexpr (std::same_as) { + wgpu::ShaderSourceWGSL chained{wgpu::ShaderSourceWGSL::Init{ + nullptr, std::string_view(source.code)}}; + descriptor.nextInChain = &chained; + return context.device.CreateShaderModule(&descriptor); + } else { +#if defined(__EMSCRIPTEN__) + throw std::invalid_argument("SPIR-V input is unavailable in browsers"); +#else + if (source.code.size() > UINT32_MAX) + throw std::invalid_argument("SPIR-V module is too large"); + wgpu::ShaderSourceSPIRV chained{wgpu::ShaderSourceSPIRV::Init{ + nullptr, static_cast(source.code.size()), + source.code.data()}}; + descriptor.nextInChain = &chained; + return context.device.CreateShaderModule(&descriptor); +#endif + } + }, + shader); +} + +inline std::string entryPoint(const Shader &shader) { + return std::visit([](const auto &source) { return source.entryPoint; }, shader); +} + +inline std::string label(const Shader &shader) { + return std::visit([](const auto &source) { return source.label; }, shader); +} + +inline Kernel createKernel(Context &context, const Shader &shader, + std::span bindings, + const Shape &workgroups, + std::span parameters) { + if (workgroups.rank != 3) + throw std::invalid_argument("dispatch size must have three dimensions"); + + context.device.PushErrorScope(wgpu::ErrorFilter::Validation); + auto module = createShaderModule(context, shader); + + std::vector layoutEntries(bindings.size()); + std::vector groupEntries(bindings.size()); + for (size_t i = 0; i < bindings.size(); ++i) { + layoutEntries[i].binding = i; + layoutEntries[i].visibility = wgpu::ShaderStage::Compute; + layoutEntries[i].buffer.type = bindings[i].type; + layoutEntries[i].buffer.minBindingSize = bindings[i].bytes; + groupEntries[i].binding = i; + groupEntries[i].buffer = bindings[i].buffer; + groupEntries[i].offset = bindings[i].offset; + groupEntries[i].size = bindings[i].bytes; + } + + Kernel kernel{.workgroups = workgroups}; + if (!parameters.empty()) { + const size_t alignedBytes = (parameters.size() + 15) & ~size_t(15); + wgpu::BufferDescriptor bufferDescriptor{}; + bufferDescriptor.label = "gpu.cpp kernel parameters"; + bufferDescriptor.size = alignedBytes; + bufferDescriptor.usage = + wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst; + kernel.parameters = context.device.CreateBuffer(&bufferDescriptor); + kernel.parameterBytes = parameters.size(); + context.queue.WriteBuffer(kernel.parameters, 0, parameters.data(), + parameters.size()); + + wgpu::BindGroupLayoutEntry layout{}; + layout.binding = layoutEntries.size(); + layout.visibility = wgpu::ShaderStage::Compute; + layout.buffer.type = wgpu::BufferBindingType::Uniform; + layout.buffer.minBindingSize = alignedBytes; + layoutEntries.push_back(layout); + + wgpu::BindGroupEntry entry{}; + entry.binding = groupEntries.size(); + entry.buffer = kernel.parameters; + entry.size = alignedBytes; + groupEntries.push_back(entry); + } + + wgpu::BindGroupLayoutDescriptor bindLayoutDescriptor{}; + bindLayoutDescriptor.label = "gpu.cpp bind group layout"; + bindLayoutDescriptor.entryCount = layoutEntries.size(); + bindLayoutDescriptor.entries = layoutEntries.data(); + auto bindLayout = context.device.CreateBindGroupLayout(&bindLayoutDescriptor); + + wgpu::PipelineLayoutDescriptor pipelineLayoutDescriptor{}; + pipelineLayoutDescriptor.label = "gpu.cpp pipeline layout"; + pipelineLayoutDescriptor.bindGroupLayoutCount = 1; + pipelineLayoutDescriptor.bindGroupLayouts = &bindLayout; + auto pipelineLayout = + context.device.CreatePipelineLayout(&pipelineLayoutDescriptor); + + const auto entry = entryPoint(shader); + const auto pipelineLabel = label(shader); + wgpu::ComputePipelineDescriptor pipelineDescriptor{}; + pipelineDescriptor.label = std::string_view(pipelineLabel); + pipelineDescriptor.layout = pipelineLayout; + pipelineDescriptor.compute.module = module; + pipelineDescriptor.compute.entryPoint = std::string_view(entry); + kernel.pipeline = context.device.CreateComputePipeline(&pipelineDescriptor); + + wgpu::BindGroupDescriptor bindGroupDescriptor{}; + bindGroupDescriptor.label = "gpu.cpp bind group"; + bindGroupDescriptor.layout = bindLayout; + bindGroupDescriptor.entryCount = groupEntries.size(); + bindGroupDescriptor.entries = groupEntries.data(); + kernel.bindGroup = context.device.CreateBindGroup(&bindGroupDescriptor); + + auto errorPromise = std::make_shared>(); + auto errorFuture = errorPromise->get_future(); + auto errorEvent = context.device.PopErrorScope( + detail::completionMode(), + [errorPromise](wgpu::PopErrorScopeStatus status, wgpu::ErrorType type, + wgpu::StringView message) { + if (status != wgpu::PopErrorScopeStatus::Success) + errorPromise->set_value("could not read WebGPU validation result: " + + string(message)); + else if (type != wgpu::ErrorType::NoError) + errorPromise->set_value(string(message)); + else + errorPromise->set_value({}); + }); + if (auto error = await(context.instance, errorFuture, errorEvent); + !error.empty()) + throw std::runtime_error(error); + context.check(); + return kernel; +} + +} // namespace detail + +inline Kernel createKernel( + Context &context, const Shader &shader, std::span bindings, + const Shape &workgroups = {1, 1, 1}, + std::span parameters = {}) { + return detail::createKernel(context, shader, bindings, workgroups, parameters); } -/** - * @brief Asynchronously submits a kernel to the GPU queue for execution. - * It also sets up a callback to notify when the kernel has finished executing - * by setting the value of the promise in the kernel instance argument. - * - * dispatchKernel does *not* wait for the kernel to finish executing and - * returns immediately. The caller can wait for the kernel to finish executing - * by calling wait() on the future in the kernel instance. - * - * @param[in] ctx Context instance to manage the kernel, from which the queue - * for the GPU is obtained - * @param[in] kernel Kernel instance to dispatch - * @param[in] promise Promise to set when the kernel has finished executing - * - * @code - * dispatchKernel(ctx, kernel); - * @endcode - */ -inline void dispatchKernel(Context &ctx, Kernel &kernel, - std::promise &promise) { - if (kernel->used) { - resetCommandBuffer(ctx.device, kernel); - } - wgpuQueueSubmit(ctx.queue, 1, &kernel->commandBuffer); - wgpuCommandBufferRelease(kernel->commandBuffer); - kernel->used = true; - - WGPUQueueWorkDoneCallbackInfo workDoneCallbackInfo = { - .mode = WGPUCallbackMode_AllowSpontaneous, - .callback = - [](WGPUQueueWorkDoneStatus status, void *userdata1, void *userdata2) { - check(status == WGPUQueueWorkDoneStatus_Success, "Queue work done", - __FILE__, __LINE__); - auto *promise = static_cast *>(userdata1); - promise->set_value(); - }, - .userdata1 = &promise, - .userdata2 = nullptr}; - wgpuQueueOnSubmittedWorkDone(ctx.queue, workDoneCallbackInfo); +template +Kernel createKernel(Context &context, const Shader &shader, + const Bindings &bindings, + const Shape &workgroups = {1, 1, 1}) { + return detail::createKernel(context, shader, bindings.data, workgroups, {}); +} + +template + requires std::is_trivially_copyable_v +Kernel createKernel(Context &context, const Shader &shader, + const Bindings &bindings, const Shape &workgroups, + const Parameters ¶meters) { + const auto *data = reinterpret_cast(¶meters); + return detail::createKernel(context, shader, bindings.data, workgroups, + {data, sizeof(parameters)}); +} + +template + requires std::is_trivially_copyable_v +void toGPU(Context &context, const Parameters ¶meters, Kernel &kernel) { + if (sizeof(parameters) != kernel.parameterBytes) + throw std::invalid_argument("kernel parameter size changed"); + context.queue.WriteBuffer(kernel.parameters, 0, ¶meters, + sizeof(parameters)); + context.check(); +} + +inline Future dispatchKernel(Context &context, const Kernel &kernel) { + auto encoder = context.device.CreateCommandEncoder(); + auto pass = encoder.BeginComputePass(); + pass.SetPipeline(kernel.pipeline); + pass.SetBindGroup(0, kernel.bindGroup); + pass.DispatchWorkgroups(kernel.workgroups[0], kernel.workgroups[1], + kernel.workgroups[2]); + pass.End(); + auto commands = encoder.Finish(); + if (!commands) throw std::runtime_error("could not encode WebGPU dispatch"); + context.queue.Submit(1, &commands); + + auto promise = std::make_shared>(); + auto future = promise->get_future(); + auto event = context.queue.OnSubmittedWorkDone( + detail::completionMode(), + [promise](wgpu::QueueWorkDoneStatus status, wgpu::StringView message) { + if (status == wgpu::QueueWorkDoneStatus::Success) + promise->set_value(); + else + promise->set_exception(std::make_exception_ptr(std::runtime_error( + "WebGPU dispatch failed: " + detail::string(message)))); + }); + return {std::move(future), event}; +} + +inline void wait(Context &context, Future &future) { + detail::await(context.instance, future.result, future.event); + context.check(); +} + +template +Future toCPU(Context &context, const Tensor &tensor, std::span output) { + if (sizeof(T) != sizeBytes(tensor.type)) + throw std::invalid_argument("host and tensor element sizes differ"); + if (output.size_bytes() != tensor.bytes) + throw std::invalid_argument("host output size does not match tensor"); + + wgpu::BufferDescriptor descriptor{}; + descriptor.label = "gpu.cpp readback"; + descriptor.size = tensor.bytes; + descriptor.usage = wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead; + auto readback = context.device.CreateBuffer(&descriptor); + auto encoder = context.device.CreateCommandEncoder(); + encoder.CopyBufferToBuffer(tensor.buffer, 0, readback, 0, tensor.bytes); + auto commands = encoder.Finish(); + if (!commands) throw std::runtime_error("could not encode WebGPU readback"); + context.queue.Submit(1, &commands); + + auto promise = std::make_shared>(); + auto future = promise->get_future(); + auto event = readback.MapAsync( + wgpu::MapMode::Read, 0, tensor.bytes, + detail::completionMode(), + [promise, readback, output](wgpu::MapAsyncStatus status, + wgpu::StringView message) { + if (status != wgpu::MapAsyncStatus::Success) { + promise->set_exception(std::make_exception_ptr(std::runtime_error( + "WebGPU readback failed: " + detail::string(message)))); + return; + } + std::memcpy(static_cast(output.data()), + readback.GetConstMappedRange(0, output.size_bytes()), + output.size_bytes()); + readback.Unmap(); + promise->set_value(); + }); + return {std::move(future), event}; +} + +template +Future toCPU(Context &context, const Tensor &tensor, std::vector &output) { + return toCPU(context, tensor, std::span(output)); +} + +template +Future toCPU(Context &context, const Tensor &tensor, + std::array &output) { + return toCPU(context, tensor, std::span(output)); } } // namespace gpu -#endif // GPU_H +#endif diff --git a/numeric_types/half.cpp b/numeric_types/half.cpp deleted file mode 100644 index e5bdaf0..0000000 --- a/numeric_types/half.cpp +++ /dev/null @@ -1,279 +0,0 @@ -#include -#include -#include -#include -#include -#include - -#include "gpu.hpp" -#include "numeric_types/half.hpp" - -using namespace gpu; -using std::isinf; -using std::isnan; - -#define EPSILON 0.01f -#define COLOR_RESET "\033[0m" -#define COLOR_RED "\033[31m" -#define COLOR_GREEN "\033[32m" - -int approximatelyEqual(float a, float b, float epsilon) { - return fabsf(a - b) <= epsilon; -} - -void printResult(bool passed, const char *message, float input, float output) { - if (passed) { - printf("[" COLOR_GREEN "PASSED" COLOR_RESET "]" - " : %s (in: %.10f, out: %.10f)\n", - message, input, output); - } else { - printf("[" COLOR_RED "FAILED" COLOR_RESET "]" - " : %s (in: %.10f, out: %.10f)\n", - message, input, output); - } -} - -void printResult(bool passed, const char *message, float input, - uint16_t output) { - if (passed) { - printf("[" COLOR_GREEN "PASSED" COLOR_RESET "]" - " : %s (input: %.10f, output: 0x%04x)\n", - message, input, output); - } else { - printf("[" COLOR_RED "FAILED" COLOR_RESET "]" - " : %s (input: %.10f, output: 0x%04x)\n", - message, input, output); - } -} - -void testRoundTrip(float value) { - // half h = halfFromFloat(value); - half h = half(value); - float result = static_cast(h); - char message[1024]; - - if (isnan(value)) { - snprintf(message, sizeof(message), "NaN correctly round tripped"); - printResult(isnan(result), message, value, result); - } else if (isinf(value)) { - snprintf(message, sizeof(message), "Infinity correctly round tripped"); - printResult(isinf(result) && - ((value > 0 && result > 0) || (value < 0 && result < 0)), - message, value, result); - } else { - snprintf(message, sizeof(message), "%.10f correctly round tripped", value); - printResult(approximatelyEqual(result, value, EPSILON), message, value, - result); - } -} - -void testRoundTrip(uint16_t value) { - half h; - h.data = value; - float f = halfToFloat(h); - half result = halfFromFloat(f); - char message[512]; - snprintf(message, sizeof(message), "half 0x%04x correctly round tripped", value); - printResult(result.data == value, message, (float)value, result.data); -} - -void testRoundTrip(half value) { - float f = static_cast(value); - half result = half(f); - char message[512]; - snprintf(message, sizeof(message), "half 0x%04x correctly round tripped", value.data); - printResult(result.data == value.data, message, (float)value, result.data); -} - -void testSpecialCases() { - half h; - char message[512]; - - // Zero - h.data = 0x0000; - snprintf(message, sizeof(message), "0x0000 correctly converted to 0.0f"); - printResult(halfToFloat(h) == 0.0f, message, 0.0f, halfToFloat(h)); - - // Negative zero - h.data = 0x8000; - snprintf(message, sizeof(message), "0x8000 correctly converted to -0.0f"); - printResult(halfToFloat(h) == -0.0f, message, -0.0f, halfToFloat(h)); - - // Infinity - h.data = 0x7c00; - snprintf(message, sizeof(message), "0x7c00 correctly converted to positive infinity"); - printResult(isinf(halfToFloat(h)) && halfToFloat(h) > 0, message, INFINITY, - halfToFloat(h)); - - // Negative infinity - h.data = 0xfc00; - snprintf(message, sizeof(message), "0xfc00 correctly converted to negative infinity"); - printResult(isinf(halfToFloat(h)) && halfToFloat(h) < 0, message, -INFINITY, - halfToFloat(h)); - - // NaN - h.data = 0x7e00; - snprintf(message, sizeof(message), "0x7e00 correctly converted to NaN"); - printResult(isnan(halfToFloat(h)), message, NAN, halfToFloat(h)); - - // Smallest positive normal number - h.data = 0x0400; - snprintf(message, sizeof(message), "0x0400 correctly converted to 6.10352e-05f"); - printResult(approximatelyEqual(halfToFloat(h), 6.10352e-05f, EPSILON), - message, 6.10352e-05f, halfToFloat(h)); - - // Largest denormalized number - h.data = 0x03ff; - snprintf(message, sizeof(message), "0x03ff correctly converted to 6.09756e-05f"); - printResult(approximatelyEqual(halfToFloat(h), 6.09756e-05f, EPSILON), - message, 6.09756e-05f, halfToFloat(h)); - - // Smallest positive denormalized number - h.data = 0x0001; - snprintf(message, sizeof(message), "0x0001 correctly converted to 5.96046e-08f"); - printResult(approximatelyEqual(halfToFloat(h), 5.96046e-08f, EPSILON), - message, 5.96046e-08f, halfToFloat(h)); - - // Zero - h = halfFromFloat(0.0f); - snprintf(message, sizeof(message), "0.0f correctly converted to 0x0000"); - printResult(h.data == 0x0000, message, 0.0f, h.data); - - // Negative zero - h = halfFromFloat(-0.0f); - snprintf(message, sizeof(message), "-0.0f correctly converted to 0x8000"); - printResult(h.data == 0x8000, message, -0.0f, h.data); - - // Infinity - h = halfFromFloat(INFINITY); - snprintf(message, sizeof(message), "positive infinity correctly converted to 0x7c00"); - printResult(h.data == 0x7c00, message, INFINITY, h.data); - - // Negative infinity - h = halfFromFloat(-INFINITY); - snprintf(message, sizeof(message), "negative infinity correctly converted to 0xfc00"); - printResult(h.data == 0xfc00, message, -INFINITY, h.data); - - // NaN - h = halfFromFloat(NAN); - snprintf(message, sizeof(message), "NaN correctly converted to NaN representation"); - printResult((h.data & 0x7c00) == 0x7c00 && (h.data & 0x03ff) != 0x0000, - message, NAN, h.data); - - // Smallest positive normal number - h = halfFromFloat(6.10352e-05f); - snprintf(message, sizeof(message), "6.10352e-05f correctly converted to 0x0400"); - printResult(h.data == 0x0400, message, 6.10352e-05f, h.data); - - // Largest denormalized number - h = halfFromFloat(6.09756e-05f); - snprintf(message, sizeof(message), "6.09756e-05f correctly converted to 0x03ff"); - printResult(h.data == 0x03ff, message, 6.09756e-05f, h.data); - - // Smallest positive denormalized number - h = halfFromFloat(5.96046448e-08f); - snprintf(message, sizeof(message), "5.96046448e-08f correctly converted to 0x0001"); - printResult(h.data == 0x0001, message, 5.96046e-08f, h.data); -} - -void testContainers() { - { - std::array h = {0.0f, -0.0f, INFINITY, NAN}; - testRoundTrip(h[0]); - testRoundTrip(h[1]); - testRoundTrip(h[2]); - testRoundTrip(h[3]); - } - { - Context ctx = createContext(); - std::array h = {1.0f, 0.5f, 2.0f, 3.14f, 1.0, 2.0, 3.0, 4.0}; - Tensor devH = createTensor(ctx, {h.size()}, kf16, h.data()); - std::array h2; - toCPU(ctx, devH, h2.data(), sizeof(h2)); - for (int i = 0; i < 8; ++i) { - printResult(h[i].data == h2[i].data, "Container round trip", - static_cast(h[i]), static_cast(h2[i])); - } - } -} - -void testWGSL() { - static const char *kGelu = R"( -const GELU_SCALING_FACTOR: f16 = 0.7978845608028654; // sqrt(2.0 / PI) -@group(0) @binding(0) var inp: array<{{precision}}>; -@group(0) @binding(1) var out: array<{{precision}}>; -@group(0) @binding(1) var dummy: array<{{precision}}>; -@compute @workgroup_size({{workgroupSize}}) -fn main( - @builtin(global_invocation_id) GlobalInvocationID: vec3) { - let i: u32 = GlobalInvocationID.x; - if (i < arrayLength(&inp)) { - let x: f16 = inp[i]; - out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR - * (x + .044715 * x * x * x))), x, x > 10.0); - } -} -)"; - Context ctx = createContext( - {}, {}, - /*device descriptor, enabling f16 in WGSL*/ - { - .requiredFeatureCount = 1, - .requiredFeatures = std::array{WGPUFeatureName_ShaderF16}.data(), - }); - static constexpr size_t N = 10000; - std::array inputArr, outputArr; - for (int i = 0; i < N; ++i) { - inputArr[i] = half(static_cast(i) / 10.0f); // dummy input data - } - Tensor input = createTensor(ctx, Shape{N}, kf16, inputArr.data()); - Tensor output = createTensor(ctx, Shape{N}, kf16); - std::promise promise; - std::future future = promise.get_future(); - Kernel op = createKernel(ctx, {kGelu, 256, kf16}, Bindings{input, output}, - {cdiv(N, 256), 1, 1}); - dispatchKernel(ctx, op, promise); - wait(ctx, future); - toCPU(ctx, output, outputArr.data(), sizeof(outputArr)); - for (int i = 0; i < 12; ++i) { - printf(" gelu(%.2f) = %.2f\n", static_cast(inputArr[i]), - static_cast(outputArr[i])); - } -} - -int main() { - printf("\nHalf-precision float tests\n==========================\n"); - - printf("\nRegular values float round trips\n\n"); - testRoundTrip(1.0f); - testRoundTrip(0.5f); - testRoundTrip(2.0f); - testRoundTrip(3.14f); - testRoundTrip(-1.0f); - testRoundTrip(-0.5f); - testRoundTrip(-2.0f); - testRoundTrip(-3.14f); - - printf("\nEdge Case float round trips\n\n"); - testRoundTrip(0.0f); - testRoundTrip(-0.0f); - testRoundTrip(INFINITY); - testRoundTrip(-INFINITY); - testRoundTrip(NAN); - // testRoundTrip(FLT_MAX); // since FLT_MAX is not representable as half it - // is not expected to round-trip correctly testRoundTrip(FLT_MIN); - testRoundTrip(FLT_TRUE_MIN); - - printf("\nSpecial half values\n\n"); - testSpecialCases(); - - printf("\nContainers and CPU/GPU round trip\n\n"); - testContainers(); - - printf("\nWGSL f16 extension test\n\n"); - testWGSL(); - - printf("\nTests completed.\n"); - - return 0; -} diff --git a/numeric_types/half.hpp b/numeric_types/half.hpp index f78e61a..20353d8 100644 --- a/numeric_types/half.hpp +++ b/numeric_types/half.hpp @@ -1,321 +1,15 @@ -#ifndef HALF_H -#define HALF_H +#ifndef GPU_CPP_NUMERIC_TYPES_HALF_HPP +#define GPU_CPP_NUMERIC_TYPES_HALF_HPP -#include -#include -#include -#include -#include +#include -#ifdef _MSC_VER -#include - -static inline uint32_t __builtin_clz(uint32_t value) -{ - unsigned long leading_zero = 0; - if (value == 0) - { - return 32; - } - _BitScanReverse(&leading_zero, value); - return 31 - leading_zero; -} - -static inline uint16_t __builtin_clz(uint16_t value) -{ - return __builtin_clz(static_cast(value)) - 16; -} - -static inline uint64_t __builtin_clz(uint64_t value) -{ - unsigned long leading_zero = 0; - if (value == 0) - { - return 64; - } -#if defined(_WIN64) - _BitScanReverse64(&leading_zero, value); - return 63 - leading_zero; -#else - uint32_t high = static_cast(value >> 32); - uint32_t low = static_cast(value); - if (high != 0) - { - return __builtin_clz(high); - } - else - { - return 32 + __builtin_clz(low); - } +#if !defined(__FLT16_MANT_DIG__) +#error "gpu.cpp requires compiler support for IEEE 754 _Float16" #endif -} -#endif - -struct half; -static inline half halfFromFloat(float f); -static inline float halfToFloat(half h); - -/** - * Experimental implementation of half-precision 16-bit floating point numbers. - */ -struct half -{ - uint16_t data; - - // Default constructor - half() : data(0) {} - - // Constructor from float - half(float f) { *this = halfFromFloat(f); } - - // Constructor from uint16_t - explicit half(uint16_t value) : data(value) {} - - operator float() const { return halfToFloat(*this); } - - // Conversion operator to uint16_t - operator uint16_t() const { return data; } - - // Overload assignment operator from uint16_t - half &operator=(uint16_t value) - { - data = value; - return *this; - } - - // Overload assignment operator from another half - half &operator=(const half &other) - { - data = other.data; - return *this; - } - - // Overload assignment operator from float - half &operator=(float value) - { - data = halfFromFloat(value); - return *this; - } -}; - -/** - * @brief Converts a 32-bit float to a 16-bit half-precision float. - * - * Based on Mike Acton's half.c implementation. - */ -half halfFromFloat(float f) -{ - union - { - float f; - uint32_t u; - } floatUnion = {f}; - - uint32_t float32 = floatUnion.u; - - // Constants for bit masks, shifts, and biases - const uint16_t ONE = 0x0001; - const uint32_t FLOAT_SIGN_MASK = 0x80000000; - const uint32_t FLOAT_EXP_MASK = 0x7f800000; - const uint32_t FLOAT_MANTISSA_MASK = 0x007fffff; - const uint32_t FLOAT_HIDDEN_BIT = 0x00800000; - const uint32_t FLOAT_ROUND_BIT = 0x00001000; - const uint16_t FLOAT_EXP_BIAS = 0x007f; - const uint16_t HALF_EXP_BIAS = 0x000f; - const uint16_t FLOAT_SIGN_POS = 0x001f; - const uint16_t HALF_SIGN_POS = 0x000f; - const uint16_t FLOAT_EXP_POS = 0x0017; - const uint16_t HALF_EXP_POS = 0x000a; - const uint16_t HALF_EXP_MASK = 0x7c00; - const uint16_t FLOAT_EXP_FLAGGED_VALUE = 0x00ff; - const uint16_t HALF_EXP_MASK_VALUE = HALF_EXP_MASK >> HALF_EXP_POS; - const uint16_t HALF_EXP_MAX_VALUE = HALF_EXP_MASK_VALUE - ONE; - const uint16_t FLOAT_HALF_SIGN_POS_OFFSET = FLOAT_SIGN_POS - HALF_SIGN_POS; - const uint16_t FLOAT_HALF_BIAS_OFFSET = FLOAT_EXP_BIAS - HALF_EXP_BIAS; - const uint16_t FLOAT_HALF_MANTISSA_POS_OFFSET = FLOAT_EXP_POS - HALF_EXP_POS; - const uint16_t HALF_NAN_MIN = HALF_EXP_MASK | ONE; - - // Extracting the sign, exponent, and mantissa from the 32-bit float - const uint32_t floatSignMasked = float32 & FLOAT_SIGN_MASK; - const uint32_t floatExpMasked = float32 & FLOAT_EXP_MASK; - const uint16_t halfSign = - static_cast(floatSignMasked >> FLOAT_HALF_SIGN_POS_OFFSET); - const uint16_t floatExp = - static_cast(floatExpMasked >> FLOAT_EXP_POS); - const uint32_t floatMantissa = float32 & FLOAT_MANTISSA_MASK; - - // Check for NaN - if ((floatExpMasked == FLOAT_EXP_MASK) && (floatMantissa != 0)) - { - half result; - result.data = - HALF_EXP_MASK | (floatMantissa >> FLOAT_HALF_MANTISSA_POS_OFFSET); - return result; - } - - // Adjusting the exponent and rounding the mantissa - const uint16_t floatExpHalfBias = floatExp - FLOAT_HALF_BIAS_OFFSET; - const uint32_t floatMantissaRoundMask = floatMantissa & FLOAT_ROUND_BIT; - const uint32_t floatMantissaRoundOffset = floatMantissaRoundMask << ONE; - const uint32_t floatMantissaRounded = - floatMantissa + floatMantissaRoundOffset; - - // Handling denormalized numbers - const uint32_t floatMantissaDenormShiftAmount = ONE - floatExpHalfBias; - const uint32_t floatMantissaWithHidden = - floatMantissaRounded | FLOAT_HIDDEN_BIT; - const uint32_t floatMantissaDenorm = - floatMantissaWithHidden >> floatMantissaDenormShiftAmount; - const uint16_t halfMantissaDenorm = static_cast( - floatMantissaDenorm >> FLOAT_HALF_MANTISSA_POS_OFFSET); - const uint16_t halfDenorm = halfSign | halfMantissaDenorm; - // Handling special cases: infinity and NaN - const uint16_t halfInf = halfSign | HALF_EXP_MASK; - const uint16_t mantissaNan = - static_cast(floatMantissa >> FLOAT_HALF_MANTISSA_POS_OFFSET); - const uint16_t halfNan = halfSign | HALF_EXP_MASK | mantissaNan; - const uint16_t halfNanNotInf = halfSign | HALF_NAN_MIN; +using half = _Float16; - // Handling overflow - const uint16_t halfExpNormOverflowOffset = floatExpHalfBias + ONE; - const uint16_t halfExpNormOverflow = halfExpNormOverflowOffset - << HALF_EXP_POS; - const uint16_t halfNormOverflow = halfSign | halfExpNormOverflow; +static_assert(sizeof(half) == 2); +static_assert(std::is_trivially_copyable_v); - // Handling normalized numbers - const uint16_t halfExpNorm = floatExpHalfBias << HALF_EXP_POS; - const uint16_t halfMantissaNorm = static_cast( - floatMantissaRounded >> FLOAT_HALF_MANTISSA_POS_OFFSET); - const uint16_t halfNorm = halfSign | halfExpNorm | halfMantissaNorm; - - // Checks and conditions - const uint16_t halfIsDenorm = FLOAT_HALF_BIAS_OFFSET >= floatExp; - const uint16_t floatHalfExpBiasedFlag = - FLOAT_EXP_FLAGGED_VALUE - FLOAT_HALF_BIAS_OFFSET; - const uint16_t floatExpIsFlagged = floatExpHalfBias == floatHalfExpBiasedFlag; - const uint16_t isFloatMantissaZero = floatMantissa == 0; - const uint16_t isHalfNanZero = mantissaNan == 0; - const uint16_t floatIsInf = floatExpIsFlagged && isFloatMantissaZero; - const uint16_t floatIsNanUnderflow = floatExpIsFlagged && isHalfNanZero; - const uint16_t floatIsNan = floatExpIsFlagged; - const uint16_t expIsOverflow = floatExpHalfBias > HALF_EXP_MAX_VALUE; - const uint32_t floatMantissaRoundedOverflow = - floatMantissaRounded & FLOAT_HIDDEN_BIT; - const uint32_t mantissaNormIsOverflow = floatMantissaRoundedOverflow != 0; - const uint16_t halfIsInf = expIsOverflow || floatIsInf; - - // Selecting final result based on conditions - const uint16_t checkOverflowResult = - mantissaNormIsOverflow ? halfNormOverflow : halfNorm; - const uint16_t checkNanResult = floatIsNan ? halfNan : checkOverflowResult; - const uint16_t checkNanUnderflowResult = - floatIsNanUnderflow ? halfNanNotInf : checkNanResult; - const uint16_t checkInfResult = halfIsInf ? halfInf : checkNanUnderflowResult; - const uint16_t checkDenormResult = halfIsDenorm ? halfDenorm : checkInfResult; - - // Final result after all checks - half result; - result.data = checkDenormResult; - - return result; -} - -/** - * @brief Converts a 16-bit half-precision float to a 32-bit float. - * - * Based on Mike Acton's half.c implementation. - */ -float halfToFloat(half h) -{ - // Constants for bit masks, shifts, and biases - const uint16_t ONE = 0x0001; - const uint16_t TWO = 0x0002; - const uint32_t FLOAT_EXP_MASK = 0x7f800000; - const uint32_t FLOAT_MANTISSA_MASK = 0x007fffff; - const uint16_t FLOAT_EXP_BIAS = 0x007f; - const uint16_t HALF_EXP_BIAS = 0x000f; - const uint16_t HALF_SIGN_MASK = 0x8000; - const uint16_t HALF_EXP_MASK = 0x7c00; - const uint16_t HALF_MANTISSA_MASK = 0x03ff; - const uint16_t HALF_EXP_POS = 0x000a; - const uint16_t FLOAT_EXP_POS = 0x0017; - const uint16_t FLOAT_SIGN_POS = 0x001f; - const uint16_t HALF_SIGN_POS = 0x000f; - const uint16_t HALF_FLOAT_DENORM_SA_OFFSET = 0x000a; - const uint32_t HALF_FLOAT_BIAS_OFFSET = HALF_EXP_BIAS - FLOAT_EXP_BIAS; - const uint16_t HALF_FLOAT_SIGN_POS_OFFSET = FLOAT_SIGN_POS - HALF_SIGN_POS; - const uint16_t HALF_FLOAT_MANTISSA_POS_OFFSET = FLOAT_EXP_POS - HALF_EXP_POS; - - // Extracting the sign, exponent, and mantissa from the 16-bit float - const uint32_t halfSignMasked = h.data & HALF_SIGN_MASK; - const uint32_t halfExpMasked = h.data & HALF_EXP_MASK; - const uint16_t halfMantissa = h.data & HALF_MANTISSA_MASK; - - // Shifting the sign bit to the correct position for the 32-bit float - const uint32_t floatSign = halfSignMasked << HALF_FLOAT_SIGN_POS_OFFSET; - - // Adjusting the exponent - const uint16_t halfExpHalfBias = halfExpMasked >> HALF_EXP_POS; - const uint32_t floatExp = halfExpHalfBias - HALF_FLOAT_BIAS_OFFSET; - - // Shifting the mantissa to the correct position for the 32-bit float - const uint32_t floatMantissa = halfMantissa << HALF_FLOAT_MANTISSA_POS_OFFSET; - - // Checking conditions for zero, denormalized, infinity, and NaN - const uint32_t isExpNonZero = halfExpMasked != 0; - const uint32_t isMantissaNonZero = halfMantissa != 0; - const uint32_t isZero = !(isExpNonZero || isMantissaNonZero); - const uint32_t isDenorm = !isZero && !isExpNonZero; - const uint32_t isExpFlagged = halfExpMasked == HALF_EXP_MASK; - const uint32_t isInf = isExpFlagged && !isMantissaNonZero; - const uint32_t isNan = isExpFlagged && isMantissaNonZero; - - // Handling denormalized numbers - const uint16_t halfMantissaLeadingZeros = __builtin_clz(halfMantissa) - 16; - const uint16_t halfDenormShiftAmount = - halfMantissaLeadingZeros + HALF_FLOAT_DENORM_SA_OFFSET; - const uint32_t halfFloatDenormMantissaShiftAmount = - halfDenormShiftAmount - TWO; - const uint32_t halfFloatDenormMantissa = - halfMantissa << halfFloatDenormMantissaShiftAmount; - const uint32_t floatDenormMantissa = - halfFloatDenormMantissa & FLOAT_MANTISSA_MASK; - const uint32_t halfFloatDenormShiftAmount = ONE - halfDenormShiftAmount; - const uint32_t floatDenormExp = halfFloatDenormShiftAmount + FLOAT_EXP_BIAS; - const uint32_t floatDenormExpPacked = floatDenormExp << FLOAT_EXP_POS; - const uint32_t floatDenorm = - floatSign | floatDenormExpPacked | floatDenormMantissa; - - // Handling special cases: infinity and NaN - const uint32_t floatInf = floatSign | FLOAT_EXP_MASK; - const uint32_t floatNan = floatSign | FLOAT_EXP_MASK | floatMantissa; - - // Handling zero - const uint32_t floatZero = floatSign; - - // Handling normalized numbers - const uint32_t floatExpPacked = floatExp << FLOAT_EXP_POS; - const uint32_t packed = floatSign | floatExpPacked | floatMantissa; - - // Selecting final result based on conditions - const uint32_t checkZeroResult = isZero ? floatZero : packed; - const uint32_t checkDenormResult = isDenorm ? floatDenorm : checkZeroResult; - const uint32_t checkInfResult = isInf ? floatInf : checkDenormResult; - const uint32_t checkNanResult = isNan ? floatNan : checkInfResult; - - // Final result after all checks - const uint32_t result = checkNanResult; - - // Reinterpret the uint32_t result as a float using a union - union - { - uint32_t u; - float f; - } floatUnion; - floatUnion.u = result; - - return floatUnion.f; -} - -#endif // HALF_H +#endif diff --git a/setup.py b/setup.py deleted file mode 100644 index 40cc5cc..0000000 --- a/setup.py +++ /dev/null @@ -1,114 +0,0 @@ -import os -import platform -import sys -import ssl -import urllib.request -from pathlib import Path - -def get_os_name(): - system = platform.system() - if system == "Windows": - return "Windows 64-bit" if platform.machine().endswith('64') else "Windows 32-bit" - elif system == "Darwin": - return "macOS" - elif system == "Linux": - return "Linux" - elif system == "FreeBSD": - return "FreeBSD" - elif system.startswith("CYGWIN"): - return "Cygwin" - else: - return "Other" - -def download_file(url, output_filename): - total_downloaded = 0 - total_truncated = 0 # only print download progress every 2MB to avoid spamming logs - - def report_progress(block_num, block_size, total_size): - nonlocal total_downloaded - nonlocal total_truncated - total_downloaded += block_size - if total_downloaded // (1024 * 1024) > total_truncated: - total_truncated = total_downloaded // (1024 * 1024) - if total_truncated % 2 == 0: - print(f"\rDownloaded {total_downloaded // (1024 * 1024)} MB", end="") - - try: - ssl._create_default_https_context = ssl._create_stdlib_context - urllib.request.urlretrieve(url, output_filename, reporthook=report_progress) - print(f"\nDownloaded {output_filename}") - return True - except Exception as e: - print(f"\nFailed to download {output_filename}") - print(f"Error: {str(e)}") - sys.exit(1) - -def check_os(os_name): - print("\nChecking System") - print("===============\n") - print(f" Operating System : {os_name}") - supported = {"macOS", "Linux"} - if os_name not in supported: - print("Unsupported operating system") - sys.exit(1) - -def download_dawn(os_name): - print("\nDownload Dawn Library") - print("=====================\n") - - outfile_map = { - "macOS": "third_party/lib/libwebgpu_dawn.dylib", - "Linux": "third_party/lib/libwebgpu_dawn.so", - } - url_map = { - "macOS": "https://github.com/austinvhuang/dawn-artifacts/releases/download/0.2.0/libwebgpu_dawn.dylib", - "Linux": "https://github.com/austinvhuang/dawn-artifacts/releases/download/0.2.0/libwebgpu_dawn.so", - } - - outfile = outfile_map.get(os_name) - url = url_map.get(os_name) - - if not outfile or not url: - print(f"No download information for {os_name}") - sys.exit(1) - - print(f" URL : {url}") - print(f" Download File : {outfile}\n") - print(" Downloading ...\n") - - if Path(outfile).exists(): - print(f" File {outfile} already exists, skipping.") - sys.exit(0) - - Path(outfile).parent.mkdir(parents=True, exist_ok=True) - download_file(url, outfile) - -def setup_env(os_name): - print("\nEnvironment Setup") - print("=================\n") - - current_dir = os.getcwd() - lib_dir = os.path.join(current_dir, "third_party", "lib") - - if os_name == "macOS": - print(" Before running the program, run the following command or add it to your shell profile:") - print(f" export DYLD_LIBRARY_PATH=$DYLD_LIBRARY_PATH:{lib_dir}") - - with open("source", "w") as f: - f.write(f"export DYLD_LIBRARY_PATH=$DYLD_LIBRARY_PATH:{lib_dir}\n") - if os_name == "Linux": - print(" Before running the program, run the following command or add it to your shell profile:") - print(f" export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:{lib_dir}") - - with open("source", "w") as f: - f.write(f"export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:{lib_dir}\n") - -def main(): - os_name = get_os_name() - check_os(os_name) - download_dawn(os_name) - setup_env(os_name) - print() - -if __name__ == "__main__": - main() diff --git a/test b/test new file mode 100755 index 0000000..e704673 --- /dev/null +++ b/test @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ ${1:-} == --web || ${1:-} == --build-web || \ + ${1:-} == --rebuild-web ]]; then + if [[ ${1:-} == --rebuild-web || \ + ! -f third_party/emdawnwebgpu/emdawnwebgpu.port.py ]]; then + tools/build_emdawn.sh + fi + export EMSDK_QUIET=1 + source "${GPUCPP_EMSDK_ROOT:-../emsdk}/emsdk_env.sh" >/dev/null + if [[ $(uname -s) == Darwin ]]; then + if [[ -x /opt/homebrew/opt/openjdk/bin/java ]]; then + export PATH="/opt/homebrew/opt/openjdk/bin:$PATH" + fi + if ! java -version >/dev/null 2>&1; then + echo "Closure requires Java on macOS; install it with brew install openjdk" >&2 + exit 1 + fi + fi + emcmake cmake -S . -B build-web -G Ninja -DCMAKE_BUILD_TYPE=Release + cmake --build build-web --target gpu_cpp_web web_binding_test + if [[ ${1:-} != --build-web ]]; then + emrun --browser chrome --timeout 60 build-web/web_binding_test.html + fi + exit +elif [[ ${1:-} == --rebuild-dawn ]]; then + tools/build_dawn.sh +elif [[ $# -ne 0 ]]; then + echo "usage: ./test [--rebuild-dawn|--web|--build-web|--rebuild-web]" >&2 + exit 2 +elif [[ ! -f third_party/dawn/include/dawn/webgpu_cpp.h && + ! -f third_party/local/dawn/build-latest/gen/include/dawn/webgpu_cpp.h ]]; then + tools/build_dawn.sh +fi + +cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release +cmake --build build +ctest --test-dir build --output-on-failure diff --git a/tests/gpu_test.cpp b/tests/gpu_test.cpp new file mode 100644 index 0000000..3b0dc68 --- /dev/null +++ b/tests/gpu_test.cpp @@ -0,0 +1,102 @@ +#include "gpu.hpp" +#include "numeric_types/half.hpp" + +#include +#include +#include +#include + +using namespace gpu; + +static constexpr auto twicePlusOne = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +@compute @workgroup_size({{workgroupSize}}) +fn main(@builtin(global_invocation_id) id: vec3) { + if (id.x < arrayLength(&input)) { + output[id.x] = input[id.x] * 2 + 1; + } +} +)"; + +static constexpr auto doubleF16 = R"( +@group(0) @binding(0) var input: array<{{precision}}>; +@group(0) @binding(1) var output: array<{{precision}}>; + +@compute @workgroup_size({{workgroupSize}}) +fn main(@builtin(global_invocation_id) id: vec3) { + if (id.x < arrayLength(&input)) { + output[id.x] = input[id.x] * 2.0; + } +} +)"; + +static std::vector readSPIRV(const char *path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) throw std::runtime_error("could not open assembled SPIR-V test"); + const auto bytes = file.tellg(); + if (bytes <= 0 || bytes % sizeof(uint32_t)) + throw std::runtime_error("assembled SPIR-V test has an invalid size"); + std::vector code(bytes / sizeof(uint32_t)); + file.seekg(0); + file.read(reinterpret_cast(code.data()), bytes); + return code; +} + +template +static void expect(const std::vector &actual, + const std::vector &expected, const char *story) { + if (actual != expected) throw std::runtime_error(std::string(story) + " failed"); +} + +int main(int argc, char **argv) { + if (argc != 2) throw std::runtime_error("expected assembled SPIR-V path"); + + // Ordinary WGSL covers upload, explicit binding access, dispatch, and readback. + auto context = createContext(); + std::vector input{1, 2, 3, 4}; + std::vector output(input.size()); + auto gpuInput = createTensor(context, {input.size()}, ki32, input); + auto gpuOutput = createTensor(context, {output.size()}, ki32); + auto wgsl = createKernel(context, WGSL{twicePlusOne, 4}, + Bindings{read(gpuInput), readWrite(gpuOutput)}); + auto dispatched = dispatchKernel(context, wgsl); + wait(context, dispatched); + auto downloaded = toCPU(context, gpuOutput, output); + wait(context, downloaded); + expect(output, std::vector{3, 5, 7, 9}, "WGSL compute"); + + // Native IEEE binary16 values round-trip through an f16 WGSL pipeline. + auto f16Context = createContext( + {.requiredFeatures = {wgpu::FeatureName::ShaderF16}}); + std::vector f16Input{half(0.5f), half(-2.0f), half(3.25f), + half(10.0f)}; + std::vector f16Output(f16Input.size()); + auto gpuF16Input = createTensor(f16Context, {f16Input.size()}, kf16, f16Input); + auto gpuF16Output = createTensor(f16Context, {f16Output.size()}, kf16); + auto f16Kernel = + createKernel(f16Context, WGSL{doubleF16, 4, kf16}, + Bindings{read(gpuF16Input), readWrite(gpuF16Output)}); + auto f16Dispatched = dispatchKernel(f16Context, f16Kernel); + wait(f16Context, f16Dispatched); + auto f16Downloaded = toCPU(f16Context, gpuF16Output, f16Output); + wait(f16Context, f16Downloaded); + expect(f16Output, + std::vector{half(1.0f), half(-4.0f), half(6.5f), half(20.0f)}, + "f16 compute"); + + // The same runtime accepts SPIR-V conforming to gpu.cpp's WebGPU profile. + auto spirvContext = createContext({.enableSPIRV = true}); + std::vector answer(1); + auto gpuAnswer = createTensor(spirvContext, {1}, ki32); + auto spirv = createKernel(spirvContext, SPIRV{readSPIRV(argv[1])}, + Bindings{readWrite(gpuAnswer)}); + auto spirvDispatched = dispatchKernel(spirvContext, spirv); + wait(spirvContext, spirvDispatched); + auto answerDownloaded = toCPU(spirvContext, gpuAnswer, answer); + wait(spirvContext, answerDownloaded); + expect(answer, std::vector{42}, "SPIR-V compute"); + + std::cout << "WGSL, f16, and SPIR-V compute stories passed\n"; +} diff --git a/tests/web_binding_test.cpp b/tests/web_binding_test.cpp new file mode 100644 index 0000000..00aa7ca --- /dev/null +++ b/tests/web_binding_test.cpp @@ -0,0 +1,47 @@ +#include + +EM_ASYNC_JS(int, runStory, (), { + try { + const context = await Module.createContext(false); + const input = new Int32Array([1, 2, 3, 4]); + const output = new Int32Array(4); + const spec = { + "code": ` + @group(0) @binding(0) var input: array; + @group(0) @binding(1) var output: array; + @compute @workgroup_size({{workgroupSize}}) + fn main(@builtin(global_invocation_id) id: vec3) { + if (id.x < arrayLength(&input)) { + output[id.x] = input[id.x] * 2 + 1; + } + } + `, + "workgroupSize": [4, 1, 1], + "workgroups": [1, 1, 1], + "bindings": [ + {"data": input, "access": "read"}, + {"data": output, "access": "readWrite"}, + ], + }; + + let rejected = false; + try { + await context.run({...spec, "code": "this is not WGSL"}); + } catch (error) { + rejected = true; + } + if (!rejected) throw new Error("invalid WGSL did not reject the run"); + + await context.run(spec); + context.delete(); + if (output.toString() !== "3,5,7,9") + throw new Error(`unexpected result: ${output}`); + out("Browser JS binding compute story passed"); + return 0; + } catch (error) { + out(`Browser JS binding failed: ${String(error && error.stack || error)}`); + return 1; + } +}); + +int main() { return runStory(); } diff --git a/tests/write42.spvasm b/tests/write42.spvasm new file mode 100644 index 0000000..2a1b828 --- /dev/null +++ b/tests/write42.spvasm @@ -0,0 +1,27 @@ +OpCapability Shader +OpMemoryModel Logical GLSL450 +OpEntryPoint GLCompute %main "main" +OpExecutionMode %main LocalSize 1 1 1 +OpDecorate %output DescriptorSet 0 +OpDecorate %output Binding 0 +OpDecorate %values ArrayStride 4 +OpMemberDecorate %block 0 Offset 0 +OpDecorate %block Block + +%void = OpTypeVoid +%function = OpTypeFunction %void +%i32 = OpTypeInt 32 1 +%values = OpTypeRuntimeArray %i32 +%block = OpTypeStruct %values +%block_ptr = OpTypePointer StorageBuffer %block +%i32_ptr = OpTypePointer StorageBuffer %i32 +%zero = OpConstant %i32 0 +%answer = OpConstant %i32 42 +%output = OpVariable %block_ptr StorageBuffer + +%main = OpFunction %void None %function +%entry = OpLabel +%item = OpAccessChain %i32_ptr %output %zero %zero +OpStore %item %answer +OpReturn +OpFunctionEnd diff --git a/tools/build_dawn.sh b/tools/build_dawn.sh new file mode 100755 index 0000000..4e9a9ae --- /dev/null +++ b/tools/build_dawn.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) +source "$ROOT/tools/dependencies.sh" +readonly LOCAL="$ROOT/third_party/local/dawn" +readonly SOURCE="$LOCAL/source" +readonly BUILD="$LOCAL/build-latest" +readonly TOOLS_BUILD="$LOCAL/build-spirv-tools" +readonly STAGE="$ROOT/third_party/dawn" + +"$ROOT/tools/fetch_dawn.sh" + +backend_flags=(-DDAWN_ENABLE_NULL=OFF) +case "$(uname -s)" in + Darwin) backend_flags+=(-DDAWN_ENABLE_METAL=ON -DDAWN_ENABLE_VULKAN=OFF) ;; + Linux) backend_flags+=(-DDAWN_ENABLE_METAL=OFF -DDAWN_ENABLE_VULKAN=ON) ;; + *) echo "gpu.cpp supports Dawn on macOS and Linux" >&2; exit 1 ;; +esac + +cmake -S "$SOURCE" -B "$BUILD" -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_SHARED_LIBS=OFF \ + -DDAWN_BUILD_MONOLITHIC_LIBRARY=SHARED \ + -DDAWN_BUILD_SAMPLES=OFF \ + -DDAWN_BUILD_TESTS=OFF \ + -DDAWN_BUILD_BENCHMARKS=OFF \ + -DDAWN_BUILD_PROTOBUF=OFF \ + -DDAWN_FETCH_DEPENDENCIES=OFF \ + -DDAWN_USE_GLFW=OFF \ + -DTINT_BUILD_CMD_TOOLS=OFF \ + -DTINT_BUILD_SPV_READER=ON \ + -DTINT_BUILD_SPV_WRITER=OFF \ + "${backend_flags[@]}" +cmake --build "$BUILD" --target webgpu_dawn + +cmake -S "$SOURCE/third_party/spirv-tools/src" -B "$TOOLS_BUILD" -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DSPIRV-Headers_SOURCE_DIR="$SOURCE/third_party/spirv-headers/src" \ + -DSPIRV_SKIP_TESTS=ON \ + -DSPIRV_SKIP_EXECUTABLES=OFF +cmake --build "$TOOLS_BUILD" --target spirv-as + +cmake -E remove_directory "$STAGE" +cmake -E make_directory "$STAGE/include/dawn" "$STAGE/include/webgpu" \ + "$STAGE/lib" "$STAGE/bin" +cmake -E copy "$BUILD/gen/include/dawn/webgpu.h" \ + "$BUILD/gen/include/dawn/webgpu_cpp.h" "$STAGE/include/dawn" +cmake -E copy_directory "$SOURCE/include/webgpu" "$STAGE/include/webgpu" +cmake -E copy "$BUILD/gen/include/webgpu/webgpu_cpp_chained_struct.h" \ + "$STAGE/include/webgpu" +cmake -E copy "$TOOLS_BUILD/tools/spirv-as" "$STAGE/bin" +if [[ $(uname -s) == Darwin ]]; then + cmake -E copy "$BUILD/src/dawn/native/libwebgpu_dawn.dylib" "$STAGE/lib" +else + cmake -E copy "$BUILD/src/dawn/native/libwebgpu_dawn.so" "$STAGE/lib" +fi + +echo "Staged Dawn $DAWN_REV in $STAGE" diff --git a/tools/build_emdawn.sh b/tools/build_emdawn.sh new file mode 100755 index 0000000..b7b619d --- /dev/null +++ b/tools/build_emdawn.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) +source "$ROOT/tools/dependencies.sh" +readonly EMSDK_ROOT_DIR=${GPUCPP_EMSDK_ROOT:-"$ROOT/../emsdk"} +readonly SOURCE="$ROOT/third_party/local/dawn/source" +readonly BUILD="$ROOT/third_party/local/dawn/build-web" +readonly STAGE="$ROOT/third_party/emdawnwebgpu" + +if [[ ! -d "$EMSDK_ROOT_DIR/.git" ]]; then + echo "emsdk is not cloned at $EMSDK_ROOT_DIR" >&2 + exit 1 +fi + +if [[ $(git -C "$EMSDK_ROOT_DIR" rev-parse HEAD) != "$EMSDK_REV" ]]; then + git -C "$EMSDK_ROOT_DIR" fetch --depth 1 origin "$EMSDK_REV" + git -C "$EMSDK_ROOT_DIR" checkout --detach -q FETCH_HEAD +fi +export EMSDK_QUIET=1 +"$EMSDK_ROOT_DIR/emsdk" install "$EMSCRIPTEN_VERSION" +"$EMSDK_ROOT_DIR/emsdk" activate "$EMSCRIPTEN_VERSION" +source "$EMSDK_ROOT_DIR/emsdk_env.sh" >/dev/null + +"$ROOT/tools/fetch_dawn.sh" +emcmake cmake -S "$SOURCE" -B "$BUILD" -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DDAWN_BUILD_SAMPLES=OFF \ + -DDAWN_BUILD_TESTS=OFF \ + -DDAWN_BUILD_BENCHMARKS=OFF \ + -DDAWN_BUILD_PROTOBUF=OFF \ + -DDAWN_FETCH_DEPENDENCIES=OFF \ + -DDAWN_USE_GLFW=OFF \ + -DDAWN_SUPPORTS_CXX_MODULES=OFF \ + -DTINT_BUILD_CMD_TOOLS=OFF \ + -DTINT_BUILD_IR_BINARY=OFF +cmake --build "$BUILD" --target emdawnwebgpu_pkg + +cmake -E remove_directory "$STAGE" +cmake -E copy_directory "$BUILD/emdawnwebgpu_pkg" "$STAGE" +echo "Staged Emdawnwebgpu from Dawn $DAWN_REV in $STAGE" diff --git a/tools/dependencies.sh b/tools/dependencies.sh new file mode 100644 index 0000000..15875ef --- /dev/null +++ b/tools/dependencies.sh @@ -0,0 +1,5 @@ +readonly DAWN_REV=0790e933451fcf14ce2ec8f1c88fcefffdf5d218 +readonly SPIRV_TOOLS_REV=1c336172641682bab6e066767d09fdff1d826467 +readonly SPIRV_HEADERS_REV=29981f65241605e08b0ede4cfeb999fe3b723c6a +readonly EMSDK_REV=948c31acd3f369a5da276e33ab2ed57108c165e5 +readonly EMSCRIPTEN_VERSION=5.0.6 diff --git a/tools/fetch_dawn.sh b/tools/fetch_dawn.sh new file mode 100755 index 0000000..b246729 --- /dev/null +++ b/tools/fetch_dawn.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) +source "$ROOT/tools/dependencies.sh" +readonly SOURCE="$ROOT/third_party/local/dawn/source" + +checkout() { + local directory=$1 url=$2 revision=$3 + if [[ ! -d "$directory/.git" ]]; then + mkdir -p "$directory" + git -C "$directory" init -q + fi + if [[ $(git -C "$directory" rev-parse HEAD 2>/dev/null || true) != "$revision" ]]; then + git -C "$directory" fetch --depth 1 "$url" "$revision" + git -C "$directory" checkout --detach -q FETCH_HEAD + fi +} + +checkout "$SOURCE" https://dawn.googlesource.com/dawn "$DAWN_REV" +python3 "$SOURCE/tools/fetch_dawn_dependencies.py" --directory "$SOURCE" --shallow +checkout "$SOURCE/third_party/spirv-tools/src" \ + https://github.com/KhronosGroup/SPIRV-Tools.git "$SPIRV_TOOLS_REV" +checkout "$SOURCE/third_party/spirv-headers/src" \ + https://github.com/KhronosGroup/SPIRV-Headers.git "$SPIRV_HEADERS_REV" diff --git a/utils/array_utils.hpp b/utils/array_utils.hpp index 6e28fec..0207a41 100644 --- a/utils/array_utils.hpp +++ b/utils/array_utils.hpp @@ -86,7 +86,7 @@ std::string show(const numtype *a, size_t rows, size_t cols, } else snprintf(buffer, 16, "%10.2e", a[i * cols + j]); } else if constexpr (std::is_same::value) { - float tmp = halfToFloat(a[i * cols + j]); + float tmp = static_cast(a[i * cols + j]); if (std::abs(tmp) < 1000 && std::abs(tmp) > 0.01 || tmp == 0.0) { @@ -222,7 +222,7 @@ inline void randn(half *a, size_t N, std::mt19937 &gen, float mean = 0.0, float std = 1.0) { std::normal_distribution dist(mean, std); for (int i = 0; i < N; i++) { - a[i] = halfFromFloat(dist(gen)); + a[i] = static_cast(dist(gen)); } } @@ -324,8 +324,8 @@ inline bool isclose(float *a, float *b, size_t n, float tol = 1e-3) { inline bool isclose(half *a, half *b, size_t n, float tol = 1) { for (size_t i = 0; i < n; i++) { - float ai = halfToFloat(a[i]); - float bi = halfToFloat(b[i]); + float ai = static_cast(a[i]); + float bi = static_cast(b[i]); if (std::abs(ai - bi) > tol || std::isnan(ai) || std::isnan(bi)) { LOG(kDefLog, kError, "Mismatch at index %d: %f != %f", i, ai, bi); return false; diff --git a/experimental/tui.h b/utils/tui.hpp similarity index 80% rename from experimental/tui.h rename to utils/tui.hpp index 66fad54..ed8f171 100644 --- a/experimental/tui.h +++ b/utils/tui.hpp @@ -1,15 +1,17 @@ -#ifndef TUI_H -#define TUI_H +#ifndef GPU_CPP_UTILS_TUI_HPP +#define GPU_CPP_UTILS_TUI_HPP +#include #include #include #include +#include // Work-in-progress - various terminal UI visualization functions namespace gpu { -void cls() { printf("\033[2J\033[H"); } +inline void cls() { printf("\033[2J\033[H"); } template void canvas(const std::array &raster) { @@ -33,12 +35,13 @@ void canvas(const std::array &raster) { } // double pendulum rasterizer -void rasterize(float *pos, size_t n, float maxX, float maxY, std::string &screen, - size_t screenWidth, size_t screenHeight) { +inline void rasterize(const float *pos, size_t n, float maxX, float maxY, + std::string &screen, size_t screenWidth, + size_t screenHeight) { static const char intensity[] = " .`'^-+=*x17X$8#%@"; const size_t eps = 1; // maximum number of simulations to display on the screen - const size_t nShow = std::min(static_cast(n), 2000); + const size_t nShow = std::min(n, size_t{2000}); for (size_t i = 0; i < screenHeight; ++i) { for (size_t j = 0; j < screenWidth - 2; ++j) { int count = 0; diff --git a/utils/wgsl.hpp b/utils/wgsl.hpp new file mode 100644 index 0000000..8057745 --- /dev/null +++ b/utils/wgsl.hpp @@ -0,0 +1,39 @@ +#ifndef GPU_CPP_UTILS_WGSL_HPP +#define GPU_CPP_UTILS_WGSL_HPP + +#include +#include + +namespace gpu { + +// Unrolls flat WGSL loops of the form +// `for (var i: u32 = START; i < END; i++) { BODY }`. +inline std::string unrollLoops(const std::string &code, int threshold = 32) { + static const std::regex loop( + R"(for\s*\(\s*var\s+(\w+):\s*u32\s*=\s*(\d+)\s*;\s*\1\s*<\s*(\d+)\s*;\s*\1\+\+\s*\)\s*\{\s*([^{}]*)\})"); + std::string result; + size_t previous = 0; + for (std::sregex_iterator it(code.begin(), code.end(), loop), sentinel; + it != sentinel; ++it) { + const std::smatch &match = *it; + result.append(code, previous, match.position() - previous); + const auto name = match[1].str(); + const int start = std::stoi(match[2]); + const int end = std::stoi(match[3]); + if (end - start > threshold) { + result += match.str(); + } else { + const std::regex variable("\\b" + name + "\\b"); + for (int i = start; i < end; ++i) + result += + std::regex_replace(match[4].str(), variable, std::to_string(i)); + } + previous = match.position() + match.length(); + } + result.append(code, previous, std::string::npos); + return result; +} + +} // namespace gpu + +#endif