diff --git a/.gitignore b/.gitignore index 98234268c1..740d5464fb 100644 --- a/.gitignore +++ b/.gitignore @@ -92,3 +92,7 @@ test_data/* # The experimental/builder directory should be tracked despite matching build* !experimental/builder !experimental/builder/** +experimental/grouped_convolution_tile_instances/instances/* +!experimental/grouped_convolution_tile_instances/instances/*.in +!experimental/grouped_convolution_tile_instances/instances/*.inc +experimental/grouped_convolution_tile_instances/*.inc diff --git a/CHANGELOG.md b/CHANGELOG.md index 066dc9aa3b..c3a257e464 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj * Added support for gfx1153 target. * Added FMHA batch prefill kernel support for several KV cache layouts, flexible page sizes, and different lookup table configurations. * Added gpt-oss sink support for FMHA FWD, include qr_ks_vs, qr_async, qr_async_trload and splitkv pipelines. +* Added persistent async input scheduler for CK Tile universal GEMM kernels to support asynchronous input streaming. ### Changed diff --git a/CMakeLists.txt b/CMakeLists.txt index 121c663f64..54464d6809 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,7 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -cmake_minimum_required(VERSION 3.14) +cmake_minimum_required(VERSION 3.21) if(POLICY CMP0140) # policies CMP0140 not known to CMake until 3.25 cmake_policy(SET CMP0140 NEW) @@ -704,6 +704,11 @@ option(BUILD_MHA_LIB "Build the static library for flash attention" OFF) add_subdirectory(library) +if (CK_EXPERIMENTAL_BUILDER) + add_subdirectory(experimental/builder) + add_subdirectory(experimental/grouped_convolution_tile_instances) +endif() + if(NOT GPU_ARCHS AND USER_GPU_TARGETS AND NOT MIOPEN_REQ_LIBS_ONLY) rocm_package_setup_component(tests LIBRARY_NAME composablekernel @@ -735,10 +740,6 @@ if (NOT MIOPEN_REQ_LIBS_ONLY) add_subdirectory(profiler) endif() -if (CK_EXPERIMENTAL_BUILDER) - add_subdirectory(experimental/builder) -endif() - if(CK_USE_CODEGEN AND (SUPPORTED_GPU_TARGETS MATCHES "gfx9" OR GPU_ARCHS)) add_subdirectory(codegen) endif() diff --git a/CMakePresets.json b/CMakePresets.json new file mode 100644 index 0000000000..f81dbadb19 --- /dev/null +++ b/CMakePresets.json @@ -0,0 +1,91 @@ +{ + "version": 3, + "cmakeMinimumRequired": { + "major": 3, + "minor": 21, + "patch": 0 + }, + "configurePresets": [ + { + "name": "use-gfx908", + "hidden": true, + "cacheVariables": { + "GPU_TARGETS": "gfx908" + } + }, + { + "name": "use-gfx90a", + "hidden": true, + "cacheVariables": { + "GPU_TARGETS": "gfx90a" + } + }, + { + "name": "use-gfx942", + "hidden": true, + "cacheVariables": { + "GPU_TARGETS": "gfx942" + } + }, + { + "name": "use-gfx950", + "hidden": true, + "cacheVariables": { + "GPU_TARGETS": "gfx950" + } + }, + { + "name": "dev", + "binaryDir": "${sourceDir}/build", + "displayName": "CK Dev", + "environment": {}, + "cacheVariables": { + "CMAKE_PREFIX_PATH": "/opt/rocm/", + "CMAKE_CXX_COMPILER": "/opt/rocm/llvm/bin/clang++", + "CMAKE_HIP_COMPILER": "/opt/rocm/llvm/bin/clang++", + "CMAKE_CXX_FLAGS": "-ftemplate-backtrace-limit=0 -fPIE -Wno-gnu-line-marker -fbracket-depth=512", + "CMAKE_BUILD_TYPE": "Release", + "BUILD_DEV": "ON", + "CMAKE_VERBOSE_MAKEFILE": "ON", + "USE_BITINT_EXTENSION_INT4": "OFF", + "GPU_TARGETS": "gfx908;gfx90a;gfx942" + } + }, + { + "name": "dev-gfx908", + "displayName": "CK Dev - gfx908", + "description": "Development build for AMD GPU gfx908", + "inherits": [ + "use-gfx908", + "dev" + ] + }, + { + "name": "dev-gfx90a", + "displayName": "CK Dev - gfx90a", + "description": "Development build for AMD GPU gfx90a", + "inherits": [ + "use-gfx90a", + "dev" + ] + }, + { + "name": "dev-gfx942", + "displayName": "CK Dev - gfx942", + "description": "Development build for AMD GPU gfx942", + "inherits": [ + "use-gfx942", + "dev" + ] + }, + { + "name": "dev-gfx950", + "displayName": "CK Dev - gfx950", + "description": "Development build for AMD GPU gfx950", + "inherits": [ + "use-gfx950", + "dev" + ] + } + ] +} \ No newline at end of file diff --git a/Jenkinsfile b/Jenkinsfile index e01cfcbf01..f3a597e404 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -115,7 +115,7 @@ def generateAndArchiveBuildTraceVisualization(String buildTraceFileName) { // Run container to get snapshot def dockerOpts = "--cap-add=SYS_ADMIN -v \"\$(pwd)/workspace:/workspace\" -e NODE_PATH=/home/pptruser/node_modules -e BUILD_TRACE_FILE=${buildTraceFileName}" // Create unique image name by sanitizing job name - def sanitizedJobName = env.JOB_NAME.replaceAll(/[\/\\:*?"<>| ]/, '_') + def sanitizedJobName = env.JOB_NAME.replaceAll(/[\/\\:*?"<>| ]/, '_').replaceAll('%2F', '_') def architectureName = (buildTraceFileName =~ /(gfx[0-9a-zA-Z]+)/)[0][1] def imageName = "perfetto_snapshot_${sanitizedJobName}_build_${env.BUILD_NUMBER}_${architectureName}.png" sh """ @@ -580,7 +580,7 @@ def cmake_build(Map conf=[:]){ if (params.NINJA_BUILD_TRACE) { echo "running ninja build trace" } - if (params.RUN_BUILDER_TESTS && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { + if ((params.RUN_BUILDER_TESTS || params.RUN_FULL_CONV_TILE_TESTS) && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { setup_args = " -D CK_EXPERIMENTAL_BUILDER=ON " + setup_args } setup_cmd = conf.get( @@ -1091,7 +1091,7 @@ CRON_SETTINGS = BRANCH_NAME == "develop" ? '''0 23 * * * % RUN_FULL_QA=true;RUN_ 0 19 * * * % BUILD_DOCKER=true;COMPILER_VERSION=amd-staging;BUILD_COMPILER=/llvm-project/build/bin/clang++;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true 0 17 * * * % BUILD_DOCKER=true;COMPILER_VERSION=amd-mainline;BUILD_COMPILER=/llvm-project/build/bin/clang++;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true 0 15 * * * % BUILD_INSTANCES_ONLY=true;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;FORCE_CI=true - 0 13 * * * % RUN_AITER_TESTS=true;BUILD_LEGACY_OS=true;USE_SCCACHE=false;RUN_PERFORMANCE_TESTS=false;FORCE_CI=true + 0 13 * * * % RUN_FULL_CONV_TILE_TESTS=true;RUN_AITER_TESTS=true;BUILD_LEGACY_OS=true;USE_SCCACHE=false;RUN_PERFORMANCE_TESTS=false;FORCE_CI=true 0 11 * * * % RUN_PYTORCH_TESTS=true;RUN_CODEGEN_TESTS=false;USE_SCCACHE=false;RUN_PERFORMANCE_TESTS=false;BUILD_GFX101=false;BUILD_GFX103=false;BUILD_GFX11=false;BUILD_GFX12=false;BUILD_GFX90A=false;FORCE_CI=true''' : "" pipeline { @@ -1255,6 +1255,10 @@ pipeline { name: "RUN_AITER_TESTS", defaultValue: false, description: "Run AITER tests with latest CK develop branch (default: OFF)") + booleanParam( + name: "RUN_FULL_CONV_TILE_TESTS", + defaultValue: false, + description: "Run CK Tile grouped convolution tests with latest CK develop branch (default: OFF)") string( name: 'aiter_branch', defaultValue: 'main', @@ -1318,21 +1322,15 @@ pipeline { agent{ label rocmnode("nogpu") } environment{ setup_args = "NO_CK_BUILD" - execute_cmd = "(cd .. && git ls-files \'*.h\' \ - \'*.hpp\' \ - \'*.cpp\' \ - \'*.h.in\' \ - \'*.hpp.in\' \ - \'*.cpp.in\' \ - \'*.cl\' \ - | grep -v 'build/' \ - | grep -v 'include/rapidjson' \ - | xargs -n 1 -P 1 -I{} -t sh -c \'clang-format-18 -style=file {} | diff - {}\') && \ + execute_cmd = """cd .. && \ + find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ + -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ + xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)' && \ /cppcheck/build/bin/cppcheck ../* -v -j \$(nproc) -I ../include -I ../profiler/include -I ../library/include \ -D CK_ENABLE_FP64 -D CK_ENABLE_FP32 -D CK_ENABLE_FP16 -D CK_ENABLE_FP8 -D CK_ENABLE_BF16 -D CK_ENABLE_BF8 -D CK_ENABLE_INT8 \ -D __gfx908__ -D __gfx90a__ -D __gfx942__ -D __gfx1030__ -D __gfx1100__ -D __gfx1101__ -D __gfx1102__ \ -U __gfx803__ -U __gfx900__ -U __gfx906__ -U CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4 \ - --file-filter=*.cpp --force --enable=all --output-file=ck_cppcheck.log" + --file-filter=*.cpp --force --enable=all --output-file=ck_cppcheck.log""" } steps{ buildHipClangJobAndReboot(setup_args:setup_args, setup_cmd: "", build_cmd: "", execute_cmd: execute_cmd) @@ -1348,17 +1346,10 @@ pipeline { agent{ label rocmnode("nogpu") } environment{ setup_args = "NO_CK_BUILD" - execute_cmd = "(cd .. && git ls-files \ - \'*.h\' \ - \'*.hpp\' \ - \'*.cpp\' \ - \'*.h.in\' \ - \'*.hpp.in\' \ - \'*.cpp.in\' \ - \'*.cl\' \ - | grep -v 'build/' \ - | grep -v 'include/rapidjson' \ - | xargs -n 1 -P 1 -I{} -t sh -c \'clang-format-18 -style=file {} | diff - {}\')" + execute_cmd = """cd .. && \ + find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ + -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ + xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)'""" } steps{ buildHipClangJobAndReboot(setup_args:setup_args, setup_cmd: "", build_cmd: "", execute_cmd: execute_cmd) @@ -1423,6 +1414,36 @@ pipeline { } } } + stage("Run Full Grouped Conv Tile Tests") + { + when { + beforeAgent true + expression { env.SHOULD_RUN_CI.toBoolean() } + } + parallel + { + stage("Run Full Grouped Conv Tile Tests on gfx90a") + { + when { + beforeAgent true + expression { params.RUN_FULL_CONV_TILE_TESTS.toBoolean() } + } + agent{ label rocmnode("gfx90a")} + environment{ + setup_args = "NO_CK_BUILD" + execute_args = """ python3 ../experimental/builder/src/generate_instances.py --mode=profiler && \ + ../script/cmake-ck-dev.sh ../ gfx90a && \ + make -j64 test_grouped_convnd_fwd_tile && \ + ./bin/test_grouped_convnd_fwd_tile""" + } + steps{ + // TODO: Reenable after the instance fixes + // buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + cleanWs() + } + } + } + } stage("Run Grouped Conv Large Case Tests") { when { @@ -1440,7 +1461,7 @@ pipeline { agent{ label rocmnode("gfx90a")} environment{ setup_args = "NO_CK_BUILD" - execute_args = """ ../script/cmake-ck-dev.sh ../ gfx90a && \ + execute_args = """ cmake .. --preset dev-gfx90a && \ make -j64 test_grouped_convnd_fwd_large_cases test_grouped_convnd_bwd_data_large_cases test_grouped_convnd_fwd_bias_clamp_large_cases && \ ./bin/test_grouped_convnd_fwd_large_cases && ./bin/test_grouped_convnd_bwd_data_large_cases && ./bin/test_grouped_convnd_fwd_bias_clamp_large_cases""" } @@ -1469,8 +1490,8 @@ pipeline { environment{ setup_args = "NO_CK_BUILD" execute_args = """ cd ../build && \ - ../script/cmake-ck-dev.sh ../ gfx90a && \ - make -j64 test_grouped_convnd_fwd_dataset_xdl \ + cmake .. --preset dev-gfx90a && \ + make -j64 test_grouped_convnd_fwd_dataset_xdl && \ test_grouped_convnd_bwd_data_dataset_xdl \ test_grouped_convnd_bwd_weight_dataset_xdl && \ cd ../test_data && \ @@ -1731,7 +1752,10 @@ pipeline { } agent{ label rocmnode("gfx90a") } environment{ - setup_args = """ -DGPU_TARGETS="gfx942" -DCK_USE_ALTERNATIVE_PYTHON=/opt/Python-3.8.13/bin/python3.8 """ + // SLES15 is a legacy platform with limited C++20 ecosystem support (older system libraries, + // standard library implementation). While the ROCm compiler supports C++20, the experimental + // CK Builder requires full C++20 feature support that does not be reliably available on SLES15. + setup_args = """ -DGPU_TARGETS="gfx942" -DCK_USE_ALTERNATIVE_PYTHON=/opt/Python-3.8.13/bin/python3.8 -DCK_EXPERIMENTAL_BUILDER=OFF """ execute_args = " " } steps{ diff --git a/docs/conceptual/ck_tile/buffer_views.rst b/docs/conceptual/ck_tile/buffer_views.rst index 14b8309504..03b8e87b1b 100644 --- a/docs/conceptual/ck_tile/buffer_views.rst +++ b/docs/conceptual/ck_tile/buffer_views.rst @@ -1,35 +1,13 @@ -.. meta:: - :description: Composable Kernel CK Tile buffer views - :keywords: composable kernel, CK, CK Tile, ROCm, API, buffer view, raw memory - .. _ck_tile_buffer_views: -CK Tile buffer view -======================= - -Buffer view is an abstraction that provides structured access to memory. The ``buffer_view`` class is exposed in ``include/ck_tile/core/tensor/buffer_view.hpp``. - -Buffer view serves as the foundation for :ref:`ck_tile_tensor_views`. BufferView handles memory addressing and type safety, while TensorView builds upon this to add multi-dimensional coordinates (shape and strides). - - -Buffer view provides the following advantages: - -* A unified interface across global, shared, and register memory -* Address spaces encoded in types, taking advantage of compile-time type checking -* Configurable handling of invalid values, out-of-bounds operations, and conditional access patterns -* Atomic operations for parallel algorithms -* AMD GPU-specific optimizations -* Automatic application of appropriate memory ordering constraints and cache control directives based on the target address space and operation type - - -[TO DO: do we want to say more about these items? There wasn't a lot of detail in the original text, so I put them in a list for now] - - +Buffer Views - Raw Memory Access Address Space Usage Patterns ---------------------------- -[TO DO: explain in words what the diagram shows] +.. + Original mermaid diagram (edit here, then run update_diagrams.py) + .. Original mermaid diagram (edit here, then run update_diagrams.py) @@ -66,18 +44,26 @@ Address Space Usage Patterns style Compute fill:#e0e7ff,stroke:#4338ca,stroke-width:2px + + + + .. image:: diagrams/buffer_views_1.svg :alt: Diagram :align: center +C++ Implementation +------------------ +**File**: ``include/ck_tile/core/tensor/buffer_view.hpp`` Basic Creation ~~~~~~~~~~~~~~ -[TO DO: remove "modern C++ template metaprogramming" and "zero-overhead abstraction"] +By encoding critical properties such as buffer size and address space as template parameters, BufferView transforms what would traditionally be runtime decisions into compile-time constants. This design philosophy enables the compiler to perform aggressive optimizations, including constant propagation, loop unrolling, and instruction selection, that would be impossible with runtime parameters. -[TO DO: might want to move the implementation details to a separate section under "reference"] +The use of compile-time constants extends beyond mere optimization. When the buffer size is encoded in the type system using constructs like ``number<8>{}``, the compiler can statically verify that array accesses are within bounds, eliminate unnecessary bounds checks, and even restructure algorithms to better match the known data dimensions. This compile-time knowledge propagates through the entire computation, enabling optimizations at every level of the abstraction hierarchy. +The address space template parameter represents another crucial design decision. By making the memory space part of the type system, BufferView ensures that operations appropriate for one memory space cannot be accidentally applied to another. This type safety prevents common errors such as attempting atomic operations on register memory or using global memory synchronization primitives on local memory. The compiler enforces these constraints at compile time, transforming potential runtime errors into compile-time diagnostics. .. code-block:: cpp @@ -98,7 +84,6 @@ Basic Creation buffer_size // number of elements ); - // Implementation detail: The actual C++ template is: // template (data, buffer_size, custom_invalid); - - // Invalid element access with is_valid_element=false - // Returns custom_invalid due to custom invalid value mode - auto invalid_value = buffer_view.template get(0, 0, false); - printf("Invalid element: %.1f\n", invalid_value.get(0)); - - // Out of bounds access - AMD buffer addressing handles bounds checking - // Will return custom_invalid when accessing beyond buffer_size - auto oob_value = buffer_view.template get(0, 100, true); - printf("Out of bounds: %.1f\n", oob_value.get(0)); - - - - - Get Operations -------------- -[TO DO: might want to put this implementation detail in the reference section] +Scalar Access +~~~~~~~~~~~~~ -The signature for the ``buffer_view`` ``get()`` takes four parameters: +The get operations in BufferView form the cornerstone of memory access patterns in CK Tile. These operations embody a advanced understanding of GPU memory systems and the patterns that lead to optimal performance. The scalar access interface incorporates multiple layers of optimization and safety mechanisms that work together to provide both performance and correctness. -``i``: the primary offset into the buffer expressed in terms of elements of type T rather than raw bytes. +The parameter structure of scalar access operations reflects careful design choices aimed at maximizing flexibility while maintaining efficiency. The base index parameter ``i`` represents the primary offset into the buffer, expressed in terms of elements of type T rather than raw bytes. This type-aware indexing prevents common errors related to pointer arithmetic and ensures that vector types are handled correctly. The additional ``linear_offset`` parameter provides fine-grained control over the final access location, enabling complex access patterns without requiring expensive index calculations in the kernel code. -``linear_offset``: [TO DO: what is this?] +The ``is_valid_element`` parameter provides a solution to conditional memory access. Rather than using traditional if-statements that would cause warp divergence, this boolean parameter enables predicated execution where the memory access occurs unconditionally but the result is conditionally used. This approach maintains uniform control flow across all threads in a warp, preserving the SIMD execution model that is fundamental to GPU performance. -``is_valid_element``: [TO DO: what is this?] +The invalid value modes provide a mechanism for handling the boundary conditions that arise in parallel algorithms. When ``InvalidElementUseNumericalZeroValue`` is set to true, the system returns zero for any invalid access, whether due to the ``is_valid_element`` flag or out-of-bounds indexing. This mode is important for algorithms where zero serves as a natural extension value, such as in image processing with zero-padding or sparse matrix operations where missing elements are implicitly zero. -[TO DO: the last param, that's the out of bounds handling, yes? -.. code:: cpp +The custom invalid value mode, activated when ``InvalidElementUseNumericalZeroValue`` is false, offers additional flexibility for algorithms with specific boundary requirements. This mode returns a user-specified value for invalid accesses, accommodating use cases such as sentinel values in sorting algorithms, infinity values in optimization problems, or special markers in data processing pipelines. The implementation ensures that this flexibility comes without performance penalty, using the same branchless execution strategies as the zero mode. - get(index_t i, - index_t linear_offset, - bool is_valid_element, - bool_constant = {}) +Out-of-bounds handling leverages AMD GPU hardware capabilities to provide safety with minimal impact to performance. When AMD buffer addressing is enabled, the hardware automatically clamps memory accesses to valid ranges, preventing the segmentation faults that would occur on CPU systems. This hardware-assisted bounds checking operates at wire speed, adding no overhead to the memory access path while ensuring that kernels cannot corrupt memory outside their allocated regions. +Vector Access +~~~~~~~~~~~~~ -[TO DO: need some context around the code] +Vector memory operations represent one of the most critical optimizations available in modern GPU programming, and BufferView's vector access interface exposes this capability. By using template parameters to specify vector types through constructs like ``ext_vector_t``, the interface enables compile-time selection of optimal load and store instructions that can transfer multiple data elements in a single memory transaction. This vectorization is crucial for :ref:`ck_tile_load_store_traits`, which automatically selects optimal access patterns. -[TO DO: code chunks need to have detail and explanation so that the reader can see what they're trying to demonstrate.] +The significance of vector operations extends beyond bandwidth improvements. GPUs are designed with wide memory buses that can transfer 128, 256, or even 512 bits per transaction. When scalar operations access only 32 bits at a time, they utilize only a fraction of this available bandwidth. Vector operations align with these wide buses, enabling full bandwidth utilization and reducing the total number of memory transactions required. +The implementation of vector access maintains the same parameter structure as scalar operations, providing consistency across the API while automatically handling the complexities of multi-element transfers. The system manages alignment requirements, ensures that vector loads and stores use the optimal hardware instructions, and handles cases where vector operations extend beyond buffer boundaries. This transparent handling of edge cases allows developers to use vector operations confidently without manual boundary checks or special-case code for partial vectors. -.. code-block:: cpp - - // Create buffer view - float data[8] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; - auto buffer_view = make_buffer_view(data, 8); - - // Simple get - compile-time bounds checking when possible - auto value_buf = buffer_view.template get(0,1,true); //get the buffer from the buffer view - float value = value_buf.get(0); //get the value from the buffer - - // Get with valid flag - branchless conditional access - bool valid_flag = false; - value_buf = buffer_view.template get(0,1,valid_flag); - value = value_buf.get(0); - // Returns 0 valid_flag is false - - // vectorized get - using float2 = ext_vector_t; - auto vector_buf = buffer_view.template get(0, 0, true); - // Loads 2 floats in a single instruction - float val1 = vector_buf.get(0); - float val2 = vector_buf.get(1); - } - -``ext_vector_t`` enables compile-time selection of optimal load and store instructions that can transfer multiple data elements in a single memory transaction. - -[TO DO: what is it actually doing? When does one use scalars vs vectors? Is it application specific or are there ] +Scalar vs Vectorized Memory Access +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. + Original mermaid diagram (edit here, then run update_diagrams.py) + .. Original mermaid diagram (edit here, then run update_diagrams.py) @@ -287,8 +216,9 @@ The signature for the ``buffer_view`` ``get()`` takes four parameters: Understanding BufferView Indexing ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -[TO DO: an explanation of the diagram is needed] - +.. + Original mermaid diagram (edit here, then run update_diagrams.py) + .. Original mermaid diagram (edit here, then run update_diagrams.py) @@ -335,14 +265,69 @@ Understanding BufferView Indexing .. image:: diagrams/buffer_views_3.svg :alt: Diagram :align: center - - + +C++ Get Operations +~~~~~~~~~~~~~~~~~~ + +.. code-block:: cpp + + __device__ void example_get_operations() + { + // Create buffer view + float data[8] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + auto buffer_view = make_buffer_view(data, 8); + + // Simple get - compile-time bounds checking when possible + auto value_buf = buffer_view.template get(0,1,true); //get the buffer from the buffer view + float value = value_buf.get(0); //get the value from the buffer + + // Get with valid flag - branchless conditional access + bool valid_flag = false; + value_buf = buffer_view.template get(0,1,valid_flag); + value = value_buf.get(0); + // Returns 0 valid_flag is false + + // vectorized get + using float2 = ext_vector_t; + auto vector_buf = buffer_view.template get(0, 0, true); + // Loads 2 floats in a single instruction + float val1 = vector_buf.get(0); + float val2 = vector_buf.get(1); + } + +Custom Value Return Mode for OOB & Invalid Access +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: cpp + + void scalar_get_operations_example() { + + // Create data array + constexpr size_t buffer_size = 8; + float data[buffer_size] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + float custom_invalid = 13.0f; + + // Create global memory buffer view with zero invalid value mode (default) + auto buffer_view = make_buffer_view(data, buffer_size, custom_invalid); + + // Invalid element access with is_valid_element=false + // Returns custom_invalid due to custom invalid value mode + auto invalid_value = buffer_view.template get(0, 0, false); + printf("Invalid element: %.1f\n", invalid_value.get(0)); + + // Out of bounds access - AMD buffer addressing handles bounds checking + // Will return custom_invalid when accessing beyond buffer_size + auto oob_value = buffer_view.template get(0, 100, true); + printf("Out of bounds: %.1f\n", oob_value.get(0)); + } + +.. note:: + + Partial Out Of Bound (OOB) access during vector reads will return 'junk' values for the OOB access. Zero or custom invalid value is only returned for complete invalid/OOB access, in other words, it is only returned when the first address of the vector is invalid. Update Operations ----------------- -Update operations modify the buffer content. The ``set()`` method writes a value to a specific location. - .. code-block:: cpp void scalar_set_operations_example() { @@ -373,8 +358,6 @@ Update operations modify the buffer content. The ``set()`` method writes a value Atomic Operations ----------------- -[TO DO: this needs information] - Atomic vs Non-Atomic Operations ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -441,3 +424,21 @@ C++ Atomic Operations __syncthreads(); } + +Summary +------- + +BufferView abstracts GPU memory hierarchies behind a concise interface. The approach is intended to keep overhead small while enabling optimizations that are otherwise awkward in low-level code. + +BufferView offers a unified interface across global, shared, and register memory. Using the same API for each space can lower cognitive overhead, reduce certain classes of mistakes, and support code reuse via template parameters. + +Address spaces are encoded in types so that common errors are reported at compile time. Consistent with CK Tile’s zero-overhead design aim, compile-time checks are favored over runtime guards. The C++ type system enforces memory-space constraints and can make valid cases more amenable to compiler optimization. + +BufferView supports configurable handling of invalid values, optional runtime bounds checks, and conditional access patterns. It also provides atomic operations for thread-safe updates. These features are intended to cover common edge cases without adding unnecessary overhead. + +By hiding the complexity of different memory spaces while exposing the operations needed for high-performance GPU computing, BufferView establishes a pattern that the rest of CK Tile follows: compile-time abstractions that enhance rather than compromise performance. The :ref:`ck_tile_tensor_views` and :ref:`ck_tile_distribution` add capability while maintaining the efficiency established at the base. For hardware-specific details about memory hierarchies, see :ref:`ck_tile_gpu_basics`. + +Next Steps +---------- + +Continue to :ref:`ck_tile_tensor_views` to learn how to build structured tensor views on top of buffer views. diff --git a/docs/sphinx/requirements.in b/docs/sphinx/requirements.in index b1ab09e6f7..f2fb27e2b9 100644 --- a/docs/sphinx/requirements.in +++ b/docs/sphinx/requirements.in @@ -1,2 +1,2 @@ -rocm-docs-core[api_reference]==1.31.1 +rocm-docs-core[api_reference]==1.31.3 sphinxcontrib-bibtex==2.6.5 diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index 099e9e439f..23397503df 100644 --- a/docs/sphinx/requirements.txt +++ b/docs/sphinx/requirements.txt @@ -237,7 +237,7 @@ requests==2.32.3 # via # pygithub # sphinx -rocm-docs-core[api-reference]==1.31.1 +rocm-docs-core[api-reference]==1.31.3 # via -r requirements.in rpds-py==0.24.0 # via diff --git a/example/01_gemm/CMakeLists.txt b/example/01_gemm/CMakeLists.txt index 2d65368d4f..aba462638e 100644 --- a/example/01_gemm/CMakeLists.txt +++ b/example/01_gemm/CMakeLists.txt @@ -149,3 +149,7 @@ add_example_executable(example_gemm_wmma_fp16_fp8_v3 gemm_wmma_fp16_fp8_v3.cpp) add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp16_fp8_v3) add_example_executable(example_gemm_wmma_fp16_pk_i4_v3_b_scale gemm_wmma_fp16_pk_i4_v3_b_scale.cpp) add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp16_pk_i4_v3_b_scale) +add_example_executable(example_gemm_wmma_fp8_bpreshuffle gemm_wmma_fp8_bpreshuffle.cpp) +add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp8_bpreshuffle) +add_example_executable(example_gemm_wmma_fp16_bpreshuffle gemm_wmma_fp16_bpreshuffle.cpp) +add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp16_bpreshuffle) diff --git a/example/01_gemm/gemm_wmma_fp16_bpreshuffle.cpp b/example/01_gemm/gemm_wmma_fp16_bpreshuffle.cpp new file mode 100644 index 0000000000..d03971e6ec --- /dev/null +++ b/example/01_gemm/gemm_wmma_fp16_bpreshuffle.cpp @@ -0,0 +1,70 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "common.hpp" + +#include "ck/ck.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp" +#include "ck/library/utility/check_err.hpp" +#include "ck/library/utility/device_memory.hpp" +#include "ck/library/utility/host_tensor.hpp" +#include "ck/library/utility/host_tensor_generator.hpp" +#include "ck/library/utility/literals.hpp" +#include "ck/stream_config.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3_b_preshuffle.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/utility/data_type.hpp" +#include "ck/utility/get_id.hpp" +#include "ck/utility/scheduler_enum.hpp" + +#include +#include +#include + +using F16 = ck::half_t; +using F32 = float; + +using ADataType = F16; +using BDataType = F16; +using AccDataType = F32; +using CShuffleDataType = F32; +using CDataType = F16; +using ComputeTypeA = F16; +using ComputeTypeB = F16; + +using ALayout = Row; +using BLayout = Col; +using CLayout = Row; + +using AElementOp = PassThrough; +using BElementOp = PassThrough; +using CElementOp = PassThrough; + +static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default; + +static constexpr bool PermuteA = false; +static constexpr bool PermuteB = false; +static constexpr int KPack = 8; // int4 -> 32, fp8 -> 16, fp16 -> 8 +// clang-format off +using DeviceOpInstance = + ck::tensor_operation::device::DeviceGemm_Wmma_CShuffleV3_BPreshuffle< + ALayout, BLayout, CLayout, + ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, + AElementOp, BElementOp, CElementOp, GemmDefault, + 128, + 32, 128, 128, + 8, 8, + 16, 16, + 2, 2, + S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, + 2, 8, 8, 0, + S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, + 2, 8, 8, 0, + 1, 1, S<1, 16, 1, 8>, S<4, 4, 1>, + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ComputeTypeA, ComputeTypeB, PermuteA, PermuteB>; +// clang-format on + +#include "run_gemm_wmma_bpreshuffle_example.inc" + +int main(int argc, char* argv[]) { return !run_gemm_splitk_example(argc, argv); } diff --git a/example/01_gemm/gemm_wmma_fp8_bpreshuffle.cpp b/example/01_gemm/gemm_wmma_fp8_bpreshuffle.cpp new file mode 100644 index 0000000000..8f8b380b93 --- /dev/null +++ b/example/01_gemm/gemm_wmma_fp8_bpreshuffle.cpp @@ -0,0 +1,72 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "common.hpp" + +#include "ck/ck.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp" +#include "ck/library/utility/check_err.hpp" +#include "ck/library/utility/device_memory.hpp" +#include "ck/library/utility/host_tensor.hpp" +#include "ck/library/utility/host_tensor_generator.hpp" +#include "ck/library/utility/literals.hpp" +#include "ck/stream_config.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3_b_preshuffle.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/utility/amd_ck_fp8.hpp" +#include "ck/utility/data_type.hpp" +#include "ck/utility/get_id.hpp" +#include "ck/utility/scheduler_enum.hpp" + +#include +#include +#include + +using F8 = ck::f8_t; +using F16 = ck::half_t; +using F32 = float; + +using ADataType = F8; +using BDataType = F8; +using AccDataType = F32; +using CShuffleDataType = F32; +using CDataType = F16; +using ComputeTypeA = F8; +using ComputeTypeB = F8; + +using ALayout = Row; +using BLayout = Col; +using CLayout = Row; + +using AElementOp = PassThrough; +using BElementOp = PassThrough; +using CElementOp = PassThrough; + +static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default; + +static constexpr bool PermuteA = false; +static constexpr bool PermuteB = false; +static constexpr int KPack = 16; // int4 -> 32, fp8 -> 16, fp16 -> 8 +// clang-format off +using DeviceOpInstance = + ck::tensor_operation::device::DeviceGemm_Wmma_CShuffleV3_BPreshuffle< + ALayout, BLayout, CLayout, + ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, + AElementOp, BElementOp, CElementOp, GemmDefault, + 256, + 32, 128, 256, + 16, 16, + 16, 16, + 2, 1, + S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, + 2, 16, 16, 0, + S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, + 2, 16, 16, 0, + 1, 1, S<1, 16, 1, 16>, S<8, 8, 1>, + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ComputeTypeA, ComputeTypeB, PermuteA, PermuteB>; +// clang-format on + +#include "run_gemm_wmma_bpreshuffle_example.inc" + +int main(int argc, char* argv[]) { return !run_gemm_splitk_example(argc, argv); } diff --git a/example/01_gemm/gemm_xdl_fp16_fp8_v3.cpp b/example/01_gemm/gemm_xdl_fp16_fp8_v3.cpp index 84ea93ad43..d93e7c9177 100644 --- a/example/01_gemm/gemm_xdl_fp16_fp8_v3.cpp +++ b/example/01_gemm/gemm_xdl_fp16_fp8_v3.cpp @@ -28,10 +28,10 @@ using DeviceGemmV2Instance = ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, AElementOp, BElementOp, CElementOp, GemmDefault, 64, - 16, 16, + 32, 32, 256, 8, 16, 16, 16, - 1, 1, + 2, 2, S<32, 2, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, diff --git a/example/01_gemm/gemm_xdl_fp16_v3.cpp b/example/01_gemm/gemm_xdl_fp16_v3.cpp index e696daf8f0..99691064e2 100644 --- a/example/01_gemm/gemm_xdl_fp16_v3.cpp +++ b/example/01_gemm/gemm_xdl_fp16_v3.cpp @@ -28,10 +28,10 @@ using DeviceGemmV2Instance = ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, PassThrough, PassThrough, PassThrough, GemmDefault, 64, - 16, 16, + 32, 32, 256, 8, 8, 16, 16, - 1, 1, + 2, 2, S<32, 2, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<32, 2, 1>, S<1, 0, 2>, S<1, 0, 2>, diff --git a/example/01_gemm/run_gemm_wmma_bpreshuffle_example.inc b/example/01_gemm/run_gemm_wmma_bpreshuffle_example.inc new file mode 100644 index 0000000000..b1d73cfe10 --- /dev/null +++ b/example/01_gemm/run_gemm_wmma_bpreshuffle_example.inc @@ -0,0 +1,206 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +template +bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) +{ + using namespace ck::literals; + + auto M = problem_size.M; + auto N = problem_size.N; + auto K = problem_size.K; + auto StrideA = problem_size.StrideA; + auto StrideB = problem_size.StrideB; + auto StrideC = problem_size.StrideC; + auto KBatch = problem_size.KBatch; + + auto f_host_tensor_descriptor = + [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { + if constexpr(std::is_same_v) + { + return HostTensorDescriptor({row, col}, {stride, 1_uz}); + } + else + { + return HostTensorDescriptor({row, col}, {1_uz, stride}); + } + }; + + auto f_get_default_stride = + [](std::size_t row, std::size_t col, ck::index_t stride, auto layout) { + if(stride == -1) + { + // give a chance if stride is -1, return a default packed stride + if constexpr(std::is_same_v) + { + return static_cast(col); + } + else + { + return static_cast(row); + } + } + else + return static_cast(stride); + }; + + StrideA = f_get_default_stride(M, K, StrideA, ALayout{}); + StrideB = f_get_default_stride(K, N, StrideB, BLayout{}); + StrideC = f_get_default_stride(M, N, StrideC, CLayout{}); + + Tensor a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{})); + Tensor b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{})); + Tensor b_k_n_preshuffled(f_host_tensor_descriptor(K, N, StrideB, BLayout{})); + + switch(config.init_method) + { + case 0: break; + case 1: + a_m_k.GenerateTensorValue(GeneratorTensor_2{-2, 2}); + b_k_n.GenerateTensorValue(GeneratorTensor_2{0, 2}); + break; + case 2: + a_m_k.GenerateTensorValue(GeneratorTensor_1{}); + b_k_n.GenerateTensorValue(GeneratorTensor_1{}); + break; + default: + a_m_k.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + b_k_n.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + } + + Tensor c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); + Tensor c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); + + std::cout << "a_m_k: " << a_m_k.mDesc << std::endl; + std::cout << "b_k_n: " << b_k_n.mDesc << std::endl; + std::cout << "b_k_n_preshuffled: " << b_k_n_preshuffled.mDesc << std::endl; + std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl; + + DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize()); + DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize()); + DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize()); + + // do GEMM + auto device_op = DeviceOpInstance{}; + + // weight pre-shuffle + int NPerWmma = device_op.GetPreShuffleParameters(); + int KLane = ck::get_warp_size() / NPerWmma; + + int K0 = K / (KLane * KPack); + // K -> K0 KLane KPack + // N -> N0 NPerWmma + // N, K -> N0 K0 KLane NPerWmma KPack + int tempk; + for(int n = 0; n < N; ++n) + { + for(int k = 0; k < K; ++k) + { + int n0 = n / NPerWmma; + int n1 = n % NPerWmma; + + int k0 = k / (KLane * KPack); + tempk = k % (KLane * KPack); + int k1 = tempk / KPack; + int k2 = tempk % KPack; + + int outputIndex = n0 * KPack * NPerWmma * KLane * K0 + k0 * KPack * NPerWmma * KLane + + k1 * KPack * NPerWmma + n1 * KPack + k2; + + b_k_n_preshuffled(outputIndex) = b_k_n(n * K + k); + } + } + + a_m_k_device_buf.ToDevice(a_m_k.mData.data()); + b_k_n_device_buf.ToDevice(b_k_n_preshuffled.mData.data()); + c_m_n_device_buf.ToDevice(c_m_n_device_result.mData.data()); + + auto a_element_op = AElementOp{}; + auto b_element_op = BElementOp{}; + auto c_element_op = CElementOp{}; + + auto invoker = device_op.MakeInvoker(); + + auto argument = + device_op.MakeArgument(static_cast(a_m_k_device_buf.GetDeviceBuffer()), + static_cast(b_k_n_device_buf.GetDeviceBuffer()), + static_cast(c_m_n_device_buf.GetDeviceBuffer()), + M, + N, + K, + StrideA, + StrideB, + StrideC, + KBatch, + a_element_op, + b_element_op, + c_element_op); + + if(!device_op.IsSupportedArgument(argument)) + { + std::cerr << device_op.GetTypeString() << " does not support this problem" << std::endl; + + return true; + } + + float ave_time = + invoker.Run(argument, StreamConfig{nullptr, config.time_kernel, 0, 50, 50, false, 1}); + + bool pass = true; + if(config.do_verification) + { + using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm; + + auto ref_gemm = ReferenceGemmInstance{}; + auto ref_invoker = ref_gemm.MakeInvoker(); + + auto ref_argument = ref_gemm.MakeArgument( + a_m_k, b_k_n, c_m_n_host_result, PassThrough{}, PassThrough{}, PassThrough{}); + + ref_invoker.Run(ref_argument); + + invoker.Run(argument, StreamConfig{nullptr, false, 0}); + c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data()); + + pass &= ck::utils::check_err(c_m_n_device_result, + c_m_n_host_result, + "Error: Incorrect results!", + get_rtol(), + get_atol()); + } + + if(config.time_kernel) + { + ave_time = + invoker.Run(argument, StreamConfig{nullptr, config.time_kernel, 0, 20, 50, true, 50}); + + std::size_t flop = 2_uz * M * N * K; + std::size_t num_btype = + sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N; + + float tflops = static_cast(flop) / 1.E9 / ave_time; + + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec + << " GB/s, " << device_op.GetTypeString() << std::endl; + } + + return pass; +} + +bool run_gemm_splitk_example(int argc, char* argv[]) +{ + ProblemSizeSplitK problem_size{3840, 4096, 4096, 4096, 4096, 4096, 1}; + ExecutionConfig config; + + return parse_cmd_args(argc, argv, problem_size, config) && run_gemm(problem_size, config); +} diff --git a/example/14_gemm_quantization/gemm_wmma_quantization_int8.cpp b/example/14_gemm_quantization/gemm_wmma_quantization_int8.cpp index cc5e3616ff..7437d0be9d 100644 --- a/example/14_gemm_quantization/gemm_wmma_quantization_int8.cpp +++ b/example/14_gemm_quantization/gemm_wmma_quantization_int8.cpp @@ -27,10 +27,11 @@ using ::ck::Tensor; template using S = ck::Sequence; -using I8 = int8_t; -using I32 = int32_t; -using Row = ck::tensor_layout::gemm::RowMajor; -using Col = ck::tensor_layout::gemm::ColumnMajor; +using I8 = int8_t; +using I32 = int32_t; +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; +using Bypass = ck::tensor_layout::BypassLayoutVerification; using PassThrough = ck::tensor_operation::element_wise::PassThrough; using ActivationOp = PassThrough; @@ -125,11 +126,11 @@ int main(int /* argc */, char* /* argv */[]) if(std::is_same::value) { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); + return HostTensorDescriptor({row, col}, {stride, 1_uz}, Bypass{}); } else { - return HostTensorDescriptor({row, col}, {1_uz, stride}); + return HostTensorDescriptor({row, col}, {1_uz, stride}, Bypass{}); } }; diff --git a/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m2n3k1_xdl_fp16.cpp b/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m2n3k1_xdl_fp16.cpp index f7663cbd0a..6295cfdd04 100644 --- a/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m2n3k1_xdl_fp16.cpp +++ b/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m2n3k1_xdl_fp16.cpp @@ -17,7 +17,7 @@ #include "ck/library/utility/host_tensor_generator.hpp" #include "ck/library/utility/numeric.hpp" -#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_contraction.hpp" using ::ck::DeviceMem; using ::ck::HostTensorDescriptor; @@ -69,142 +69,6 @@ using DeviceOpInstanceKKNN = ck::tensor_operation::device:: using DeviceOpInstance = DeviceOpInstanceKKNN; -// hardcoded for NumDimM == NumDimN == NumDimK == 2 -template = - false> -struct ReferenceContraction_G1_M2_N3_K1 : public ck::tensor_operation::device::BaseOperator -{ - // Argument - struct Argument : public ck::tensor_operation::device::BaseArgument - { - Argument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - : a_gs_ms_ks_{a_gs_ms_ks}, - b_gs_ns_ks_{b_gs_ns_ks}, - e_gs_ms_ns_{e_gs_ms_ns}, - a_element_op_{a_element_op}, - b_element_op_{b_element_op}, - cde_element_op_{cde_element_op} - { - } - - const Tensor& a_gs_ms_ks_; - const Tensor& b_gs_ns_ks_; - Tensor& e_gs_ms_ns_; - - AElementwiseOperation a_element_op_; - BElementwiseOperation b_element_op_; - CDEElementwiseOperation cde_element_op_; - }; - - // Invoker - struct Invoker : public ck::tensor_operation::device::BaseInvoker - { - using Argument = ReferenceContraction_G1_M2_N3_K1::Argument; - - float Run(const Argument& arg) - { - auto f_gs_ms_ns = [&](auto g0, auto m0, auto m1, auto n0, auto n1, auto n2) { - const int K0 = arg.a_gs_ms_ks_.mDesc.GetLengths()[3]; - - AccDataType v_acc = 0; - - for(int k0 = 0; k0 < K0; ++k0) - { - AccDataType v_a; - AccDataType v_b; - - arg.a_element_op_( - v_a, ck::type_convert(arg.a_gs_ms_ks_(g0, m0, m1, k0))); - arg.b_element_op_( - v_b, - ck::type_convert(arg.b_gs_ns_ks_(g0, n0, n1, n2, k0))); - - v_acc += v_a * v_b; - } - - AccDataType v_c; - - arg.cde_element_op_(v_c, v_acc); - - arg.e_gs_ms_ns_(g0, m0, m1, n0, n1, n2) = v_c; - }; - - make_ParallelTensorFunctor(f_gs_ms_ns, - arg.e_gs_ms_ns_.mDesc.GetLengths()[0], - arg.e_gs_ms_ns_.mDesc.GetLengths()[1], - arg.e_gs_ms_ns_.mDesc.GetLengths()[2], - arg.e_gs_ms_ns_.mDesc.GetLengths()[3], - arg.e_gs_ms_ns_.mDesc.GetLengths()[4], - arg.e_gs_ms_ns_.mDesc.GetLengths()[5])( - std::thread::hardware_concurrency()); - - return 0; - } - - float Run(const ck::tensor_operation::device::BaseArgument* p_arg, - const StreamConfig& /* stream_config */ = StreamConfig{}) override - { - return Run(*dynamic_cast(p_arg)); - } - }; - - static constexpr bool IsValidCompilationParameter() - { - // TODO: properly implement this check - return true; - } - - bool IsSupportedArgument(const ck::tensor_operation::device::BaseArgument*) override - { - return true; - } - - static auto MakeArgument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - { - return Argument{ - a_gs_ms_ks, b_gs_ns_ks, e_gs_ms_ns, a_element_op, b_element_op, cde_element_op}; - } - - static auto MakeInvoker() { return Invoker{}; } - - virtual std::unique_ptr MakeInvokerPointer() - { - return std::make_unique(Invoker{}); - } - - std::string GetTypeString() const override - { - auto str = std::stringstream(); - - // clang-format off - str << "ReferenceContraction_M3_N2_K1" - << std::endl; - // clang-format on - - return str.str(); - } -}; - int main(int argc, char* argv[]) { bool do_verification = true; @@ -353,16 +217,18 @@ int main(int argc, char* argv[]) Tensor c_gs_ms_ns_host_result( e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - using ReferenceOpInstance = ReferenceContraction_G1_M2_N3_K1; + using ReferenceOpInstance = + ck::tensor_operation::host::ReferenceBatchedContraction_G1_M2_N3_K1; auto ref_gemm = ReferenceOpInstance{}; auto ref_invoker = ref_gemm.MakeInvoker(); @@ -399,7 +265,13 @@ int main(int argc, char* argv[]) } } - return ck::utils::check_err(e_gs_ms_ns_device_result, e_gs_ms_ns_host_result) ? 0 : 1; + bool pass = ck::utils::check_err(e_gs_ms_ns_device_result, e_gs_ms_ns_host_result); + std::cout << "Verification: " << (pass ? "SUCCESS" : "FAILURE") << "!" << std::endl; + + if(!pass) + { + return 1; + } } return 0; diff --git a/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m3n2k1_xdl_fp16.cpp b/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m3n2k1_xdl_fp16.cpp index 736dc09867..3adfecc7ae 100644 --- a/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m3n2k1_xdl_fp16.cpp +++ b/example/25_gemm_bias_e_permute/gemm_bias_e_permute_g1m3n2k1_xdl_fp16.cpp @@ -17,6 +17,8 @@ #include "ck/library/utility/host_tensor_generator.hpp" #include "ck/library/utility/numeric.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_contraction.hpp" + using ::ck::DeviceMem; using ::ck::HostTensorDescriptor; using ::ck::make_ParallelTensorFunctor; @@ -67,142 +69,6 @@ using DeviceOpInstanceKKNN = ck::tensor_operation::device:: using DeviceOpInstance = DeviceOpInstanceKKNN; -template = - false> -struct ReferenceContraction_G1_M3_N2_K1 : public ck::tensor_operation::device::BaseOperator -{ - // Argument - struct Argument : public ck::tensor_operation::device::BaseArgument - { - Argument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - : a_gs_ms_ks_{a_gs_ms_ks}, - b_gs_ns_ks_{b_gs_ns_ks}, - e_gs_ms_ns_{e_gs_ms_ns}, - a_element_op_{a_element_op}, - b_element_op_{b_element_op}, - cde_element_op_{cde_element_op} - { - } - - const Tensor& a_gs_ms_ks_; - const Tensor& b_gs_ns_ks_; - Tensor& e_gs_ms_ns_; - - AElementwiseOperation a_element_op_; - BElementwiseOperation b_element_op_; - CDEElementwiseOperation cde_element_op_; - }; - - // Invoker - struct Invoker : public ck::tensor_operation::device::BaseInvoker - { - using Argument = ReferenceContraction_G1_M3_N2_K1::Argument; - - float Run(const Argument& arg) - { - auto f_gs_ms_ns = [&](auto g0, auto m0, auto m1, auto m2, auto n0, auto n1) { - const int K0 = arg.a_gs_ms_ks_.mDesc.GetLengths()[4]; - - AccDataType v_acc = 0; - - for(int k0 = 0; k0 < K0; ++k0) - { - AccDataType v_a; - AccDataType v_b; - - arg.a_element_op_( - v_a, - ck::type_convert(arg.a_gs_ms_ks_(g0, m0, m1, m2, k0))); - arg.b_element_op_( - v_b, ck::type_convert(arg.b_gs_ns_ks_(g0, n0, n1, k0))); - - v_acc += v_a * v_b; - } - - AccDataType v_c; - - arg.cde_element_op_(v_c, v_acc); - - arg.e_gs_ms_ns_(g0, m0, m1, m2, n0, n1) = v_c; - }; - - make_ParallelTensorFunctor(f_gs_ms_ns, - arg.e_gs_ms_ns_.mDesc.GetLengths()[0], - arg.e_gs_ms_ns_.mDesc.GetLengths()[1], - arg.e_gs_ms_ns_.mDesc.GetLengths()[2], - arg.e_gs_ms_ns_.mDesc.GetLengths()[3], - arg.e_gs_ms_ns_.mDesc.GetLengths()[4], - arg.e_gs_ms_ns_.mDesc.GetLengths()[5])( - std::thread::hardware_concurrency()); - - return 0; - } - - float Run(const ck::tensor_operation::device::BaseArgument* p_arg, - const StreamConfig& /* stream_config */ = StreamConfig{}) override - { - return Run(*dynamic_cast(p_arg)); - } - }; - - static constexpr bool IsValidCompilationParameter() - { - // TODO: properly implement this check - return true; - } - - bool IsSupportedArgument(const ck::tensor_operation::device::BaseArgument*) override - { - return true; - } - - static auto MakeArgument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - { - return Argument{ - a_gs_ms_ks, b_gs_ns_ks, e_gs_ms_ns, a_element_op, b_element_op, cde_element_op}; - } - - static auto MakeInvoker() { return Invoker{}; } - - virtual std::unique_ptr MakeInvokerPointer() - { - return std::make_unique(Invoker{}); - } - - std::string GetTypeString() const override - { - auto str = std::stringstream(); - - // clang-format off - str << "ReferenceContraction_G1_M3_N2_K1" - << std::endl; - // clang-format on - - return str.str(); - } -}; - int main(int argc, char* argv[]) { bool do_verification = true; @@ -353,17 +219,18 @@ int main(int argc, char* argv[]) Tensor c_gs_ms_ns_host_result( e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - using ReferenceOpInstance = ReferenceContraction_G1_M3_N2_K1; + using ReferenceOpInstance = + ck::tensor_operation::host::ReferenceBatchedContraction_G1_M3_N2_K1; auto ref_gemm = ReferenceOpInstance{}; auto ref_invoker = ref_gemm.MakeInvoker(); @@ -400,7 +267,13 @@ int main(int argc, char* argv[]) } } - return ck::utils::check_err(e_gs_ms_ns_device_result, e_gs_ms_ns_host_result) ? 0 : 1; + bool pass = ck::utils::check_err(e_gs_ms_ns_device_result, e_gs_ms_ns_host_result); + std::cout << "Verification: " << (pass ? "SUCCESS" : "FAILURE") << "!" << std::endl; + + if(!pass) + { + return 1; + } } return 0; diff --git a/example/26_contraction/CMakeLists.txt b/example/26_contraction/CMakeLists.txt index 4a41bc5e65..c8a516bae6 100644 --- a/example/26_contraction/CMakeLists.txt +++ b/example/26_contraction/CMakeLists.txt @@ -38,16 +38,28 @@ add_example_executable(example_contraction_scale_xdl_fp64_compute_fp32 contracti add_example_dependencies(example_contraction_scale example_contraction_scale_xdl_fp64_compute_fp32) # FP16 +add_example_executable(example_contraction_bilinear_xdl_fp16 contraction_bilinear_xdl_fp16.cpp) +add_example_dependencies(example_contraction_bilinear example_contraction_bilinear_xdl_fp16) + add_example_executable(example_contraction_bilinear_xdl_fp16_compute_fp32 contraction_bilinear_xdl_fp16_compute_fp32.cpp) add_example_dependencies(example_contraction_bilinear example_contraction_bilinear_xdl_fp16_compute_fp32) +add_example_executable(example_contraction_scale_xdl_fp16 contraction_scale_xdl_fp16.cpp) +add_example_dependencies(example_contraction_scale example_contraction_scale_xdl_fp16) + add_example_executable(example_contraction_scale_xdl_fp16_compute_fp32 contraction_scale_xdl_fp16_compute_fp32.cpp) add_example_dependencies(example_contraction_scale example_contraction_scale_xdl_fp16_compute_fp32) # BF16 +add_example_executable(example_contraction_bilinear_xdl_bf16 contraction_bilinear_xdl_bf16.cpp) +add_example_dependencies(example_contraction_bilinear example_contraction_bilinear_xdl_bf16) + add_example_executable(example_contraction_bilinear_xdl_bf16_compute_fp32 contraction_bilinear_xdl_bf16_compute_fp32.cpp) add_example_dependencies(example_contraction_bilinear example_contraction_bilinear_xdl_bf16_compute_fp32) +add_example_executable(example_contraction_scale_xdl_bf16 contraction_scale_xdl_bf16.cpp) +add_example_dependencies(example_contraction_scale example_contraction_scale_xdl_bf16) + add_example_executable(example_contraction_scale_xdl_bf16_compute_fp32 contraction_scale_xdl_bf16_compute_fp32.cpp) add_example_dependencies(example_contraction_scale example_contraction_scale_xdl_bf16_compute_fp32) diff --git a/example/26_contraction/contraction_bilinear_xdl_bf16.cpp b/example/26_contraction/contraction_bilinear_xdl_bf16.cpp new file mode 100644 index 0000000000..8899b54fbf --- /dev/null +++ b/example/26_contraction/contraction_bilinear_xdl_bf16.cpp @@ -0,0 +1,86 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" + +#include "common_instances.hpp" + +using ADataType = BF16; +using BDataType = BF16; +using AccDataType = F32; +using CShuffleDataType = BF16; +using DDataType = BF16; +using DsDataType = ck::Tuple; +using EDataType = BF16; +using ComputeDataType = BF16; + +static constexpr ck::index_t NumDimM = 2; +static constexpr ck::index_t NumDimN = 2; +static constexpr ck::index_t NumDimK = 2; + +using AElementOp = ck::tensor_operation::element_wise::PassThrough; +using BElementOp = ck::tensor_operation::element_wise::PassThrough; +using CDEElementOp = ck::tensor_operation::element_wise::Bilinear; + +using DeviceOpInstanceKKNN = DeviceOpInstanceKK_Generic; + +using DeviceOpInstanceKNNN = DeviceOpInstanceKN_Generic; + +using DeviceOpInstanceMKNN = DeviceOpInstanceMK_Generic; + +using DeviceOpInstanceMNNN = DeviceOpInstanceMN_Generic; + +using DeviceOpInstance = DeviceOpInstanceKKNN; + +#include "run_contraction_bilinear_example.inc" + +int main(int argc, char* argv[]) { return run_contraction_bilinear_example(argc, argv); } diff --git a/example/26_contraction/contraction_bilinear_xdl_fp16.cpp b/example/26_contraction/contraction_bilinear_xdl_fp16.cpp new file mode 100644 index 0000000000..16e33e0886 --- /dev/null +++ b/example/26_contraction/contraction_bilinear_xdl_fp16.cpp @@ -0,0 +1,86 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" + +#include "common_instances.hpp" + +using ADataType = F16; +using BDataType = F16; +using AccDataType = F32; +using CShuffleDataType = F16; +using DDataType = F16; +using DsDataType = ck::Tuple; +using EDataType = F16; +using ComputeDataType = F16; + +static constexpr ck::index_t NumDimM = 2; +static constexpr ck::index_t NumDimN = 2; +static constexpr ck::index_t NumDimK = 2; + +using AElementOp = ck::tensor_operation::element_wise::PassThrough; +using BElementOp = ck::tensor_operation::element_wise::PassThrough; +using CDEElementOp = ck::tensor_operation::element_wise::Bilinear; + +using DeviceOpInstanceKKNN = DeviceOpInstanceKK_Generic; + +using DeviceOpInstanceKNNN = DeviceOpInstanceKN_Generic; + +using DeviceOpInstanceMKNN = DeviceOpInstanceMK_Generic; + +using DeviceOpInstanceMNNN = DeviceOpInstanceMN_Generic; + +using DeviceOpInstance = DeviceOpInstanceKKNN; + +#include "run_contraction_bilinear_example.inc" + +int main(int argc, char* argv[]) { return run_contraction_bilinear_example(argc, argv); } diff --git a/example/26_contraction/contraction_scale_xdl_bf16.cpp b/example/26_contraction/contraction_scale_xdl_bf16.cpp new file mode 100644 index 0000000000..586b022397 --- /dev/null +++ b/example/26_contraction/contraction_scale_xdl_bf16.cpp @@ -0,0 +1,85 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" + +#include "common_instances.hpp" + +using ADataType = BF16; +using BDataType = BF16; +using AccDataType = F32; +using CShuffleDataType = BF16; +using DsDataType = ck::Tuple<>; +using EDataType = BF16; +using ComputeDataType = BF16; + +static constexpr ck::index_t NumDimM = 2; +static constexpr ck::index_t NumDimN = 2; +static constexpr ck::index_t NumDimK = 2; + +using AElementOp = ck::tensor_operation::element_wise::PassThrough; +using BElementOp = ck::tensor_operation::element_wise::PassThrough; +using CDEElementOp = ck::tensor_operation::element_wise::Scale; + +using DeviceOpInstanceKKN = DeviceOpInstanceKK_Generic; + +using DeviceOpInstanceKNN = DeviceOpInstanceKN_Generic; + +using DeviceOpInstanceMKN = DeviceOpInstanceMK_Generic; + +using DeviceOpInstanceMNN = DeviceOpInstanceMN_Generic; + +using DeviceOpInstance = DeviceOpInstanceKKN; + +#include "run_contraction_scale_example.inc" + +int main(int argc, char* argv[]) { return run_contraction_scale_example(argc, argv); } diff --git a/example/26_contraction/contraction_scale_xdl_fp16.cpp b/example/26_contraction/contraction_scale_xdl_fp16.cpp new file mode 100644 index 0000000000..1f29e16223 --- /dev/null +++ b/example/26_contraction/contraction_scale_xdl_fp16.cpp @@ -0,0 +1,85 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" + +#include "common_instances.hpp" + +using ADataType = F16; +using BDataType = F16; +using AccDataType = F32; +using CShuffleDataType = F16; +using DsDataType = ck::Tuple<>; +using EDataType = F16; +using ComputeDataType = F16; + +static constexpr ck::index_t NumDimM = 2; +static constexpr ck::index_t NumDimN = 2; +static constexpr ck::index_t NumDimK = 2; + +using AElementOp = ck::tensor_operation::element_wise::PassThrough; +using BElementOp = ck::tensor_operation::element_wise::PassThrough; +using CDEElementOp = ck::tensor_operation::element_wise::Scale; + +using DeviceOpInstanceKKN = DeviceOpInstanceKK_Generic; + +using DeviceOpInstanceKNN = DeviceOpInstanceKN_Generic; + +using DeviceOpInstanceMKN = DeviceOpInstanceMK_Generic; + +using DeviceOpInstanceMNN = DeviceOpInstanceMN_Generic; + +using DeviceOpInstance = DeviceOpInstanceKKN; + +#include "run_contraction_scale_example.inc" + +int main(int argc, char* argv[]) { return run_contraction_scale_example(argc, argv); } diff --git a/example/26_contraction/run_contraction_bilinear_example.inc b/example/26_contraction/run_contraction_bilinear_example.inc index 69eb42defd..08ed098b66 100644 --- a/example/26_contraction/run_contraction_bilinear_example.inc +++ b/example/26_contraction/run_contraction_bilinear_example.inc @@ -235,13 +235,20 @@ int run_contraction_bilinear_example(int argc, char* argv[]) if(ck::is_gfx11_supported()) { - return ck::utils::check_err(e_ms_ns_device_result, - e_ms_ns_host_result, - "Error: Incorrect results!", - 1e-4, - 1e-4) - ? 0 - : 1; + if constexpr(std::is_same_v) + { + return ck::utils::check_err(e_ms_ns_device_result, + e_ms_ns_host_result, + "Error: Incorrect results!", + 1e-4, + 1e-4) + ? 0 + : 1; + } + else + { + return ck::utils::check_err(e_ms_ns_device_result, e_ms_ns_host_result) ? 0 : 1; + } } else { diff --git a/example/26_contraction/run_contraction_scale_example.inc b/example/26_contraction/run_contraction_scale_example.inc index a7451fab71..a5bcd8d447 100644 --- a/example/26_contraction/run_contraction_scale_example.inc +++ b/example/26_contraction/run_contraction_scale_example.inc @@ -218,13 +218,20 @@ int run_contraction_scale_example(int argc, char* argv[]) if(ck::is_gfx11_supported()) { - return ck::utils::check_err(e_ms_ns_device_result, - e_ms_ns_host_result, - "Error: Incorrect results!", - 1e-4, - 1e-4) - ? 0 - : 1; + if constexpr(std::is_same_v) + { + return ck::utils::check_err(e_ms_ns_device_result, + e_ms_ns_host_result, + "Error: Incorrect results!", + 1e-4, + 1e-4) + ? 0 + : 1; + } + else + { + return ck::utils::check_err(e_ms_ns_device_result, e_ms_ns_host_result) ? 0 : 1; + } } else { diff --git a/example/29_batched_gemm_bias_e_permute/CMakeLists.txt b/example/29_batched_gemm_bias_e_permute/CMakeLists.txt index d5d5521370..6cf93215f8 100644 --- a/example/29_batched_gemm_bias_e_permute/CMakeLists.txt +++ b/example/29_batched_gemm_bias_e_permute/CMakeLists.txt @@ -3,3 +3,4 @@ add_example_executable(example_batched_gemm_bias_e_permute_xdl_fp16 batched_gemm_bias_e_permute_xdl_fp16.cpp) add_example_executable(example_batched_gemm_bias_e_permute_wmma_fp16 batched_gemm_bias_e_permute_wmma_fp16.cpp) +add_example_executable(example_batched_gemm_bias_e_permute_wmma_v3_fp16 batched_gemm_bias_e_permute_wmma_v3_fp16.cpp) diff --git a/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_wmma_fp16.cpp b/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_wmma_fp16.cpp index 06bf971ac4..f102a0b132 100644 --- a/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_wmma_fp16.cpp +++ b/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_wmma_fp16.cpp @@ -106,352 +106,5 @@ using DeviceOpInstanceKKNN = using DeviceOpInstance = DeviceOpInstanceKKNN; -// hardcoded for NumDimM == NumDimN == NumDimK == 2 -template = - false> -struct ReferenceContraction_G2_M2_N2_K1 : public ck::tensor_operation::device::BaseOperator -{ - // Argument - struct Argument : public ck::tensor_operation::device::BaseArgument - { - Argument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - : a_gs_ms_ks_{a_gs_ms_ks}, - b_gs_ns_ks_{b_gs_ns_ks}, - e_gs_ms_ns_{e_gs_ms_ns}, - a_element_op_{a_element_op}, - b_element_op_{b_element_op}, - cde_element_op_{cde_element_op} - { - } - - const Tensor& a_gs_ms_ks_; - const Tensor& b_gs_ns_ks_; - Tensor& e_gs_ms_ns_; - - AElementwiseOperation a_element_op_; - BElementwiseOperation b_element_op_; - CDEElementwiseOperation cde_element_op_; - }; - - // Invoker - struct Invoker : public ck::tensor_operation::device::BaseInvoker - { - using Argument = ReferenceContraction_G2_M2_N2_K1::Argument; - - float Run(const Argument& arg) - { - auto f_ms_ns = [&](auto g0, auto g1, auto m0, auto m1, auto n0, auto n1) { - const int K0 = arg.a_gs_ms_ks_.mDesc.GetLengths()[4]; - - AccDataType v_acc = 0; - - for(int k0 = 0; k0 < K0; ++k0) - { - AccDataType v_a; - AccDataType v_b; - - arg.a_element_op_( - v_a, - ck::type_convert(arg.a_gs_ms_ks_(g0, g1, m0, m1, k0))); - arg.b_element_op_( - v_b, - ck::type_convert(arg.b_gs_ns_ks_(g0, g1, n0, n1, k0))); - - v_acc += v_a * v_b; - } - - AccDataType v_c; - - arg.cde_element_op_(v_c, v_acc); - - arg.e_gs_ms_ns_(g0, g1, m0, m1, n0, n1) = v_c; - }; - - make_ParallelTensorFunctor(f_ms_ns, - arg.e_gs_ms_ns_.mDesc.GetLengths()[0], - arg.e_gs_ms_ns_.mDesc.GetLengths()[1], - arg.e_gs_ms_ns_.mDesc.GetLengths()[2], - arg.e_gs_ms_ns_.mDesc.GetLengths()[3], - arg.e_gs_ms_ns_.mDesc.GetLengths()[4], - arg.e_gs_ms_ns_.mDesc.GetLengths()[5])( - std::thread::hardware_concurrency()); - - return 0; - } - - float Run(const ck::tensor_operation::device::BaseArgument* p_arg, - const StreamConfig& /* stream_config */ = StreamConfig{}) override - { - return Run(*dynamic_cast(p_arg)); - } - }; - - static constexpr bool IsValidCompilationParameter() - { - // TODO: properly implement this check - return true; - } - - bool IsSupportedArgument(const ck::tensor_operation::device::BaseArgument*) override - { - return true; - } - - static auto MakeArgument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - { - return Argument{ - a_gs_ms_ks, b_gs_ns_ks, e_gs_ms_ns, a_element_op, b_element_op, cde_element_op}; - } - - static auto MakeInvoker() { return Invoker{}; } - - virtual std::unique_ptr MakeInvokerPointer() - { - return std::make_unique(Invoker{}); - } - - std::string GetTypeString() const override - { - auto str = std::stringstream(); - - // clang-format off - str << "ReferenceContraction_G2_M2_N2_K1" - << std::endl; - // clang-format on - - return str.str(); - } -}; - -int main(int argc, char* argv[]) -{ - bool do_verification = true; - int init_method = 1; - bool time_kernel = false; - - ck::index_t G0 = 1; - ck::index_t G1 = 2; - - ck::index_t M0 = 4; - ck::index_t M1 = 128; - - ck::index_t N0 = 16; - ck::index_t N1 = 256; - - ck::index_t K0 = 2048; - - if(argc == 1) - { - // use default case - } - else if(argc == 4) - { - do_verification = std::stoi(argv[1]); - init_method = std::stoi(argv[2]); - time_kernel = std::stoi(argv[3]); - } - else if(argc == 11) - { - do_verification = std::stoi(argv[1]); - init_method = std::stoi(argv[2]); - time_kernel = std::stoi(argv[3]); - G0 = std::stoi(argv[4]); - G1 = std::stoi(argv[5]); - M0 = std::stoi(argv[6]); - M1 = std::stoi(argv[7]); - N0 = std::stoi(argv[8]); - N1 = std::stoi(argv[9]); - K0 = std::stoi(argv[10]); - } - else - { - printf("arg1: verification (0=no, 1=yes)\n"); - printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"); - printf("arg3: time kernel (0=no, 1=yes)\n"); - printf("arg4-10: G0, G1, M0, M1, N0, N1, K0\n"); - exit(0); - } - - // A[G0, G1, M0, M1, K0] - std::vector a_gs_ms_ks_lengths{G0, G1, M0, M1, K0}; - std::vector a_gs_ms_ks_strides{G1 * M0 * M1 * K0, M0 * M1 * K0, M1 * K0, K0, 1}; - // B[G0, G1, N0, N1, K0] - std::vector b_gs_ns_ks_lengths{G0, G1, N0, N1, K0}; - std::vector b_gs_ns_ks_strides{G1 * N0 * N1 * K0, N0 * N1 * K0, N1 * K0, K0, 1}; - - // D[G0, G1, M0, N0, M1, N1] - std::vector d_gs_ms_ns_lengths{G0, G1, M0, M1, N0, N1}; - std::vector d_gs_ms_ns_strides{G1 * N0 * N1, N0 * N1, 0, 0, N1, 1}; - // E[G0, G1, M0, N0, M1, N1] - std::vector e_gs_ms_ns_lengths{G0, G1, M0, M1, N0, N1}; - std::vector e_gs_ms_ns_strides{ - G1 * M0 * N0 * M1 * N1, M0 * N0 * M1 * N1, N0 * M1 * N1, N1, M1 * N1, 1}; - - Tensor a_gs_ms_ks(a_gs_ms_ks_lengths, a_gs_ms_ks_strides, Row{}); - Tensor b_gs_ns_ks(b_gs_ns_ks_lengths, b_gs_ns_ks_strides, Row{}); - Tensor d_gs_ms_ns(d_gs_ms_ns_lengths, d_gs_ms_ns_strides, Bypass{}); - Tensor e_gs_ms_ns_host_result(e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - Tensor e_gs_ms_ns_device_result(e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - std::cout << "a_gs_ms_ks: " << a_gs_ms_ks.mDesc << std::endl; - std::cout << "b_gs_ns_ks: " << b_gs_ns_ks.mDesc << std::endl; - std::cout << "d_gs_ms_ns: " << d_gs_ms_ns.mDesc << std::endl; - std::cout << "e_gs_ms_ns: " << e_gs_ms_ns_host_result.mDesc << std::endl; - - switch(init_method) - { - case 0: break; - case 1: - a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_2{-5, 5}); - b_gs_ns_ks.GenerateTensorValue(GeneratorTensor_2{-5, 5}); - d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_2{-5, 5}); - break; - default: - a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - b_gs_ns_ks.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); - d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); - break; - } - DeviceMem a_device_buf(sizeof(ADataType) * a_gs_ms_ks.mDesc.GetElementSpaceSize()); - DeviceMem b_device_buf(sizeof(BDataType) * b_gs_ns_ks.mDesc.GetElementSpaceSize()); - DeviceMem d_device_buf(sizeof(DDataType) * d_gs_ms_ns.mDesc.GetElementSpaceSize()); - DeviceMem e_device_buf(sizeof(EDataType) * - e_gs_ms_ns_device_result.mDesc.GetElementSpaceSize()); - - a_device_buf.ToDevice(a_gs_ms_ks.mData.data()); - b_device_buf.ToDevice(b_gs_ns_ks.mData.data()); - d_device_buf.ToDevice(d_gs_ms_ns.mData.data()); - - // set zero - e_device_buf.SetZero(); - - auto a_element_op = AElementOp{}; - auto b_element_op = BElementOp{}; - auto cde_element_op = CDEElementOp{}; - - // device operation - auto op = DeviceOpInstance{}; - auto invoker = op.MakeInvoker(); - auto argument = op.MakeArgument(a_device_buf.GetDeviceBuffer(), - b_device_buf.GetDeviceBuffer(), - std::array{d_device_buf.GetDeviceBuffer()}, - e_device_buf.GetDeviceBuffer(), - a_gs_ms_ks_lengths, - a_gs_ms_ks_strides, - b_gs_ns_ks_lengths, - b_gs_ns_ks_strides, - std::array, 1>{d_gs_ms_ns_lengths}, - std::array, 1>{d_gs_ms_ns_strides}, - e_gs_ms_ns_lengths, - e_gs_ms_ns_strides, - a_element_op, - b_element_op, - cde_element_op); - - if(!op.IsSupportedArgument(argument)) - { - std::cout << op.GetTypeString() << " does not support this problem" << std::endl; - - return 0; - } - - float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel}); - - ck::index_t G = - ck::accumulate_n(e_gs_ms_ns_lengths.begin(), NumDimG, 1, std::multiplies<>{}); - - ck::index_t M = ck::accumulate_n( - e_gs_ms_ns_lengths.begin() + NumDimG, NumDimM, 1, std::multiplies<>{}); - - ck::index_t N = ck::accumulate_n( - e_gs_ms_ns_lengths.begin() + NumDimG + NumDimM, NumDimN, 1, std::multiplies<>{}); - - ck::index_t K = ck::accumulate_n( - a_gs_ms_ks_lengths.begin() + NumDimG + NumDimM, NumDimK, 1, std::multiplies<>{}); - std::cout << "GMNK=" << G << ", " << M << ", " << N << ", " << K << std::endl; - std::size_t flop = std::size_t(2) * G * M * N * K; - std::size_t num_btype = sizeof(ADataType) * G * M * K + sizeof(BDataType) * G * K * N + - sizeof(DDataType) * G * M * N + sizeof(EDataType) * G * M * N; - - float tflops = static_cast(flop) / 1.E9 / ave_time; - - float gb_per_sec = num_btype / 1.E6 / ave_time; - - std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " - << op.GetTypeString() << std::endl; - - e_device_buf.FromDevice(e_gs_ms_ns_device_result.mData.data()); - - if(do_verification) - { - Tensor c_ms_ns_host_result( - e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - - using ReferenceOpInstance = ReferenceContraction_G2_M2_N2_K1; - - auto ref_gemm = ReferenceOpInstance{}; - auto ref_invoker = ref_gemm.MakeInvoker(); - - auto ref_argument = ref_gemm.MakeArgument( - a_gs_ms_ks, b_gs_ns_ks, c_ms_ns_host_result, a_element_op, b_element_op, PassThrough{}); - - ref_invoker.Run(ref_argument); - - for(size_t g0 = 0; g0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[0]; ++g0) - { - for(size_t g1 = 0; g1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[1]; ++g1) - { - for(size_t m0 = 0; m0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[2]; ++m0) - { - for(size_t m1 = 0; m1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[3]; ++m1) - { - for(size_t n0 = 0; n0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[4]; ++n0) - { - for(size_t n1 = 0; n1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[5]; - ++n1) - { - cde_element_op(e_gs_ms_ns_host_result(g0, g1, m0, m1, n0, n1), - c_ms_ns_host_result(g0, g1, m0, m1, n0, n1), - d_gs_ms_ns(g0, g1, m0, m1, n0, n1)); - } - } - } - } - } - } - - return ck::utils::check_err(e_gs_ms_ns_device_result, e_gs_ms_ns_host_result) ? 0 : 1; - } - - return 0; -} +#include "run_batched_gemm_bias_e_permute_example.inc" +int main(int argc, char* argv[]) { return !run_batched_gemm_bias_e_permute_example(argc, argv); } diff --git a/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_wmma_v3_fp16.cpp b/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_wmma_v3_fp16.cpp new file mode 100644 index 0000000000..4e34f18b8b --- /dev/null +++ b/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_wmma_v3_fp16.cpp @@ -0,0 +1,111 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_batched_contraction_multiple_d_wmma_cshuffle_v3.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" + +#include "ck/library/utility/check_err.hpp" +#include "ck/library/utility/device_memory.hpp" +#include "ck/library/utility/host_tensor.hpp" +#include "ck/library/utility/host_tensor_generator.hpp" +#include "ck/library/utility/numeric.hpp" + +using ::ck::DeviceMem; +using ::ck::HostTensorDescriptor; +using ::ck::make_ParallelTensorFunctor; +using ::ck::Tensor; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Bypass = ck::tensor_layout::BypassLayoutVerification; + +template +using S = ck::Sequence; + +using F16 = ck::half_t; +using F32 = float; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; +using Add = ck::tensor_operation::element_wise::Add; + +using ADataType = F16; +using BDataType = F16; +using AccDataType = F32; +using CShuffleDataType = F16; +using DDataType = F16; +using DsDataType = ck::Tuple; +using EDataType = F16; + +static constexpr ck::index_t NumDimG = 2; +static constexpr ck::index_t NumDimM = 2; +static constexpr ck::index_t NumDimN = 2; +static constexpr ck::index_t NumDimK = 1; + +using AElementOp = ck::tensor_operation::element_wise::PassThrough; +using BElementOp = ck::tensor_operation::element_wise::PassThrough; +using CDEElementOp = ck::tensor_operation::element_wise::Add; + +static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKPadding; + +static constexpr auto ASpec = ck::tensor_operation::device::TensorSpecialization::Default; +static constexpr auto BSpec = ck::tensor_operation::device::TensorSpecialization::Default; +static constexpr auto DESpec = ck::tensor_operation::device::TensorSpecialization::Default; + +using DeviceOpInstanceKKNN = + ck::tensor_operation::device::DeviceBatchedContractionMultipleD_Wmma_CShuffle_V3< + NumDimG, + NumDimM, + NumDimN, + NumDimK, + ADataType, + BDataType, + AccDataType, + CShuffleDataType, + DsDataType, + EDataType, + AElementOp, + BElementOp, + CDEElementOp, + GemmSpec, + ASpec, + BSpec, + DESpec, + 128, + 64, + 64, + 64, + 4, + 4, + 16, + 16, + 1, + 4, + S<4, 32, 1>, + S<1, 0, 2>, + S<1, 0, 2>, + 2, + 4, + 4, + false, + S<4, 32, 1>, + S<1, 0, 2>, + S<1, 0, 2>, + 2, + 4, + 4, + false, + 1, + 1, + S<1, 64, 1, 2>, + S<8, 8>>; + +using DeviceOpInstance = DeviceOpInstanceKKNN; + +#include "run_batched_gemm_bias_e_permute_example.inc" +int main(int argc, char* argv[]) { return !run_batched_gemm_bias_e_permute_example(argc, argv); } diff --git a/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_xdl_fp16.cpp b/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_xdl_fp16.cpp index d7f468bc62..4ed054faaa 100644 --- a/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_xdl_fp16.cpp +++ b/example/29_batched_gemm_bias_e_permute/batched_gemm_bias_e_permute_xdl_fp16.cpp @@ -67,340 +67,5 @@ using DeviceOpInstanceKKNN = ck::tensor_operation::device:: using DeviceOpInstance = DeviceOpInstanceKKNN; -// hardcoded for NumDimM == NumDimN == NumDimK == 2 -template = - false> -struct ReferenceContraction_G2_M2_N2_K1 : public ck::tensor_operation::device::BaseOperator -{ - // Argument - struct Argument : public ck::tensor_operation::device::BaseArgument - { - Argument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - : a_gs_ms_ks_{a_gs_ms_ks}, - b_gs_ns_ks_{b_gs_ns_ks}, - e_gs_ms_ns_{e_gs_ms_ns}, - a_element_op_{a_element_op}, - b_element_op_{b_element_op}, - cde_element_op_{cde_element_op} - { - } - - const Tensor& a_gs_ms_ks_; - const Tensor& b_gs_ns_ks_; - Tensor& e_gs_ms_ns_; - - AElementwiseOperation a_element_op_; - BElementwiseOperation b_element_op_; - CDEElementwiseOperation cde_element_op_; - }; - - // Invoker - struct Invoker : public ck::tensor_operation::device::BaseInvoker - { - using Argument = ReferenceContraction_G2_M2_N2_K1::Argument; - - float Run(const Argument& arg) - { - auto f_ms_ns = [&](auto g0, auto g1, auto m0, auto m1, auto n0, auto n1) { - const int K0 = arg.a_gs_ms_ks_.mDesc.GetLengths()[4]; - - AccDataType v_acc = 0; - - for(int k0 = 0; k0 < K0; ++k0) - { - AccDataType v_a; - AccDataType v_b; - - arg.a_element_op_( - v_a, - ck::type_convert(arg.a_gs_ms_ks_(g0, g1, m0, m1, k0))); - arg.b_element_op_( - v_b, - ck::type_convert(arg.b_gs_ns_ks_(g0, g1, n0, n1, k0))); - - v_acc += v_a * v_b; - } - - AccDataType v_c; - - arg.cde_element_op_(v_c, v_acc); - - arg.e_gs_ms_ns_(g0, g1, m0, m1, n0, n1) = v_c; - }; - - make_ParallelTensorFunctor(f_ms_ns, - arg.e_gs_ms_ns_.mDesc.GetLengths()[0], - arg.e_gs_ms_ns_.mDesc.GetLengths()[1], - arg.e_gs_ms_ns_.mDesc.GetLengths()[2], - arg.e_gs_ms_ns_.mDesc.GetLengths()[3], - arg.e_gs_ms_ns_.mDesc.GetLengths()[4], - arg.e_gs_ms_ns_.mDesc.GetLengths()[5])( - std::thread::hardware_concurrency()); - - return 0; - } - - float Run(const ck::tensor_operation::device::BaseArgument* p_arg, - const StreamConfig& /* stream_config */ = StreamConfig{}) override - { - return Run(*dynamic_cast(p_arg)); - } - }; - - static constexpr bool IsValidCompilationParameter() - { - // TODO: properly implement this check - return true; - } - - bool IsSupportedArgument(const ck::tensor_operation::device::BaseArgument*) override - { - return true; - } - - static auto MakeArgument(const Tensor& a_gs_ms_ks, - const Tensor& b_gs_ns_ks, - Tensor& e_gs_ms_ns, - AElementwiseOperation a_element_op, - BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) - { - return Argument{ - a_gs_ms_ks, b_gs_ns_ks, e_gs_ms_ns, a_element_op, b_element_op, cde_element_op}; - } - - static auto MakeInvoker() { return Invoker{}; } - - virtual std::unique_ptr MakeInvokerPointer() - { - return std::make_unique(Invoker{}); - } - - std::string GetTypeString() const override - { - auto str = std::stringstream(); - - // clang-format off - str << "ReferenceContraction_G2_M2_N2_K1" - << std::endl; - // clang-format on - - return str.str(); - } -}; - -int main(int argc, char* argv[]) -{ - bool do_verification = true; - int init_method = 1; - bool time_kernel = false; - - ck::index_t G0 = 1; - ck::index_t G1 = 2; - - ck::index_t M0 = 4; - ck::index_t M1 = 256; - - ck::index_t N0 = 16; - ck::index_t N1 = 128; - - ck::index_t K0 = 64; - - // A[G0, G1, M0, M1, K0] - std::vector a_gs_ms_ks_lengths{G0, G1, M0, M1, K0}; - std::vector a_gs_ms_ks_strides{G1 * M0 * M1 * K0, M0 * M1 * K0, M1 * K0, K0, 1}; - // B[G0, G1, N0, N1, K0] - std::vector b_gs_ns_ks_lengths{G0, G1, N0, N1, K0}; - std::vector b_gs_ns_ks_strides{G1 * N0 * N1 * K0, N0 * N1 * K0, N1 * K0, K0, 1}; - - // D[G0, G1, M0, N0, M1, N1] - std::vector d_gs_ms_ns_lengths{G0, G1, M0, M1, N0, N1}; - std::vector d_gs_ms_ns_strides{G1 * N0 * N1, N0 * N1, 0, 0, N1, 1}; - // E[G0, G1, M0, N0, M1, N1] - std::vector e_gs_ms_ns_lengths{G0, G1, M0, M1, N0, N1}; - std::vector e_gs_ms_ns_strides{ - G1 * M0 * N0 * M1 * N1, M0 * N0 * M1 * N1, N0 * M1 * N1, N1, M1 * N1, 1}; - - if(argc == 1) - { - // use default case - } - else if(argc == 4) - { - do_verification = std::stoi(argv[1]); - init_method = std::stoi(argv[2]); - time_kernel = std::stoi(argv[3]); - } - else - { - printf("arg1: verification (0=no, 1=yes)\n"); - printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"); - printf("arg3: time kernel (0=no, 1=yes)\n"); - exit(0); - } - - Tensor a_gs_ms_ks(a_gs_ms_ks_lengths, a_gs_ms_ks_strides, Row{}); - Tensor b_gs_ns_ks(b_gs_ns_ks_lengths, b_gs_ns_ks_strides, Row{}); - Tensor d_gs_ms_ns(d_gs_ms_ns_lengths, d_gs_ms_ns_strides, Bypass{}); - Tensor e_gs_ms_ns_host_result(e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - Tensor e_gs_ms_ns_device_result(e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - - std::cout << "a_gs_ms_ks: " << a_gs_ms_ks.mDesc << std::endl; - std::cout << "b_gs_ns_ks: " << b_gs_ns_ks.mDesc << std::endl; - std::cout << "d_gs_ms_ns: " << d_gs_ms_ns.mDesc << std::endl; - std::cout << "e_gs_ms_ns: " << e_gs_ms_ns_host_result.mDesc << std::endl; - - switch(init_method) - { - case 0: break; - case 1: - a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_2{-5, 5}); - b_gs_ns_ks.GenerateTensorValue(GeneratorTensor_2{-5, 5}); - d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_2{-5, 5}); - break; - default: - a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - b_gs_ns_ks.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); - d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); - break; - } - - DeviceMem a_device_buf(sizeof(ADataType) * a_gs_ms_ks.mDesc.GetElementSpaceSize()); - DeviceMem b_device_buf(sizeof(BDataType) * b_gs_ns_ks.mDesc.GetElementSpaceSize()); - DeviceMem d_device_buf(sizeof(DDataType) * d_gs_ms_ns.mDesc.GetElementSpaceSize()); - DeviceMem e_device_buf(sizeof(EDataType) * - e_gs_ms_ns_device_result.mDesc.GetElementSpaceSize()); - - a_device_buf.ToDevice(a_gs_ms_ks.mData.data()); - b_device_buf.ToDevice(b_gs_ns_ks.mData.data()); - d_device_buf.ToDevice(d_gs_ms_ns.mData.data()); - - // set zero - e_device_buf.SetZero(); - - auto a_element_op = AElementOp{}; - auto b_element_op = BElementOp{}; - auto cde_element_op = CDEElementOp{}; - - // device operation - auto op = DeviceOpInstance{}; - auto invoker = op.MakeInvoker(); - auto argument = op.MakeArgument(a_device_buf.GetDeviceBuffer(), - b_device_buf.GetDeviceBuffer(), - std::array{d_device_buf.GetDeviceBuffer()}, - e_device_buf.GetDeviceBuffer(), - a_gs_ms_ks_lengths, - a_gs_ms_ks_strides, - b_gs_ns_ks_lengths, - b_gs_ns_ks_strides, - std::array, 1>{d_gs_ms_ns_lengths}, - std::array, 1>{d_gs_ms_ns_strides}, - e_gs_ms_ns_lengths, - e_gs_ms_ns_strides, - a_element_op, - b_element_op, - cde_element_op); - - if(!op.IsSupportedArgument(argument)) - { - std::cout << op.GetTypeString() << " does not support this problem" << std::endl; - - return 0; - } - - float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel}); - - ck::index_t G = - ck::accumulate_n(e_gs_ms_ns_lengths.begin(), NumDimG, 1, std::multiplies<>{}); - - ck::index_t M = ck::accumulate_n( - e_gs_ms_ns_lengths.begin() + NumDimG, NumDimM, 1, std::multiplies<>{}); - - ck::index_t N = ck::accumulate_n( - e_gs_ms_ns_lengths.begin() + NumDimG + NumDimM, NumDimN, 1, std::multiplies<>{}); - - ck::index_t K = ck::accumulate_n( - a_gs_ms_ks_lengths.begin() + NumDimG + NumDimM, NumDimK, 1, std::multiplies<>{}); - - std::size_t flop = std::size_t(2) * G * M * N * K; - std::size_t num_btype = sizeof(ADataType) * G * M * K + sizeof(BDataType) * G * K * N + - sizeof(DDataType) * G * M * N + sizeof(EDataType) * G * M * N; - - float tflops = static_cast(flop) / 1.E9 / ave_time; - - float gb_per_sec = num_btype / 1.E6 / ave_time; - - std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " - << op.GetTypeString() << std::endl; - - e_device_buf.FromDevice(e_gs_ms_ns_device_result.mData.data()); - - if(do_verification) - { - Tensor c_ms_ns_host_result( - e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); - - using ReferenceOpInstance = ReferenceContraction_G2_M2_N2_K1; - - auto ref_gemm = ReferenceOpInstance{}; - auto ref_invoker = ref_gemm.MakeInvoker(); - - auto ref_argument = ref_gemm.MakeArgument( - a_gs_ms_ks, b_gs_ns_ks, c_ms_ns_host_result, a_element_op, b_element_op, PassThrough{}); - - ref_invoker.Run(ref_argument); - - for(size_t g0 = 0; g0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[0]; ++g0) - { - for(size_t g1 = 0; g1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[1]; ++g1) - { - for(size_t m0 = 0; m0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[2]; ++m0) - { - for(size_t m1 = 0; m1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[3]; ++m1) - { - for(size_t n0 = 0; n0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[4]; ++n0) - { - for(size_t n1 = 0; n1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[5]; - ++n1) - { - cde_element_op(e_gs_ms_ns_host_result(g0, g1, m0, m1, n0, n1), - c_ms_ns_host_result(g0, g1, m0, m1, n0, n1), - d_gs_ms_ns(g0, g1, m0, m1, n0, n1)); - } - } - } - } - } - } - - return ck::utils::check_err(e_gs_ms_ns_device_result, e_gs_ms_ns_host_result) ? 0 : 1; - } - - return 0; -} +#include "run_batched_gemm_bias_e_permute_example.inc" +int main(int argc, char* argv[]) { return !run_batched_gemm_bias_e_permute_example(argc, argv); } diff --git a/example/29_batched_gemm_bias_e_permute/run_batched_gemm_bias_e_permute_example.inc b/example/29_batched_gemm_bias_e_permute/run_batched_gemm_bias_e_permute_example.inc new file mode 100644 index 0000000000..803c1eb0bf --- /dev/null +++ b/example/29_batched_gemm_bias_e_permute/run_batched_gemm_bias_e_permute_example.inc @@ -0,0 +1,350 @@ + +// hardcoded for NumDimM == NumDimN == NumDimK == 2 +template = + false> +struct ReferenceContraction_G2_M2_N2_K1 : public ck::tensor_operation::device::BaseOperator +{ + // Argument + struct Argument : public ck::tensor_operation::device::BaseArgument + { + Argument(const Tensor& a_gs_ms_ks, + const Tensor& b_gs_ns_ks, + Tensor& e_gs_ms_ns, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) + : a_gs_ms_ks_{a_gs_ms_ks}, + b_gs_ns_ks_{b_gs_ns_ks}, + e_gs_ms_ns_{e_gs_ms_ns}, + a_element_op_{a_element_op}, + b_element_op_{b_element_op}, + cde_element_op_{cde_element_op} + { + } + + const Tensor& a_gs_ms_ks_; + const Tensor& b_gs_ns_ks_; + Tensor& e_gs_ms_ns_; + + AElementwiseOperation a_element_op_; + BElementwiseOperation b_element_op_; + CDEElementwiseOperation cde_element_op_; + }; + + // Invoker + struct Invoker : public ck::tensor_operation::device::BaseInvoker + { + using Argument = ReferenceContraction_G2_M2_N2_K1::Argument; + + float Run(const Argument& arg) + { + auto f_ms_ns = [&](auto g0, auto g1, auto m0, auto m1, auto n0, auto n1) { + const int K0 = arg.a_gs_ms_ks_.mDesc.GetLengths()[4]; + + AccDataType v_acc = 0; + + for(int k0 = 0; k0 < K0; ++k0) + { + AccDataType v_a; + AccDataType v_b; + + arg.a_element_op_( + v_a, + ck::type_convert(arg.a_gs_ms_ks_(g0, g1, m0, m1, k0))); + arg.b_element_op_( + v_b, + ck::type_convert(arg.b_gs_ns_ks_(g0, g1, n0, n1, k0))); + + v_acc += v_a * v_b; + } + + AccDataType v_c; + + arg.cde_element_op_(v_c, v_acc); + + arg.e_gs_ms_ns_(g0, g1, m0, m1, n0, n1) = v_c; + }; + + make_ParallelTensorFunctor(f_ms_ns, + arg.e_gs_ms_ns_.mDesc.GetLengths()[0], + arg.e_gs_ms_ns_.mDesc.GetLengths()[1], + arg.e_gs_ms_ns_.mDesc.GetLengths()[2], + arg.e_gs_ms_ns_.mDesc.GetLengths()[3], + arg.e_gs_ms_ns_.mDesc.GetLengths()[4], + arg.e_gs_ms_ns_.mDesc.GetLengths()[5])( + std::thread::hardware_concurrency()); + + return 0; + } + + float Run(const ck::tensor_operation::device::BaseArgument* p_arg, + const StreamConfig& /* stream_config */ = StreamConfig{}) override + { + return Run(*dynamic_cast(p_arg)); + } + }; + + static constexpr bool IsValidCompilationParameter() + { + // TODO: properly implement this check + return true; + } + + bool IsSupportedArgument(const ck::tensor_operation::device::BaseArgument*) override + { + return true; + } + + static auto MakeArgument(const Tensor& a_gs_ms_ks, + const Tensor& b_gs_ns_ks, + Tensor& e_gs_ms_ns, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) + { + return Argument{ + a_gs_ms_ks, b_gs_ns_ks, e_gs_ms_ns, a_element_op, b_element_op, cde_element_op}; + } + + static auto MakeInvoker() { return Invoker{}; } + + virtual std::unique_ptr MakeInvokerPointer() + { + return std::make_unique(Invoker{}); + } + + std::string GetTypeString() const override + { + auto str = std::stringstream(); + + // clang-format off + str << "ReferenceContraction_G2_M2_N2_K1" + << std::endl; + // clang-format on + + return str.str(); + } +}; + +int run_batched_gemm_bias_e_permute_example(int argc, char* argv[]) +{ + bool do_verification = true; + int init_method = 1; + bool time_kernel = false; + + ck::index_t G0 = 1; + ck::index_t G1 = 2; + + ck::index_t M0 = 4; + ck::index_t M1 = 128; + + ck::index_t N0 = 16; + ck::index_t N1 = 256; + + ck::index_t K0 = 2048; + + if(argc == 1) + { + // use default case + } + else if(argc == 4) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + } + else if(argc == 11) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + G0 = std::stoi(argv[4]); + G1 = std::stoi(argv[5]); + M0 = std::stoi(argv[6]); + M1 = std::stoi(argv[7]); + N0 = std::stoi(argv[8]); + N1 = std::stoi(argv[9]); + K0 = std::stoi(argv[10]); + } + else + { + printf("arg1: verification (0=no, 1=yes)\n"); + printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"); + printf("arg3: time kernel (0=no, 1=yes)\n"); + printf("arg4-10: G0, G1, M0, M1, N0, N1, K0\n"); + exit(0); + } + + // A[G0, G1, M0, M1, K0] + std::vector a_gs_ms_ks_lengths{G0, G1, M0, M1, K0}; + std::vector a_gs_ms_ks_strides{G1 * M0 * M1 * K0, M0 * M1 * K0, M1 * K0, K0, 1}; + // B[G0, G1, N0, N1, K0] + std::vector b_gs_ns_ks_lengths{G0, G1, N0, N1, K0}; + std::vector b_gs_ns_ks_strides{G1 * N0 * N1 * K0, N0 * N1 * K0, N1 * K0, K0, 1}; + + // D[G0, G1, M0, N0, M1, N1] + std::vector d_gs_ms_ns_lengths{G0, G1, M0, M1, N0, N1}; + std::vector d_gs_ms_ns_strides{G1 * N0 * N1, N0 * N1, 0, 0, N1, 1}; + // E[G0, G1, M0, N0, M1, N1] + std::vector e_gs_ms_ns_lengths{G0, G1, M0, M1, N0, N1}; + std::vector e_gs_ms_ns_strides{ + G1 * M0 * N0 * M1 * N1, M0 * N0 * M1 * N1, N0 * M1 * N1, N1, M1 * N1, 1}; + + Tensor a_gs_ms_ks(a_gs_ms_ks_lengths, a_gs_ms_ks_strides, Row{}); + Tensor b_gs_ns_ks(b_gs_ns_ks_lengths, b_gs_ns_ks_strides, Row{}); + Tensor d_gs_ms_ns(d_gs_ms_ns_lengths, d_gs_ms_ns_strides, Bypass{}); + Tensor e_gs_ms_ns_host_result(e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); + Tensor e_gs_ms_ns_device_result(e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); + std::cout << "a_gs_ms_ks: " << a_gs_ms_ks.mDesc << std::endl; + std::cout << "b_gs_ns_ks: " << b_gs_ns_ks.mDesc << std::endl; + std::cout << "d_gs_ms_ns: " << d_gs_ms_ns.mDesc << std::endl; + std::cout << "e_gs_ms_ns: " << e_gs_ms_ns_host_result.mDesc << std::endl; + + switch(init_method) + { + case 0: break; + case 1: + a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_2{-5, 5}); + b_gs_ns_ks.GenerateTensorValue(GeneratorTensor_2{-5, 5}); + d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_2{-5, 5}); + break; + default: + a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + b_gs_ns_ks.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + break; + } + DeviceMem a_device_buf(sizeof(ADataType) * a_gs_ms_ks.mDesc.GetElementSpaceSize()); + DeviceMem b_device_buf(sizeof(BDataType) * b_gs_ns_ks.mDesc.GetElementSpaceSize()); + DeviceMem d_device_buf(sizeof(DDataType) * d_gs_ms_ns.mDesc.GetElementSpaceSize()); + DeviceMem e_device_buf(sizeof(EDataType) * + e_gs_ms_ns_device_result.mDesc.GetElementSpaceSize()); + + a_device_buf.ToDevice(a_gs_ms_ks.mData.data()); + b_device_buf.ToDevice(b_gs_ns_ks.mData.data()); + d_device_buf.ToDevice(d_gs_ms_ns.mData.data()); + + // set zero + e_device_buf.SetZero(); + + auto a_element_op = AElementOp{}; + auto b_element_op = BElementOp{}; + auto cde_element_op = CDEElementOp{}; + + // device operation + auto op = DeviceOpInstance{}; + auto invoker = op.MakeInvoker(); + auto argument = op.MakeArgument(a_device_buf.GetDeviceBuffer(), + b_device_buf.GetDeviceBuffer(), + std::array{d_device_buf.GetDeviceBuffer()}, + e_device_buf.GetDeviceBuffer(), + a_gs_ms_ks_lengths, + a_gs_ms_ks_strides, + b_gs_ns_ks_lengths, + b_gs_ns_ks_strides, + std::array, 1>{d_gs_ms_ns_lengths}, + std::array, 1>{d_gs_ms_ns_strides}, + e_gs_ms_ns_lengths, + e_gs_ms_ns_strides, + a_element_op, + b_element_op, + cde_element_op); + + if(!op.IsSupportedArgument(argument)) + { + std::cout << op.GetTypeString() << " does not support this problem" << std::endl; + + return 0; + } + + float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel}); + + ck::index_t G = + ck::accumulate_n(e_gs_ms_ns_lengths.begin(), NumDimG, 1, std::multiplies<>{}); + + ck::index_t M = ck::accumulate_n( + e_gs_ms_ns_lengths.begin() + NumDimG, NumDimM, 1, std::multiplies<>{}); + + ck::index_t N = ck::accumulate_n( + e_gs_ms_ns_lengths.begin() + NumDimG + NumDimM, NumDimN, 1, std::multiplies<>{}); + + ck::index_t K = ck::accumulate_n( + a_gs_ms_ks_lengths.begin() + NumDimG + NumDimM, NumDimK, 1, std::multiplies<>{}); + std::cout << "GMNK=" << G << ", " << M << ", " << N << ", " << K << std::endl; + std::size_t flop = std::size_t(2) * G * M * N * K; + std::size_t num_btype = sizeof(ADataType) * G * M * K + sizeof(BDataType) * G * K * N + + sizeof(DDataType) * G * M * N + sizeof(EDataType) * G * M * N; + + float tflops = static_cast(flop) / 1.E9 / ave_time; + + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " + << op.GetTypeString() << std::endl; + + e_device_buf.FromDevice(e_gs_ms_ns_device_result.mData.data()); + + if(do_verification) + { + Tensor c_ms_ns_host_result( + e_gs_ms_ns_lengths, e_gs_ms_ns_strides, Bypass{}); + + using ReferenceOpInstance = ReferenceContraction_G2_M2_N2_K1; + + auto ref_gemm = ReferenceOpInstance{}; + auto ref_invoker = ref_gemm.MakeInvoker(); + + auto ref_argument = ref_gemm.MakeArgument( + a_gs_ms_ks, b_gs_ns_ks, c_ms_ns_host_result, a_element_op, b_element_op, PassThrough{}); + + ref_invoker.Run(ref_argument); + + for(size_t g0 = 0; g0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[0]; ++g0) + { + for(size_t g1 = 0; g1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[1]; ++g1) + { + for(size_t m0 = 0; m0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[2]; ++m0) + { + for(size_t m1 = 0; m1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[3]; ++m1) + { + for(size_t n0 = 0; n0 < e_gs_ms_ns_host_result.mDesc.GetLengths()[4]; ++n0) + { + for(size_t n1 = 0; n1 < e_gs_ms_ns_host_result.mDesc.GetLengths()[5]; + ++n1) + { + cde_element_op(e_gs_ms_ns_host_result(g0, g1, m0, m1, n0, n1), + c_ms_ns_host_result(g0, g1, m0, m1, n0, n1), + d_gs_ms_ns(g0, g1, m0, m1, n0, n1)); + } + } + } + } + } + } + + return ck::utils::check_err(e_gs_ms_ns_device_result, e_gs_ms_ns_host_result); + } + + return 1; +} diff --git a/example/37_batched_gemm_add_add_relu_gemm_add/CMakeLists.txt b/example/37_batched_gemm_add_add_relu_gemm_add/CMakeLists.txt index 7b6eb01413..769b0888df 100644 --- a/example/37_batched_gemm_add_add_relu_gemm_add/CMakeLists.txt +++ b/example/37_batched_gemm_add_add_relu_gemm_add/CMakeLists.txt @@ -2,3 +2,4 @@ # SPDX-License-Identifier: MIT add_example_executable(example_batched_gemm_add_add_relu_gemm_add_xdl_fp16 batched_gemm_add_add_relu_gemm_add_xdl_fp16.cpp) +add_example_executable(example_batched_gemm_add_add_relu_gemm_add_wmma_fp16 batched_gemm_add_add_relu_gemm_add_wmma_fp16.cpp) diff --git a/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_add_add_relu_gemm_add_wmma_fp16.cpp b/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_add_add_relu_gemm_add_wmma_fp16.cpp new file mode 100644 index 0000000000..bf12ee6c3e --- /dev/null +++ b/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_add_add_relu_gemm_add_wmma_fp16.cpp @@ -0,0 +1,135 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/* +Computes C_m_o = Relu(A0[m, k] * B0[n, k] + D00[m, n] + D01[mn]) * B1[n, o] + D1[m, o] +*/ + +#include +#include +#include +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_gemm_multiple_d_wmma_cshuffle_v3.hpp" +#include "ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp" + +#include "ck/library/utility/check_err.hpp" +#include "ck/library/utility/device_memory.hpp" +#include "ck/library/utility/host_tensor.hpp" +#include "ck/library/utility/host_tensor_generator.hpp" +#include "ck/library/utility/literals.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp" + +#include "element_ops.h" + +using ::ck::DeviceMem; +using ::ck::HostTensorDescriptor; +using ::ck::Tensor; + +template +using S = ck::Sequence; + +using F16 = ck::half_t; +using F32 = float; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; + +using A0DataType = F16; +using B0DataType = F16; +using D00DataType = F16; +using D01DataType = F16; +using B1DataType = F16; +using D1DataType = F16; +using AccDataType = F32; +using CShuffleDataType = F32; +using E1DataType = F16; + +using A0Layout = Row; +using B0Layout = Col; +using D00Layout = Row; +using D01Layout = Row; +using B1Layout = Row; +using D1Layout = Row; +using E1Layout = Row; + +using A0ElementOp = PassThrough; +using B0ElementOp = PassThrough; +using CDE0ElementOp = AddAddRelu; +using A1ElementOp = PassThrough; +using B1ElementOp = PassThrough; +using CDE1ElementOp = ck::tensor_operation::element_wise::Add; + +constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; + +using DeviceGemmInstance = + ck::tensor_operation::device::DeviceBatchedGemmMultipleDGemmMultipleD_Wmma_CShuffleV3< + A0Layout, + B0Layout, + ck::Tuple, + B1Layout, + ck::Tuple, + E1Layout, + A0DataType, + B0DataType, + ck::Tuple, + B1DataType, + ck::Tuple, + E1DataType, + AccDataType, + CShuffleDataType, + A0ElementOp, + B0ElementOp, + CDE0ElementOp, + B1ElementOp, + CDE1ElementOp, + GemmSpec, + + 32, // BlockSize + 16, // MPerBlock + 64, // LPerBlock + 64, // KPerBlock + 64, // NPerBlock (Gemm1NPerBlock) + 64, // LTilePerBlock (Gemm1KPerBlock) + 8, // AK1 + 8, // BK1 + 8, // L1 (B1K1) + 16, // MPerWmma + 16, // LPerWmma + 1, // MRepeat + 4, // LRepeat (Gemm0NRepeat) + 4, // NRepeat (Gemm1NRepeat) + + S<2, 16, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder + S<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + false, // ABlockLdsAddExtraM + S<2, 16, 1>, // B0BlockTransferThreadClusterLengths_K0_L_K1 + S<1, 0, 2>, // B0BlockTransferThreadClusterArrangeOrder + S<1, 0, 2>, // B0BlockTransferSrcAccessOrder + 2, // B0BlockTransferSrcVectorDim + 8, // B0BlockTransferSrcScalarPerVector + 8, // B0BlockTransferDstScalarPerVector_K1 + false, // B0BlockLdsAddExtraL + 4, // CDE0BlockTransferSrcScalarPerVector + S<2, 16, 1>, // B1BlockTransferThreadClusterLengths_L0_N_L1 + S<0, 2, 1>, // B1BlockTransferThreadClusterArrangeOrder + S<0, 2, 1>, // B1BlockTransferSrcAccessOrder + 1, // B1BlockTransferSrcVectorDim + 4, // B1BlockTransferSrcScalarPerVector + 2, // B1BlockTransferDstScalarPerVector_L1 + true, // B1BlockLdsAddExtraN + 1, // CShuffleMRepeatPerShuffle + 2, // CShuffleNRepeatPerShuffle + S<1, 16, 1, 2>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock + 8>; // CShuffleBlockTransferScalarPerVector_NPerBlock + +#include "batched_gemm_multiple_d_gemm_multiple_d.inc" +int main(int argc, char* argv[]) { return run_example(argc, argv); } diff --git a/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_add_add_relu_gemm_add_xdl_fp16.cpp b/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_add_add_relu_gemm_add_xdl_fp16.cpp index ab87124c6b..60a1316907 100644 --- a/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_add_add_relu_gemm_add_xdl_fp16.cpp +++ b/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_add_add_relu_gemm_add_xdl_fp16.cpp @@ -22,6 +22,8 @@ Computes C_m_o = Relu(A0[m, k] * B0[n, k] + D00[m, n] + D01[mn]) * B1[n, o] + D1 #include "ck/library/utility/literals.hpp" #include "ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp" +#include "element_ops.h" + using ::ck::DeviceMem; using ::ck::HostTensorDescriptor; using ::ck::Tensor; @@ -39,11 +41,10 @@ using PassThrough = ck::tensor_operation::element_wise::PassThrough; using A0DataType = F16; using B0DataType = F16; -using Acc0DataType = F32; +using AccDataType = F32; using D00DataType = F16; using D01DataType = F16; using B1DataType = F16; -using Acc1DataType = F32; using C1ShuffleDataType = F32; using D1DataType = F16; using E1DataType = F16; @@ -56,58 +57,6 @@ using B1Layout = Row; using D1Layout = Row; using E1Layout = Row; -// E = Relu(C + D0 + D1) -struct AddAddRelu -{ - __host__ __device__ void - operator()(ck::half_t& e, const ck::half_t& c, const ck::half_t& d0, const ck::half_t& d1) const - { - const ck::half_t x = c + d0 + d1; - - ck::tensor_operation::element_wise::Relu{}.operator()(e, x); - } - __host__ __device__ void - operator()(float& e, const float& c, const ck::half_t& d0, const ck::half_t& d1) const - { - const float x = c + (d0 + d1); - - ck::tensor_operation::element_wise::Relu{}.operator()(e, x); - } -}; - -// E = Gelu(C + D0 + D1) -struct AddAddGelu -{ - __host__ __device__ void - operator()(ck::half_t& e, const ck::half_t& c, const ck::half_t& d0, const ck::half_t& d1) const - { - const ck::half_t x = c + d0 + d1; - - ck::tensor_operation::element_wise::Gelu{}.template operator()(e, - x); - } - - __host__ __device__ void - operator()(float& e, const float& c, const ck::half_t& d0, const ck::half_t& d1) const - { - const float x = c + (d0 + d1); - - ck::tensor_operation::element_wise::Gelu{}.template operator()(e, x); - } -}; - -// E = FastGelu(C + D0 + D1) -struct AddAddFastGelu -{ - __host__ __device__ void - operator()(float& e, const float& c, const ck::half_t& d0, const ck::half_t& d1) const - { - const float x = c + (d0 + d1); - - ck::tensor_operation::element_wise::FastGelu{}.template operator()(e, x); - } -}; - using A0ElementOp = PassThrough; using B0ElementOp = PassThrough; using CDE0ElementOp = AddAddRelu; @@ -131,10 +80,10 @@ using DeviceGemmInstance = E1Layout, A0DataType, B0DataType, - Acc0DataType, + AccDataType, ck::Tuple, B1DataType, - Acc1DataType, + AccDataType, C1ShuffleDataType, ck::Tuple, E1DataType, @@ -191,337 +140,5 @@ using DeviceGemmInstance = S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock 4>; // CShuffleBlockTransferScalarPerVector_NPerBlock -int main(int argc, char* argv[]) -{ - bool do_verification = true; - int init_method = 1; - bool time_kernel = false; - - // GEMM shape - ck::index_t M = 1024; - ck::index_t N = 1024; - ck::index_t K = 64; - ck::index_t O = 128; - ck::index_t BatchCount = 4; - ck::index_t StrideA0 = -1; - ck::index_t StrideB0 = -1; - ck::index_t StrideD00 = -1; - ck::index_t StrideD01 = -1; - ck::index_t StrideB1 = -1; - ck::index_t StrideD1 = -1; - ck::index_t StrideE1 = -1; - ck::index_t BatchStrideA0 = -1; - ck::index_t BatchStrideB0 = -1; - ck::index_t BatchStrideD00 = -1; - ck::index_t BatchStrideD01 = -1; - ck::index_t BatchStrideB1 = -1; - ck::index_t BatchStrideD1 = -1; - ck::index_t BatchStrideE1 = -1; - - if(argc == 1) - { - // use default case - } - else if(argc == 4) - { - do_verification = std::stoi(argv[1]); - init_method = std::stoi(argv[2]); - time_kernel = std::stoi(argv[3]); - } - else if(argc == 9) - { - do_verification = std::stoi(argv[1]); - init_method = std::stoi(argv[2]); - time_kernel = std::stoi(argv[3]); - - M = std::stoi(argv[4]); - N = std::stoi(argv[5]); - K = std::stoi(argv[6]); - O = std::stoi(argv[7]); - - BatchCount = std::stoi(argv[8]); - } - else if(argc == 23) - { - do_verification = std::stoi(argv[1]); - init_method = std::stoi(argv[2]); - time_kernel = std::stoi(argv[3]); - - M = std::stoi(argv[4]); - N = std::stoi(argv[5]); - K = std::stoi(argv[6]); - O = std::stoi(argv[7]); - - BatchCount = std::stoi(argv[8]); - - StrideA0 = std::stoi(argv[9]); - StrideB0 = std::stoi(argv[10]); - StrideD00 = std::stoi(argv[11]); - StrideD01 = std::stoi(argv[12]); - StrideB1 = std::stoi(argv[13]); - StrideD1 = std::stoi(argv[14]); - StrideE1 = std::stoi(argv[15]); - - BatchStrideA0 = std::stoi(argv[16]); - BatchStrideB0 = std::stoi(argv[17]); - BatchStrideD00 = std::stoi(argv[18]); - BatchStrideD01 = std::stoi(argv[19]); - BatchStrideB1 = std::stoi(argv[20]); - BatchStrideD1 = std::stoi(argv[21]); - BatchStrideE1 = std::stoi(argv[22]); - } - else - { - printf("arg1: verification (0=no, 1=yes)\n"); - printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"); - printf("arg3: time kernel (0=no, 1=yes)\n"); - printf("arg4 to 8: M, N, K, O, Batch\n"); - printf( - "arg9 to 15: StrideA0, StrideB0, StrideD00, StrideD01, StrideB1, StrideD1, StrideE1\n"); - printf("arg16 to 22: BatchStrideA0, BatchStrideB0, BatchStrideD00, BatchStrideD01, " - "BatchStrideB1, BatchStrideD1, BatchStrideE1 \n"); - exit(0); - } - - const int DefaultStrideA0 = ck::is_same_v ? K : M; - const int DefaultStrideB0 = ck::is_same_v ? N : K; - const int DefaultStrideD00 = ck::is_same_v ? N : M; - const int DefaultStrideD01 = ck::is_same_v ? N : M; - const int DefaultStrideB1 = ck::is_same_v ? O : N; - const int DefaultStrideD1 = ck::is_same_v ? O : M; - const int DefaultStrideE1 = ck::is_same_v ? O : M; - - StrideA0 = (StrideA0 < 0) ? DefaultStrideA0 : StrideA0; - StrideB0 = (StrideB0 < 0) ? DefaultStrideB0 : StrideB0; - StrideD00 = (StrideD00 < 0) ? DefaultStrideD00 : StrideD00; - StrideD01 = (StrideD01 < 0) ? DefaultStrideD01 : StrideD01; - StrideB1 = (StrideB1 < 0) ? DefaultStrideB1 : StrideB1; - StrideD1 = (StrideD1 < 0) ? DefaultStrideD1 : StrideD1; - StrideE1 = (StrideE1 < 0) ? DefaultStrideE1 : StrideE1; - - const int DefaultBatchStrideA0 = (ck::is_same_v ? K : M) * StrideA0; - const int DefaultBatchStrideB0 = (ck::is_same_v ? N : K) * StrideB0; - const int DefaultBatchStrideD00 = (ck::is_same_v ? N : M) * StrideD00; - const int DefaultBatchStrideD01 = (ck::is_same_v ? N : M) * StrideD01; - const int DefaultBatchStrideB1 = (ck::is_same_v ? O : N) * StrideB1; - const int DefaultBatchStrideD1 = (ck::is_same_v ? O : M) * StrideD1; - const int DefaultBatchStrideE1 = (ck::is_same_v ? O : M) * StrideE1; - - BatchStrideA0 = BatchStrideA0 < 0 ? DefaultBatchStrideA0 : BatchStrideA0; - BatchStrideB0 = BatchStrideB0 < 0 ? DefaultBatchStrideB0 : BatchStrideB0; - BatchStrideD00 = BatchStrideD00 < 0 ? DefaultBatchStrideD00 : BatchStrideD00; - BatchStrideD01 = BatchStrideD01 < 0 ? DefaultBatchStrideD01 : BatchStrideD01; - BatchStrideB1 = BatchStrideB1 < 0 ? DefaultBatchStrideB1 : BatchStrideB1; - BatchStrideD1 = BatchStrideD1 < 0 ? DefaultBatchStrideD1 : BatchStrideD1; - BatchStrideE1 = BatchStrideE1 < 0 ? DefaultBatchStrideE1 : BatchStrideE1; - - auto f_host_tensor_descriptor = [](std::size_t batch_count, - std::size_t row, - std::size_t col, - std::size_t stride, - std::size_t batch_stride, - auto layout) { - using namespace ck::literals; - - if(std::is_same::value) - { - return HostTensorDescriptor( - {batch_count, row, col}, {batch_stride, stride, 1_uz}, layout); - } - else - { - return HostTensorDescriptor( - {batch_count, row, col}, {batch_stride, 1_uz, stride}, layout); - } - }; - - // E_m_o = A_m_k * B0_k_n * B1_n_o - Tensor a0_g_m_k( - f_host_tensor_descriptor(BatchCount, M, K, StrideA0, BatchStrideA0, A0Layout{})); - Tensor b0_g_k_n( - f_host_tensor_descriptor(BatchCount, K, N, StrideB0, BatchStrideB0, B0Layout{})); - Tensor d00_g_m_n( - f_host_tensor_descriptor(BatchCount, M, N, StrideD00, BatchStrideD00, D00Layout{})); - Tensor d01_g_m_n( - f_host_tensor_descriptor(BatchCount, M, N, StrideD01, BatchStrideD01, D01Layout{})); - Tensor b1_g_n_o( - f_host_tensor_descriptor(BatchCount, N, O, StrideB1, BatchStrideB1, B1Layout{})); - Tensor d1_g_m_o( - f_host_tensor_descriptor(BatchCount, M, O, StrideD1, BatchStrideD1, D1Layout{})); - Tensor e1_g_m_o_host_result( - f_host_tensor_descriptor(BatchCount, M, O, StrideE1, BatchStrideE1, E1Layout{})); - Tensor e1_g_m_o_device_result( - f_host_tensor_descriptor(BatchCount, M, O, StrideE1, BatchStrideE1, E1Layout{})); - - std::cout << "a0_g_m_k: " << a0_g_m_k.mDesc << std::endl; - std::cout << "b0_g_k_n: " << b0_g_k_n.mDesc << std::endl; - std::cout << "d00_g_m_n: " << d00_g_m_n.mDesc - << " size: " << d00_g_m_n.mDesc.GetElementSpaceSize() << std::endl; - std::cout << "d01_g_m_n: " << d01_g_m_n.mDesc - << " size: " << d01_g_m_n.mDesc.GetElementSpaceSize() << std::endl; - std::cout << "b1_g_n_o: " << b1_g_n_o.mDesc << std::endl; - std::cout << "e1_g_m_o: " << e1_g_m_o_host_result.mDesc << std::endl; - - switch(init_method) - { - case 0: break; - case 1: - a0_g_m_k.GenerateTensorValue(GeneratorTensor_2{-2, 3}); - b0_g_k_n.GenerateTensorValue(GeneratorTensor_2{-2, 3}); - d00_g_m_n.GenerateTensorValue(GeneratorTensor_2{-2, 3}); - d01_g_m_n.GenerateTensorValue(GeneratorTensor_2{-2, 3}); - b1_g_n_o.GenerateTensorValue(GeneratorTensor_2{-2, 3}); - d1_g_m_o.GenerateTensorValue(GeneratorTensor_2{-2, 3}); - break; - case 2: - a0_g_m_k.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - b0_g_k_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - d00_g_m_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - d01_g_m_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - b1_g_n_o.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); - d1_g_m_o.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - break; - default: - a0_g_m_k.GenerateTensorValue(GeneratorTensor_1{1}); - b0_g_k_n.GenerateTensorValue(GeneratorTensor_Sequential{}); - d00_g_m_n.GenerateTensorValue(GeneratorTensor_1{1}); - d01_g_m_n.GenerateTensorValue(GeneratorTensor_1{1}); - b1_g_n_o.GenerateTensorValue(GeneratorTensor_Diagonal{}); - d1_g_m_o.GenerateTensorValue(GeneratorTensor_1{1}); - } - - DeviceMem a0_g_m_k_device_buf(sizeof(A0DataType) * a0_g_m_k.mDesc.GetElementSize()); - DeviceMem b0_g_k_n_device_buf(sizeof(B0DataType) * b0_g_k_n.mDesc.GetElementSize()); - DeviceMem d00_g_m_n_device_buf(sizeof(D00DataType) * d00_g_m_n.mDesc.GetElementSpaceSize()); - DeviceMem d01_g_m_n_device_buf(sizeof(D01DataType) * d01_g_m_n.mDesc.GetElementSpaceSize()); - DeviceMem b1_g_n_o_device_buf(sizeof(B1DataType) * b1_g_n_o.mDesc.GetElementSize()); - DeviceMem e1_g_m_o_device_buf(sizeof(E1DataType) * - e1_g_m_o_device_result.mDesc.GetElementSize()); - DeviceMem d1_g_m_o_device_buf(sizeof(D1DataType) * d1_g_m_o.mDesc.GetElementSpaceSize()); - - a0_g_m_k_device_buf.ToDevice(a0_g_m_k.mData.data()); - b0_g_k_n_device_buf.ToDevice(b0_g_k_n.mData.data()); - d00_g_m_n_device_buf.ToDevice(d00_g_m_n.mData.data()); - d01_g_m_n_device_buf.ToDevice(d01_g_m_n.mData.data()); - b1_g_n_o_device_buf.ToDevice(b1_g_n_o.mData.data()); - d1_g_m_o_device_buf.ToDevice(d1_g_m_o.mData.data()); - - auto a0_element_op = A0ElementOp{}; - auto b0_element_op = B0ElementOp{}; - auto cde0_element_op = CDE0ElementOp{}; - auto b1_element_op = B1ElementOp{}; - auto cde1_element_op = CDE1ElementOp{}; - - // do GEMM - auto gemm = DeviceGemmInstance{}; - auto invoker = gemm.MakeInvoker(); - auto argument = - gemm.MakeArgument(static_cast(a0_g_m_k_device_buf.GetDeviceBuffer()), - static_cast(b0_g_k_n_device_buf.GetDeviceBuffer()), - std::array{d00_g_m_n_device_buf.GetDeviceBuffer(), - d01_g_m_n_device_buf.GetDeviceBuffer()}, - static_cast(b1_g_n_o_device_buf.GetDeviceBuffer()), - std::array{d1_g_m_o_device_buf.GetDeviceBuffer()}, - static_cast(e1_g_m_o_device_buf.GetDeviceBuffer()), - M, - N, - K, - O, - BatchCount, - StrideA0, - StrideB0, - std::array{StrideD00, StrideD01}, - StrideB1, - std::array{StrideD1}, - StrideE1, - BatchStrideA0, - BatchStrideB0, - std::array{BatchStrideD00, BatchStrideD01}, - BatchStrideB1, - std::array{BatchStrideD1}, - BatchStrideE1, - a0_element_op, - b0_element_op, - cde0_element_op, - b1_element_op, - cde1_element_op); - - if(!gemm.IsSupportedArgument(argument)) - { - std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl; - - return 0; - } - - float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel}); - - std::size_t flop = (size_t(M) * N * K * 2 + size_t(M) * N * O * 2) * BatchCount; - std::size_t num_btype = - (sizeof(A0DataType) * M * K + sizeof(B0DataType) * K * N + sizeof(D00DataType) * N + - sizeof(D01DataType) * N + sizeof(B1DataType) * N * O + sizeof(E1DataType) * M * O + - sizeof(D1DataType) * O) * - BatchCount; - - float tflops = static_cast(flop) / 1.E9 / ave_time; - - float gb_per_sec = num_btype / 1.E6 / ave_time; - - std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " - << gemm.GetTypeString() << std::endl; - - e1_g_m_o_device_buf.FromDevice(e1_g_m_o_device_result.mData.data()); - - if(do_verification) - { - using ReferenceGemm0Instance = - ck::tensor_operation::host::ReferenceBatchedGemm; - - using ReferenceGemm1Instance = - ck::tensor_operation::host::ReferenceBatchedGemm; - - // Output of Gemm0 is input A of Gemm1 - Tensor c0_g_m_n(f_host_tensor_descriptor(BatchCount, M, N, N, M * N, Row{})); - Tensor e0_g_m_n(f_host_tensor_descriptor(BatchCount, M, N, N, M * N, Row{})); - Tensor c1_g_m_o(f_host_tensor_descriptor(BatchCount, M, O, O, M * O, Row{})); - - auto ref_gemm0 = ReferenceGemm0Instance{}; - auto ref_gemm0_invoker = ref_gemm0.MakeInvoker(); - auto ref_gemm0_argument = ref_gemm0.MakeArgument( - a0_g_m_k, b0_g_k_n, c0_g_m_n, a0_element_op, b0_element_op, PassThrough{}); - - ref_gemm0_invoker.Run(ref_gemm0_argument); - - // bias+bias+relu - e0_g_m_n.ForEach([&](auto&, auto idx) { - cde0_element_op(e0_g_m_n(idx), c0_g_m_n(idx), d00_g_m_n(idx), d01_g_m_n(idx)); - }); - - auto ref_gemm1 = ReferenceGemm1Instance{}; - auto ref_gemm1_invoker = ref_gemm1.MakeInvoker(); - auto ref_gemm1_argument = ref_gemm1.MakeArgument( - e0_g_m_n, b1_g_n_o, c1_g_m_o, PassThrough{}, b1_element_op, PassThrough{}); - - ref_gemm1_invoker.Run(ref_gemm1_argument); - - // bias - e1_g_m_o_host_result.ForEach([&](auto&, auto idx) { - cde1_element_op(e1_g_m_o_host_result(idx), c1_g_m_o(idx), d1_g_m_o(idx)); - }); - - return ck::utils::check_err(e1_g_m_o_device_result, e1_g_m_o_host_result) ? 0 : 1; - } - - return 0; -} +#include "batched_gemm_multiple_d_gemm_multiple_d.inc" +int main(int argc, char* argv[]) { return run_example(argc, argv); } diff --git a/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_multiple_d_gemm_multiple_d.inc b/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_multiple_d_gemm_multiple_d.inc new file mode 100644 index 0000000000..57bdd5e9ef --- /dev/null +++ b/example/37_batched_gemm_add_add_relu_gemm_add/batched_gemm_multiple_d_gemm_multiple_d.inc @@ -0,0 +1,350 @@ + +int run_example(int argc, char* argv[]) +{ + bool do_verification = true; + int init_method = 1; + bool time_kernel = false; + + // GEMM shape + ck::index_t M = 1024; + ck::index_t N = 1024; + ck::index_t K = 256; + ck::index_t O = 512; + ck::index_t BatchCount = 4; + ck::index_t StrideA0 = -1; + ck::index_t StrideB0 = -1; + ck::index_t StrideD00 = -1; + ck::index_t StrideD01 = -1; + ck::index_t StrideB1 = -1; + ck::index_t StrideD1 = -1; + ck::index_t StrideE1 = -1; + ck::index_t BatchStrideA0 = -1; + ck::index_t BatchStrideB0 = -1; + ck::index_t BatchStrideD00 = -1; + ck::index_t BatchStrideD01 = -1; + ck::index_t BatchStrideB1 = -1; + ck::index_t BatchStrideD1 = -1; + ck::index_t BatchStrideE1 = -1; + + if(argc == 1) + { + // use default case + } + else if(argc == 4) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + } + else if(argc == 9) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + + M = std::stoi(argv[4]); + N = std::stoi(argv[5]); + K = std::stoi(argv[6]); + O = std::stoi(argv[7]); + + BatchCount = std::stoi(argv[8]); + } + else if(argc == 23) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + + M = std::stoi(argv[4]); + N = std::stoi(argv[5]); + K = std::stoi(argv[6]); + O = std::stoi(argv[7]); + + BatchCount = std::stoi(argv[8]); + + StrideA0 = std::stoi(argv[9]); + StrideB0 = std::stoi(argv[10]); + StrideD00 = std::stoi(argv[11]); + StrideD01 = std::stoi(argv[12]); + StrideB1 = std::stoi(argv[13]); + StrideD1 = std::stoi(argv[14]); + StrideE1 = std::stoi(argv[15]); + + BatchStrideA0 = std::stoi(argv[16]); + BatchStrideB0 = std::stoi(argv[17]); + BatchStrideD00 = std::stoi(argv[18]); + BatchStrideD01 = std::stoi(argv[19]); + BatchStrideB1 = std::stoi(argv[20]); + BatchStrideD1 = std::stoi(argv[21]); + BatchStrideE1 = std::stoi(argv[22]); + } + else + { + printf("arg1: verification (0=no, 1=yes)\n"); + printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"); + printf("arg3: time kernel (0=no, 1=yes)\n"); + printf("arg4 to 8: M, N, K, O, Batch\n"); + printf( + "arg9 to 15: StrideA0, StrideB0, StrideD00, StrideD01, StrideB1, StrideD1, StrideE1\n"); + printf("arg16 to 22: BatchStrideA0, BatchStrideB0, BatchStrideD00, BatchStrideD01, " + "BatchStrideB1, BatchStrideD1, BatchStrideE1 \n"); + exit(0); + } + + const int DefaultStrideA0 = ck::is_same_v ? K : M; + const int DefaultStrideB0 = ck::is_same_v ? N : K; + const int DefaultStrideD00 = ck::is_same_v ? N : M; + const int DefaultStrideD01 = ck::is_same_v ? N : M; + const int DefaultStrideB1 = ck::is_same_v ? O : N; + const int DefaultStrideD1 = ck::is_same_v ? O : M; + const int DefaultStrideE1 = ck::is_same_v ? O : M; + + StrideA0 = (StrideA0 < 0) ? DefaultStrideA0 : StrideA0; + StrideB0 = (StrideB0 < 0) ? DefaultStrideB0 : StrideB0; + StrideD00 = (StrideD00 < 0) ? DefaultStrideD00 : StrideD00; + StrideD01 = (StrideD01 < 0) ? DefaultStrideD01 : StrideD01; + StrideB1 = (StrideB1 < 0) ? DefaultStrideB1 : StrideB1; + StrideD1 = (StrideD1 < 0) ? DefaultStrideD1 : StrideD1; + StrideE1 = (StrideE1 < 0) ? DefaultStrideE1 : StrideE1; + + const int DefaultBatchStrideA0 = (ck::is_same_v ? K : M) * StrideA0; + const int DefaultBatchStrideB0 = (ck::is_same_v ? N : K) * StrideB0; + const int DefaultBatchStrideD00 = (ck::is_same_v ? N : M) * StrideD00; + const int DefaultBatchStrideD01 = (ck::is_same_v ? N : M) * StrideD01; + const int DefaultBatchStrideB1 = (ck::is_same_v ? O : N) * StrideB1; + const int DefaultBatchStrideD1 = (ck::is_same_v ? O : M) * StrideD1; + const int DefaultBatchStrideE1 = (ck::is_same_v ? O : M) * StrideE1; + + BatchStrideA0 = BatchStrideA0 < 0 ? DefaultBatchStrideA0 : BatchStrideA0; + BatchStrideB0 = BatchStrideB0 < 0 ? DefaultBatchStrideB0 : BatchStrideB0; + BatchStrideD00 = BatchStrideD00 < 0 ? DefaultBatchStrideD00 : BatchStrideD00; + BatchStrideD01 = BatchStrideD01 < 0 ? DefaultBatchStrideD01 : BatchStrideD01; + BatchStrideB1 = BatchStrideB1 < 0 ? DefaultBatchStrideB1 : BatchStrideB1; + BatchStrideD1 = BatchStrideD1 < 0 ? DefaultBatchStrideD1 : BatchStrideD1; + BatchStrideE1 = BatchStrideE1 < 0 ? DefaultBatchStrideE1 : BatchStrideE1; + + auto f_host_tensor_descriptor = [](std::size_t batch_count, + std::size_t row, + std::size_t col, + std::size_t stride, + std::size_t batch_stride, + auto layout) { + using namespace ck::literals; + + if(std::is_same::value) + { + return HostTensorDescriptor( + {batch_count, row, col}, {batch_stride, stride, 1_uz}, layout); + } + else + { + return HostTensorDescriptor( + {batch_count, row, col}, {batch_stride, 1_uz, stride}, layout); + } + }; + + // E_m_o = A_m_k * B0_k_n * B1_n_o + Tensor a0_g_m_k( + f_host_tensor_descriptor(BatchCount, M, K, StrideA0, BatchStrideA0, A0Layout{})); + Tensor b0_g_k_n( + f_host_tensor_descriptor(BatchCount, K, N, StrideB0, BatchStrideB0, B0Layout{})); + Tensor d00_g_m_n( + f_host_tensor_descriptor(BatchCount, M, N, StrideD00, BatchStrideD00, D00Layout{})); + Tensor d01_g_m_n( + f_host_tensor_descriptor(BatchCount, M, N, StrideD01, BatchStrideD01, D01Layout{})); + Tensor b1_g_n_o( + f_host_tensor_descriptor(BatchCount, N, O, StrideB1, BatchStrideB1, B1Layout{})); + Tensor d1_g_m_o( + f_host_tensor_descriptor(BatchCount, M, O, StrideD1, BatchStrideD1, D1Layout{})); + Tensor e1_g_m_o_host_result( + f_host_tensor_descriptor(BatchCount, M, O, StrideE1, BatchStrideE1, E1Layout{})); + Tensor e1_g_m_o_device_result( + f_host_tensor_descriptor(BatchCount, M, O, StrideE1, BatchStrideE1, E1Layout{})); + + std::cout << "a0_g_m_k: " << a0_g_m_k.mDesc << std::endl; + std::cout << "b0_g_k_n: " << b0_g_k_n.mDesc << std::endl; + std::cout << "d00_g_m_n: " << d00_g_m_n.mDesc + << " size: " << d00_g_m_n.mDesc.GetElementSpaceSize() << std::endl; + std::cout << "d01_g_m_n: " << d01_g_m_n.mDesc + << " size: " << d01_g_m_n.mDesc.GetElementSpaceSize() << std::endl; + std::cout << "b1_g_n_o: " << b1_g_n_o.mDesc << std::endl; + std::cout << "e1_g_m_o: " << e1_g_m_o_host_result.mDesc << std::endl; + + switch(init_method) + { + case 0: break; + case 1: + a0_g_m_k.GenerateTensorValue(GeneratorTensor_2{-2, 3}); + b0_g_k_n.GenerateTensorValue(GeneratorTensor_2{-2, 3}); + d00_g_m_n.GenerateTensorValue(GeneratorTensor_2{-2, 3}); + d01_g_m_n.GenerateTensorValue(GeneratorTensor_2{-2, 3}); + b1_g_n_o.GenerateTensorValue(GeneratorTensor_2{-2, 3}); + d1_g_m_o.GenerateTensorValue(GeneratorTensor_2{-2, 3}); + break; + case 2: + a0_g_m_k.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + b0_g_k_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + d00_g_m_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + d01_g_m_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + b1_g_n_o.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + d1_g_m_o.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + break; + default: + a0_g_m_k.GenerateTensorValue(GeneratorTensor_1{1}); + b0_g_k_n.GenerateTensorValue(GeneratorTensor_Sequential{}); + d00_g_m_n.GenerateTensorValue(GeneratorTensor_1{1}); + d01_g_m_n.GenerateTensorValue(GeneratorTensor_1{1}); + b1_g_n_o.GenerateTensorValue(GeneratorTensor_Diagonal{}); + d1_g_m_o.GenerateTensorValue(GeneratorTensor_1{1}); + } + + DeviceMem a0_g_m_k_device_buf(sizeof(A0DataType) * a0_g_m_k.mDesc.GetElementSize()); + DeviceMem b0_g_k_n_device_buf(sizeof(B0DataType) * b0_g_k_n.mDesc.GetElementSize()); + DeviceMem d00_g_m_n_device_buf(sizeof(D00DataType) * d00_g_m_n.mDesc.GetElementSpaceSize()); + DeviceMem d01_g_m_n_device_buf(sizeof(D01DataType) * d01_g_m_n.mDesc.GetElementSpaceSize()); + DeviceMem b1_g_n_o_device_buf(sizeof(B1DataType) * b1_g_n_o.mDesc.GetElementSize()); + DeviceMem e1_g_m_o_device_buf(sizeof(E1DataType) * + e1_g_m_o_device_result.mDesc.GetElementSize()); + DeviceMem d1_g_m_o_device_buf(sizeof(D1DataType) * d1_g_m_o.mDesc.GetElementSpaceSize()); + + a0_g_m_k_device_buf.ToDevice(a0_g_m_k.mData.data()); + b0_g_k_n_device_buf.ToDevice(b0_g_k_n.mData.data()); + d00_g_m_n_device_buf.ToDevice(d00_g_m_n.mData.data()); + d01_g_m_n_device_buf.ToDevice(d01_g_m_n.mData.data()); + b1_g_n_o_device_buf.ToDevice(b1_g_n_o.mData.data()); + d1_g_m_o_device_buf.ToDevice(d1_g_m_o.mData.data()); + + auto a0_element_op = A0ElementOp{}; + auto b0_element_op = B0ElementOp{}; + auto cde0_element_op = CDE0ElementOp{}; + auto b1_element_op = B1ElementOp{}; + auto cde1_element_op = CDE1ElementOp{}; + + // do GEMM + auto gemm = DeviceGemmInstance{}; + auto invoker = gemm.MakeInvoker(); + auto argument = + gemm.MakeArgument(static_cast(a0_g_m_k_device_buf.GetDeviceBuffer()), + static_cast(b0_g_k_n_device_buf.GetDeviceBuffer()), + std::array{d00_g_m_n_device_buf.GetDeviceBuffer(), + d01_g_m_n_device_buf.GetDeviceBuffer()}, + static_cast(b1_g_n_o_device_buf.GetDeviceBuffer()), + std::array{d1_g_m_o_device_buf.GetDeviceBuffer()}, + static_cast(e1_g_m_o_device_buf.GetDeviceBuffer()), + M, + N, + K, + O, + BatchCount, + StrideA0, + StrideB0, + std::array{StrideD00, StrideD01}, + StrideB1, + std::array{StrideD1}, + StrideE1, + BatchStrideA0, + BatchStrideB0, + std::array{BatchStrideD00, BatchStrideD01}, + BatchStrideB1, + std::array{BatchStrideD1}, + BatchStrideE1, + a0_element_op, + b0_element_op, + cde0_element_op, + b1_element_op, + cde1_element_op); + + if(!gemm.IsSupportedArgument(argument)) + { + std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl; + + return 0; + } + + float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel}); + + std::size_t flop = (size_t(M) * N * K * 2 + size_t(M) * N * O * 2) * BatchCount; + std::size_t num_btype = + (sizeof(A0DataType) * M * K + sizeof(B0DataType) * K * N + sizeof(D00DataType) * N + + sizeof(D01DataType) * N + sizeof(B1DataType) * N * O + sizeof(E1DataType) * M * O + + sizeof(D1DataType) * O) * + BatchCount; + + float tflops = static_cast(flop) / 1.E9 / ave_time; + + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " + << gemm.GetTypeString() << std::endl; + + e1_g_m_o_device_buf.FromDevice(e1_g_m_o_device_result.mData.data()); + + if(do_verification) + { + using ReferenceGemm0Instance = + ck::tensor_operation::host::ReferenceBatchedGemm; + + using ReferenceGemm1Instance = + ck::tensor_operation::host::ReferenceBatchedGemm; + + // Output of Gemm0 is input A of Gemm1 + Tensor c0_g_m_n(f_host_tensor_descriptor(BatchCount, M, N, N, M * N, Row{})); + Tensor e0_g_m_n(f_host_tensor_descriptor(BatchCount, M, N, N, M * N, Row{})); + Tensor c1_g_m_o(f_host_tensor_descriptor(BatchCount, M, O, O, M * O, Row{})); + + auto ref_gemm0 = ReferenceGemm0Instance{}; + auto ref_gemm0_invoker = ref_gemm0.MakeInvoker(); + auto ref_gemm0_argument = ref_gemm0.MakeArgument( + a0_g_m_k, b0_g_k_n, c0_g_m_n, a0_element_op, b0_element_op, PassThrough{}); + + ref_gemm0_invoker.Run(ref_gemm0_argument); + + // bias+bias+relu + // Note that we also convert from AccDataType to A0DataType to match what the device + // operation does + e0_g_m_n.ForEach([&](auto&, auto idx) { + AccDataType out; + cde0_element_op(out, c0_g_m_n(idx), d00_g_m_n(idx), d01_g_m_n(idx)); + e0_g_m_n(idx) = ck::type_convert(out); + }); + + auto ref_gemm1 = ReferenceGemm1Instance{}; + auto ref_gemm1_invoker = ref_gemm1.MakeInvoker(); + auto ref_gemm1_argument = ref_gemm1.MakeArgument( + e0_g_m_n, b1_g_n_o, c1_g_m_o, PassThrough{}, b1_element_op, PassThrough{}); + + ref_gemm1_invoker.Run(ref_gemm1_argument); + + // bias + e1_g_m_o_host_result.ForEach([&](auto&, auto idx) { + cde1_element_op(e1_g_m_o_host_result(idx), c1_g_m_o(idx), d1_g_m_o(idx)); + }); + + // NOTE: For float initialization (mode 2) verification currently fails due to inaccuracy. + // This seems to just be accumulating errors due to double gemm. It only seems to happen + // when using B1 tensor containing negative values, as this can get large values from gemm0 + // back to zero again but reduce the tolerance allowed by the relative tolerance. + // + // There doesn't seem to be any bug with the implementation, just a difference in order of + // operations between CPU and GPU causing an accumulating error. + bool validation_result = ck::utils::check_err(e1_g_m_o_device_result, e1_g_m_o_host_result); + std::cout << "Validation result: " << (validation_result ? "SUCCESS" : "FAIL") << "." + << std::endl; + + return validation_result ? 0 : 1; + } + + return 0; +} diff --git a/example/37_batched_gemm_add_add_relu_gemm_add/element_ops.h b/example/37_batched_gemm_add_add_relu_gemm_add/element_ops.h new file mode 100644 index 0000000000..e9f22b862e --- /dev/null +++ b/example/37_batched_gemm_add_add_relu_gemm_add/element_ops.h @@ -0,0 +1,58 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck/ck.hpp" + +// E = Relu(C + D0 + D1) +struct AddAddRelu +{ + __host__ __device__ void + operator()(ck::half_t& e, const ck::half_t& c, const ck::half_t& d0, const ck::half_t& d1) const + { + const ck::half_t x = c + d0 + d1; + + ck::tensor_operation::element_wise::Relu{}.operator()(e, x); + } + __host__ __device__ void + operator()(float& e, const float& c, const ck::half_t& d0, const ck::half_t& d1) const + { + const float x = c + ck::type_convert(d0) + ck::type_convert(d1); + + ck::tensor_operation::element_wise::Relu{}.operator()(e, x); + } +}; + +// E = Gelu(C + D0 + D1) +struct AddAddGelu +{ + __host__ __device__ void + operator()(ck::half_t& e, const ck::half_t& c, const ck::half_t& d0, const ck::half_t& d1) const + { + const ck::half_t x = c + d0 + d1; + + ck::tensor_operation::element_wise::Gelu{}.template operator()(e, + x); + } + + __host__ __device__ void + operator()(float& e, const float& c, const ck::half_t& d0, const ck::half_t& d1) const + { + const float x = c + ck::type_convert(d0) + ck::type_convert(d1); + + ck::tensor_operation::element_wise::Gelu{}.template operator()(e, x); + } +}; + +// E = FastGelu(C + D0 + D1) +struct AddAddFastGelu +{ + __host__ __device__ void + operator()(float& e, const float& c, const ck::half_t& d0, const ck::half_t& d1) const + { + const float x = c + ck::type_convert(d0) + ck::type_convert(d1); + + ck::tensor_operation::element_wise::FastGelu{}.template operator()(e, x); + } +}; diff --git a/example/65_gemm_multiply_multiply/gemm_add_add_wmma_fp16.cpp b/example/65_gemm_multiply_multiply/gemm_add_add_wmma_fp16.cpp index 24c58bb69a..1e3d946bad 100644 --- a/example/65_gemm_multiply_multiply/gemm_add_add_wmma_fp16.cpp +++ b/example/65_gemm_multiply_multiply/gemm_add_add_wmma_fp16.cpp @@ -31,8 +31,9 @@ using S = ck::Sequence; using F16 = ck::half_t; using F32 = float; -using Row = ck::tensor_layout::gemm::RowMajor; -using Col = ck::tensor_layout::gemm::ColumnMajor; +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; +using Bypass = ck::tensor_layout::BypassLayoutVerification; using A0DataType = F16; using B0DataType = F16; @@ -139,11 +140,11 @@ int main(int argc, char* argv[]) if(std::is_same::value) { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); + return HostTensorDescriptor({row, col}, {stride, 1_uz}, Bypass{}); } else { - return HostTensorDescriptor({row, col}, {1_uz, stride}); + return HostTensorDescriptor({row, col}, {1_uz, stride}, Bypass{}); } }; diff --git a/example/65_gemm_multiply_multiply/run_gemm_multiply_multiply_wp_example.inc b/example/65_gemm_multiply_multiply/run_gemm_multiply_multiply_wp_example.inc index 2de3222380..10dce7fe64 100644 --- a/example/65_gemm_multiply_multiply/run_gemm_multiply_multiply_wp_example.inc +++ b/example/65_gemm_multiply_multiply/run_gemm_multiply_multiply_wp_example.inc @@ -5,6 +5,8 @@ int run_gemm_example(int argc, char* argv[]) { + using Bypass = ck::tensor_layout::BypassLayoutVerification; + bool do_verification = true; int init_method = 1; bool time_kernel = false; @@ -64,11 +66,11 @@ int run_gemm_example(int argc, char* argv[]) if(std::is_same::value) { - return ck::HostTensorDescriptor({row, col}, {stride, 1_uz}); + return ck::HostTensorDescriptor({row, col}, {stride, 1_uz}, Bypass{}); } else { - return ck::HostTensorDescriptor({row, col}, {1_uz, stride}); + return ck::HostTensorDescriptor({row, col}, {1_uz, stride}, Bypass{}); } }; diff --git a/example/68_gemm_add/common.hpp b/example/68_gemm_add/common.hpp index 362dc2fff2..12d4b381b2 100644 --- a/example/68_gemm_add/common.hpp +++ b/example/68_gemm_add/common.hpp @@ -87,7 +87,7 @@ parse_cmd_args(int argc, char* argv[], ProblemSize& problem_size, ExecutionConfi config.init_method = std::stoi(argv[2]); config.time_kernel = std::stoi(argv[3]); } - else if(argc == 13) + else if(argc == 11) { config.do_verification = std::stoi(argv[1]); config.init_method = std::stoi(argv[2]); diff --git a/example/68_gemm_add/run_gemm_add_example_wmma.inc b/example/68_gemm_add/run_gemm_add_example_wmma.inc index ba15d03e07..0f2cc08edf 100644 --- a/example/68_gemm_add/run_gemm_add_example_wmma.inc +++ b/example/68_gemm_add/run_gemm_add_example_wmma.inc @@ -6,6 +6,7 @@ bool run_gemm_add(const ProblemSize& problem_size, const ExecutionConfig& config) { using namespace ck::literals; + using Bypass = ck::tensor_layout::BypassLayoutVerification; auto& [M, N, K, StrideA, StrideB, StrideD, StrideE] = problem_size; @@ -13,11 +14,11 @@ bool run_gemm_add(const ProblemSize& problem_size, const ExecutionConfig& config [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { if(std::is_same::value) { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); + return HostTensorDescriptor({row, col}, {stride, 1_uz}, Bypass{}); } else { - return HostTensorDescriptor({row, col}, {1_uz, stride}); + return HostTensorDescriptor({row, col}, {1_uz, stride}, Bypass{}); } }; diff --git a/example/68_gemm_add/run_gemm_add_example_xdl.inc b/example/68_gemm_add/run_gemm_add_example_xdl.inc index da22230a4e..186423d32f 100644 --- a/example/68_gemm_add/run_gemm_add_example_xdl.inc +++ b/example/68_gemm_add/run_gemm_add_example_xdl.inc @@ -6,6 +6,7 @@ bool run_gemm_add(const ProblemSize& problem_size, const ExecutionConfig& config) { using namespace ck::literals; + using Bypass = ck::tensor_layout::BypassLayoutVerification; auto& [M, N, K, StrideA, StrideB, StrideD, StrideE] = problem_size; @@ -13,11 +14,11 @@ bool run_gemm_add(const ProblemSize& problem_size, const ExecutionConfig& config [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { if(std::is_same::value) { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); + return HostTensorDescriptor({row, col}, {stride, 1_uz}, Bypass{}); } else { - return HostTensorDescriptor({row, col}, {1_uz, stride}); + return HostTensorDescriptor({row, col}, {1_uz, stride}, Bypass{}); } }; diff --git a/example/69_gemm_add_relu/common.hpp b/example/69_gemm_add_relu/common.hpp index e54c5317ae..de84d69a5e 100644 --- a/example/69_gemm_add_relu/common.hpp +++ b/example/69_gemm_add_relu/common.hpp @@ -87,7 +87,7 @@ parse_cmd_args(int argc, char* argv[], ProblemSize& problem_size, ExecutionConfi config.init_method = std::stoi(argv[2]); config.time_kernel = std::stoi(argv[3]); } - else if(argc == 13) + else if(argc == 11) { config.do_verification = std::stoi(argv[1]); config.init_method = std::stoi(argv[2]); diff --git a/example/69_gemm_add_relu/run_gemm_add_relu_example_wmma.inc b/example/69_gemm_add_relu/run_gemm_add_relu_example_wmma.inc index 8deac6dec8..c3cfd00ab3 100644 --- a/example/69_gemm_add_relu/run_gemm_add_relu_example_wmma.inc +++ b/example/69_gemm_add_relu/run_gemm_add_relu_example_wmma.inc @@ -6,6 +6,7 @@ bool run_gemm_add_relu(const ProblemSize& problem_size, const ExecutionConfig& config) { using namespace ck::literals; + using Bypass = ck::tensor_layout::BypassLayoutVerification; auto& [M, N, K, StrideA, StrideB, StrideD, StrideE] = problem_size; @@ -13,11 +14,11 @@ bool run_gemm_add_relu(const ProblemSize& problem_size, const ExecutionConfig& c [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { if(std::is_same::value) { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); + return HostTensorDescriptor({row, col}, {stride, 1_uz}, Bypass{}); } else { - return HostTensorDescriptor({row, col}, {1_uz, stride}); + return HostTensorDescriptor({row, col}, {1_uz, stride}, Bypass{}); } }; diff --git a/example/69_gemm_add_relu/run_gemm_add_relu_example_xdl.inc b/example/69_gemm_add_relu/run_gemm_add_relu_example_xdl.inc index df7474bab5..cca85aa11c 100644 --- a/example/69_gemm_add_relu/run_gemm_add_relu_example_xdl.inc +++ b/example/69_gemm_add_relu/run_gemm_add_relu_example_xdl.inc @@ -6,6 +6,7 @@ bool run_gemm_add_relu(const ProblemSize& problem_size, const ExecutionConfig& config) { using namespace ck::literals; + using Bypass = ck::tensor_layout::BypassLayoutVerification; auto& [M, N, K, StrideA, StrideB, StrideD, StrideE] = problem_size; @@ -13,11 +14,11 @@ bool run_gemm_add_relu(const ProblemSize& problem_size, const ExecutionConfig& c [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { if(std::is_same::value) { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); + return HostTensorDescriptor({row, col}, {stride, 1_uz}, Bypass{}); } else { - return HostTensorDescriptor({row, col}, {1_uz, stride}); + return HostTensorDescriptor({row, col}, {1_uz, stride}, Bypass{}); } }; diff --git a/example/ck_tile/01_fmha/fmha_bwd.hpp b/example/ck_tile/01_fmha/fmha_bwd.hpp index d1b55168e3..180d039cd4 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd.hpp @@ -189,7 +189,7 @@ struct fmha_bwd_args ck_tile::index_t nhead_stride_randval; ck_tile::index_t nhead_stride_do; ck_tile::index_t nhead_stride_lsed; - ck_tile::index_t nhead_stride_dq_acc; + ck_tile::long_index_t nhead_stride_dq_acc; ck_tile::index_t nhead_stride_dq; ck_tile::index_t nhead_stride_dk; ck_tile::index_t nhead_stride_dv; @@ -202,7 +202,7 @@ struct fmha_bwd_args ck_tile::index_t batch_stride_randval; ck_tile::index_t batch_stride_do; ck_tile::index_t batch_stride_lsed; - ck_tile::index_t batch_stride_dq_acc; + ck_tile::long_index_t batch_stride_dq_acc; ck_tile::index_t batch_stride_dq; ck_tile::index_t batch_stride_dk; ck_tile::index_t batch_stride_dv; diff --git a/example/ck_tile/01_fmha/fmha_bwd_runner.hpp b/example/ck_tile/01_fmha/fmha_bwd_runner.hpp index d62b908e33..f41f0668e5 100644 --- a/example/ck_tile/01_fmha/fmha_bwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd_runner.hpp @@ -287,9 +287,7 @@ bwd_result fmha_bwd_run(mode_enum mode, ? get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, max_seqlen_k) : std::array{1, 1, 1, 1} /* dummy shape for simplifying code */); ck_tile::HostTensor dq_acc_host( - i_perm - ? std::array{nsplits, shape_batch, nhead, shape_seqlen_q, hdim_q} - : std::array{nsplits, shape_batch, shape_seqlen_q, nhead, hdim_q}); + std::array{shape_batch, nhead, nsplits, shape_seqlen_q, hdim_q}); if(init_method == "ui" || init_method == "0") { @@ -433,6 +431,7 @@ bwd_result fmha_bwd_run(mode_enum mode, const ck_tile::index_t stride_dk = (i_perm ? hdim_q : nhead * hdim_q); const ck_tile::index_t stride_dv = (i_perm ? hdim_v : nhead * hdim_v); const ck_tile::index_t stride_dbias = (i_perm ? max_seqlen_k : nhead * max_seqlen_k); + const auto split_stride_dq_acc = (shape_seqlen_q * hdim_q); // setup nhead_stride_* arguments const ck_tile::index_t nhead_stride_q = (i_perm ? shape_seqlen_q * hdim_q : hdim_q); const ck_tile::index_t nhead_stride_k = (i_perm ? shape_seqlen_k * hdim_q : hdim_q); @@ -444,6 +443,8 @@ bwd_result fmha_bwd_run(mode_enum mode, const ck_tile::index_t nhead_stride_lsed = shape_seqlen_q; const ck_tile::index_t nhead_stride_dbias = (i_perm ? shape_seqlen_q * max_seqlen_k : max_seqlen_k); + const auto nhead_stride_dq_acc = + static_cast(split_stride_dq_acc) * nsplits; // setup batch_stride_* arguments const ck_tile::index_t batch_stride_q = (nhead * shape_seqlen_q * hdim_q); const ck_tile::index_t batch_stride_k = (nhead_k * shape_seqlen_k * hdim_q); @@ -456,8 +457,7 @@ bwd_result fmha_bwd_run(mode_enum mode, const ck_tile::index_t batch_stride_dk = (nhead * shape_seqlen_k * hdim_q); const ck_tile::index_t batch_stride_dv = (nhead * shape_seqlen_k * hdim_v); const ck_tile::index_t batch_stride_dbias = (nhead * shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t split_stride_dq_acc = - (shape_batch * nhead * shape_seqlen_q * hdim_q); + const auto batch_stride_dq_acc = nhead * nhead_stride_dq_acc; const auto drop_seed_offset = [&]() -> decltype(fmha_bwd_args::drop_seed_offset) { if(drop_prefs) @@ -513,7 +513,7 @@ bwd_result fmha_bwd_run(mode_enum mode, stride_o, stride_randval, stride_do, - stride_q, // stride_dq_acc + hdim_q, // stride_dq_acc stride_q, // stride_dq stride_dk, stride_dv, @@ -526,7 +526,7 @@ bwd_result fmha_bwd_run(mode_enum mode, nhead_stride_randval, nhead_stride_do, nhead_stride_lsed, - nhead_stride_q, // nhead_stride_dq_acc + nhead_stride_dq_acc, nhead_stride_q, // nhead_stride_dq nhead_stride_k, // nhead_stride_dk nhead_stride_v, // nhead_stride_dv @@ -539,7 +539,7 @@ bwd_result fmha_bwd_run(mode_enum mode, batch_stride_randval, batch_stride_do, batch_stride_lsed, - batch_stride_q, // batch_stride_dq_acc + batch_stride_dq_acc, batch_stride_q, // batch_stride_dq batch_stride_dk, batch_stride_dv, diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp index 8eff0e7469..c1df27ecc8 100644 --- a/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/example/ck_tile/03_gemm/gemm_utils.hpp @@ -456,7 +456,8 @@ inline auto create_args() .insert("json", "0", "0: No Json, 1: Dump Results in Json format") .insert("jsonfile", "gemm.json", "json file name to dump results") .insert("flush_cache", "true", "flush cache before running the kernel, defaults to true") - .insert("rotating_count", "1000", "rotating count, defaults to 1000"); + .insert("rotating_count", "1000", "rotating count, defaults to 1000") + .insert("test_async", "0", "0: normal gemm, 1: test async input scheduler"); return arg_parser; } diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index c1c8a2fc89..ace9152747 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -12,6 +12,169 @@ #include "run_gemm_example_common.hpp" #include "universal_gemm_invoker.hpp" +// Universal GEMM-specific wrapper that handles test_async flag +template +int run_gemm_example_with_layouts_universal(ck_tile::ArgParser& arg_parser, + const ALayout a_layout = ALayout{}, + const BLayout b_layout = BLayout{}, + const CLayout c_layout = CLayout{}) +{ + using Invoker = UniversalInvoker; + using AccDataType = typename GemmTypeConfig::AccDataType; + + // Check for async input scheduler test mode + bool test_async = arg_parser.get_int("test_async"); + if(test_async) + { + // Extract parameters for async test (same as shared implementation) + const ck_tile::index_t M = arg_parser.get_int("m"); + const ck_tile::index_t N = arg_parser.get_int("n"); + const ck_tile::index_t K = arg_parser.get_int("k"); + const ck_tile::index_t kbatch = arg_parser.get_int("split_k"); + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + constexpr bool is_a_row_major = std::is_same_v; + constexpr bool is_b_row_major = std::is_same_v; + constexpr bool is_c_row_major = std::is_same_v; + + const ck_tile::index_t stride_A = is_a_row_major ? K : M; + const ck_tile::index_t stride_B = is_b_row_major ? N : K; + const ck_tile::index_t stride_C = is_c_row_major ? N : M; + + // Allocate and initialize tensors + ck_tile::HostTensor a_m_k(ck_tile::host_tensor_descriptor( + M, K, stride_A, ck_tile::bool_constant{})); + ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor( + K, N, stride_B, ck_tile::bool_constant{})); + ck_tile::HostTensor c_m_n_dev_result(ck_tile::host_tensor_descriptor( + M, N, stride_C, ck_tile::bool_constant{})); + + ck_tile::FillUniformDistributionIntegerValue{-5, 5}(a_m_k); + ck_tile::FillUniformDistributionIntegerValue{-5, 5}(b_k_n); + + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); + + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + c_m_n_dev_buf.SetZero(); + c_m_n_dev_result.SetZero(); + + ck_tile::GemmHostArgs args = {a_m_k_dev_buf.GetDeviceBuffer(), + b_k_n_dev_buf.GetDeviceBuffer(), + c_m_n_dev_buf.GetDeviceBuffer(), + kbatch, + M, + N, + K, + stride_A, + stride_B, + stride_C}; + + Invoker::template test_async_input_scheduler, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough>( + args, ck_tile::stream_config{nullptr, false, 1}); + + // Copy result from device for verification + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); + + // Compute CPU reference + ck_tile::HostTensor c_m_n_ref(ck_tile::host_tensor_descriptor( + M, N, stride_C, ck_tile::bool_constant{})); + c_m_n_ref.SetZero(); + ck_tile::reference_gemm( + a_m_k, b_k_n, c_m_n_ref); + + // Verify results + const float max_accumulated_value = + *std::max_element(c_m_n_ref.mData.begin(), c_m_n_ref.mData.end()); + const auto rtol_atol = calculate_rtol_atol( + K, kbatch, max_accumulated_value); + bool pass = do_verify(c_m_n_dev_result, c_m_n_ref, rtol_atol, "CPU"); + + std::cout << "Async input scheduler test: " << (pass ? "PASS" : "FAIL") << std::endl; + return pass; + } + + // Normal path - delegate to shared implementation + return run_gemm_example_with_layouts( + arg_parser, a_layout, b_layout, c_layout); +} + +// Universal GEMM-specific prec_type dispatcher that uses the wrapper +template +int run_gemm_example_prec_type_universal(std::string a_layout, + std::string b_layout, + ck_tile::ArgParser& arg_parser) +{ + using Row = ck_tile::tensor_layout::gemm::RowMajor; + using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + bool preshuffle = GemmConfig::Preshuffle; + + if(preshuffle && std::is_same_v) + { + throw std::runtime_error("Preshuffle is not supported for this int4 datatype!"); + } + + if(preshuffle && a_layout != "R" && b_layout != "C") + { + throw std::runtime_error( + "Preshuffle is supported only for A(Row major), B(column major) input matrices!"); + } + + using LayoutVariant = std::variant; + + auto string_to_layout = [](const std::string& layout) -> LayoutVariant { + if(layout == "R") + return Row{}; + if(layout == "C") + return Col{}; + throw std::runtime_error("Unsupported layout: " + layout); + }; + + auto a_layout_variant = string_to_layout(a_layout); + auto b_layout_variant = string_to_layout(b_layout); + + return std::visit( + [&](auto a_layout_type, auto b_layout_type) -> int { + if constexpr(std::is_same_v && + std::is_same_v) + { + throw std::runtime_error("Unsupported memory layout for the input matrices when " + "BPrecType is ck_tile::pk_int4_t!"); + } + else + { + return run_gemm_example_with_layouts_universal( + arg_parser, a_layout_type, b_layout_type, Row{}); + } + }, + a_layout_variant, + b_layout_variant); +} + template