Spaces:
Sleeping
Sleeping
| /*************************************************************************************************** | |
| * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| * SPDX-License-Identifier: BSD-3-Clause | |
| * | |
| * Redistribution and use in source and binary forms, with or without | |
| * modification, are permitted provided that the following conditions are met: | |
| * | |
| * 1. Redistributions of source code must retain the above copyright notice, this | |
| * list of conditions and the following disclaimer. | |
| * | |
| * 2. Redistributions in binary form must reproduce the above copyright notice, | |
| * this list of conditions and the following disclaimer in the documentation | |
| * and/or other materials provided with the distribution. | |
| * | |
| * 3. Neither the name of the copyright holder nor the names of its | |
| * contributors may be used to endorse or promote products derived from | |
| * this software without specific prior written permission. | |
| * | |
| * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" | |
| * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE | |
| * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE | |
| * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE | |
| * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL | |
| * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR | |
| * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER | |
| * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, | |
| * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE | |
| * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. | |
| * | |
| **************************************************************************************************/ | |
| /*! \file | |
| \brief CUTLASS Library handle. | |
| */ | |
| #include <iostream> | |
| #include <stdexcept> | |
| #include <cstdint> | |
| #include "cutlass/library/handle.h" | |
| #include "cutlass/library/singleton.h" | |
| #include "cutlass/library/util.h" | |
| namespace cutlass { | |
| namespace library { | |
| /////////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Constructor | |
| Handle::Handle( | |
| cudaStream_t stream, | |
| size_t workspace_size | |
| ): | |
| provider_(Provider::kCUTLASS), | |
| stream_(stream), | |
| workspace_(nullptr), | |
| workspace_size_(0), | |
| scalar_pointer_mode_(ScalarPointerMode::kHost), | |
| last_operation_(nullptr) { | |
| int device_idx = -1; | |
| cudaError_t error = cudaGetDevice(&device_idx); | |
| if (error != cudaSuccess) { | |
| throw std::runtime_error("cudaGetDevice() failed"); | |
| } | |
| error = cudaGetDeviceProperties(&device_, device_idx); | |
| if (error != cudaSuccess) { | |
| throw std::runtime_error("cudaGetDeviceProperties() failed"); | |
| } | |
| set_workspace_size(workspace_size); | |
| Singleton::get(); | |
| } | |
| /// Destructor | |
| Handle::~Handle() { | |
| if (workspace_) { | |
| if (workspace_) { | |
| cudaFree(workspace_); | |
| } | |
| workspace_ = nullptr; | |
| workspace_size_ = 0; | |
| } | |
| } | |
| /// Move constructor | |
| Handle::Handle(Handle && handle) { | |
| device_ = handle.device_; | |
| workspace_size_ = handle.workspace_size_; | |
| workspace_ = handle.workspace_; | |
| stream_ = handle.stream_; | |
| scalar_pointer_mode_ = handle.scalar_pointer_mode_; | |
| handle.workspace_ = nullptr; | |
| handle.workspace_size_ = 0; | |
| } | |
| /// Move assignment operator | |
| Handle & Handle::operator=(Handle && handle) { | |
| provider_ = handle.provider_; | |
| device_ = handle.device_; | |
| workspace_size_ = handle.workspace_size_; | |
| workspace_ = handle.workspace_; | |
| stream_ = handle.stream_; | |
| scalar_pointer_mode_ = handle.scalar_pointer_mode_; | |
| handle.workspace_ = nullptr; | |
| handle.workspace_size_ = 0; | |
| return *this; | |
| } | |
| int Handle::compute_capability() const { | |
| return device_.major * 10 + device_.minor; | |
| } | |
| /// Sets the current CUDA stream | |
| void Handle::set_stream(cudaStream_t stream) { | |
| stream_ = stream; | |
| } | |
| /// Gets the current CUDA stream | |
| cudaStream_t Handle::get_stream() const { | |
| return stream_; | |
| } | |
| /// Gets the current provider | |
| Provider Handle::get_provider() const { | |
| return provider_; | |
| } | |
| /// Sets the provider of operations | |
| void Handle::set_provider(Provider provider) { | |
| provider_ = provider; | |
| } | |
| /// Gets the device workspace size | |
| size_t Handle::get_workspace_size() const { | |
| return workspace_size_; | |
| } | |
| /// Gets a pointer to the device workspace allocation in Global Memory | |
| void *Handle::get_workspace() const { | |
| return workspace_; | |
| } | |
| /// Sets the size of device workspace, invalidating previous calls to get_device_workspace() | |
| void Handle::set_workspace_size(size_t bytes) { | |
| if (bytes != workspace_size_) { | |
| if (workspace_) { | |
| cudaFree(workspace_); | |
| } | |
| workspace_ = nullptr; | |
| workspace_size_ = bytes; | |
| if (workspace_size_) { | |
| cudaError_t error = cudaMalloc((void **)&workspace_, workspace_size_); | |
| if (error != cudaSuccess) { | |
| throw std::runtime_error("Failed to allocate workspace"); | |
| } | |
| } | |
| } | |
| if (workspace_) { | |
| cudaError_t error = cudaMemset(workspace_, 0, workspace_size_); | |
| if (error != cudaSuccess) { | |
| throw std::runtime_error("Failed to clear workspace"); | |
| } | |
| } | |
| } | |
| /// Gets the scalar pointer mode | |
| ScalarPointerMode Handle::get_scalar_pointer_mode() const { | |
| return scalar_pointer_mode_; | |
| } | |
| /// Sets the scalar pointer mode | |
| void Handle::set_scalar_pointer_mode(ScalarPointerMode mode) { | |
| scalar_pointer_mode_ = mode; | |
| } | |
| /// Gets the last operation | |
| Operation const *Handle::get_last_operation() const { | |
| return last_operation_; | |
| } | |
| /////////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Returns the maximum required alignment for each operator | |
| static int maximum_alignment_requirement(GemmDescription const &desc) { | |
| return std::max( | |
| std::max(desc.A.alignment, desc.B.alignment), desc.C.alignment); | |
| } | |
| /// Returns the largest alignment (in units of elements) the problem satisfies, starting from a | |
| /// given upper limit. | |
| static int gemm_problem_alignment( | |
| int M, | |
| int N, | |
| int K, | |
| NumericTypeID element_A, | |
| void const *ptr_A, | |
| int64_t lda, | |
| int64_t batch_stride_A, | |
| NumericTypeID element_B, | |
| void const *ptr_B, | |
| int64_t ldb, | |
| int64_t batch_stride_B, | |
| NumericTypeID element_C, | |
| void const * ptr_C, | |
| int64_t ldc, | |
| int64_t batch_stride_C, | |
| void const * ptr_D, | |
| int64_t ldd, | |
| int64_t batch_stride_D, | |
| int max_alignment_in_bytes = 16 | |
| ) { | |
| void const *pointers[] = { | |
| ptr_A, ptr_B, ptr_C, ptr_D | |
| }; | |
| int64_t extents[] = { | |
| M, N, K, lda, ldb, ldc, ldd, batch_stride_A, batch_stride_B, batch_stride_C, batch_stride_D | |
| }; | |
| NumericTypeID elements[] = { | |
| element_A, element_B, element_C | |
| }; | |
| for (; max_alignment_in_bytes > 0; max_alignment_in_bytes /= 2) { | |
| bool satisfied = true; | |
| // Can pointers satisfy this? | |
| for (void const *ptr : pointers) { | |
| std::uintptr_t int_ptr = reinterpret_cast<std::uintptr_t>(ptr); | |
| if (int_ptr % max_alignment_in_bytes) { | |
| satisfied = false; | |
| break; | |
| } | |
| } | |
| if (!satisfied) { | |
| continue; | |
| } | |
| // Compute the maximum alignment based on element data types | |
| int max_element_alignment = 0; | |
| for (NumericTypeID type_id : elements) { | |
| int element_alignment = max_alignment_in_bytes * 8 / library::sizeof_bits(type_id); | |
| max_element_alignment = std::max(max_element_alignment, element_alignment); | |
| } | |
| // Can the problem size and leading dimensions satisfy this? | |
| for (int64_t extent : extents) { | |
| if (extent % max_element_alignment) { | |
| satisfied = false; | |
| break; | |
| } | |
| } | |
| if (!satisfied) { | |
| continue; | |
| } | |
| // Yes | |
| return max_element_alignment; | |
| } | |
| // No alignment satisfies this problem | |
| return 0; | |
| } | |
| /// Find the best kernel in descending order of preference. | |
| static Operation const * find_gemm_operation( | |
| GemmOperationFunctionalMap::const_iterator operators_it, | |
| GemmPreferenceKey const preference_key) { | |
| auto cc_it = operators_it->second.upper_bound(preference_key); | |
| if (cc_it == operators_it->second.begin()) { | |
| return nullptr; | |
| } | |
| Operation const *operation = nullptr; | |
| // Search in descending order of compute capability | |
| do { | |
| --cc_it; | |
| // Search tile sizes in order, for now. | |
| for (auto const * op : cc_it->second) { | |
| GemmDescription const &desc = static_cast<GemmDescription const &>(op->description()); | |
| int min_cc = desc.tile_description.minimum_compute_capability; | |
| int max_cc = desc.tile_description.maximum_compute_capability; | |
| int op_alignment = maximum_alignment_requirement(desc); | |
| if ((min_cc <= preference_key.compute_capability) && | |
| (preference_key.compute_capability <= max_cc) && | |
| (op_alignment <= preference_key.alignment)) { | |
| operation = op; | |
| break; | |
| } | |
| } | |
| } while (!operation && cc_it != operators_it->second.begin()); | |
| return operation; | |
| } | |
| /////////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Executes a GEMM computation: D <= alpha * A*B + beta * C | |
| Status Handle::gemm( | |
| int M, /// GEMM M dimension | |
| int N, /// GEMM N dimension | |
| int K, /// GEMM K dimension | |
| NumericTypeID element_compute, /// Data type of internal accumulation | |
| NumericTypeID element_scalar, /// Data type of alpha/beta scalars | |
| void const *alpha, /// Pointer to alpha scalar | |
| NumericTypeID element_A, /// Data type of A matrix elements | |
| LayoutTypeID layout_A, /// Layout of A matrix | |
| ComplexTransform transform_A, /// Complex transformation applied to A matrix - ignored for real-valued matrices | |
| void const * ptr_A, /// Pointer to A matrix in Global Memory | |
| int64_t lda, /// Leading dimension of A matrix | |
| NumericTypeID element_B, /// Data type of B matrix elements | |
| LayoutTypeID layout_B, /// Layout of B matrix | |
| ComplexTransform transform_B, /// Complex transformation applied to B matrix - ignored for real-valued matrices | |
| void const * ptr_B, /// Pointer to B matrix in Global Memory | |
| int64_t ldb, /// Leading dimension of B matrix | |
| void const * beta, /// Pointer to beta scalar | |
| NumericTypeID element_C, /// Data type of C and D matrices | |
| void const * ptr_C, /// Pointer to C matrix | |
| int64_t ldc, /// Leading dimension of C matrix | |
| void * ptr_D, /// Pointer to D matrix | |
| int64_t ldd /// Leading dimension of D matrix | |
| ) { | |
| // | |
| // Find the operation | |
| // | |
| GemmFunctionalKey key( | |
| provider_, | |
| GemmKind::kGemm, | |
| element_compute, | |
| element_scalar, | |
| element_A, | |
| layout_A, | |
| transform_A, | |
| element_B, | |
| layout_B, | |
| transform_B, | |
| element_C, // C/D are same type and col major default | |
| LayoutTypeID::kColumnMajor, | |
| element_C, | |
| LayoutTypeID::kColumnMajor | |
| ); | |
| auto operators_it = Singleton::get().operation_table.gemm_operations.find(key); | |
| if (operators_it == Singleton::get().operation_table.gemm_operations.end()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| if (operators_it->second.empty()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // | |
| // Compute the largest alignment restriction the kernel can satisfy. | |
| // | |
| // Maximum alignment expectation among all kernels (in units of bytes) | |
| int const kMaximumAlignmentSize = 16; | |
| int alignment = gemm_problem_alignment( | |
| M, N, K, | |
| element_A, ptr_A, lda, 0, | |
| element_B, ptr_B, ldb, 0, | |
| element_C, ptr_C, ldc, 0, | |
| ptr_D, ldd, 0, kMaximumAlignmentSize | |
| ); | |
| // | |
| // Find the best kernel in descending order of preference. | |
| // | |
| GemmPreferenceKey preference_key(compute_capability(), alignment); | |
| Operation const *operation = find_gemm_operation(operators_it, preference_key); | |
| if (!operation) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| last_operation_ = operation; | |
| // | |
| // Configure operation | |
| // | |
| GemmConfiguration configuration{ | |
| {M, N, K}, | |
| lda, | |
| ldb, | |
| ldc, | |
| ldd, | |
| 1 | |
| }; | |
| // Query host work space size | |
| uint64_t host_workspace_size_needed = operation->get_host_workspace_size(&configuration); | |
| if (uint64_t(kHostWorkspaceSize) < host_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| char host_workspace[kHostWorkspaceSize]; | |
| // Query device workspace size | |
| uint64_t device_workspace_size_needed = operation->get_device_workspace_size(&configuration); | |
| if (uint64_t(workspace_size_) < device_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // Initialize host and device workspaces | |
| Status status = operation->initialize( | |
| &configuration, | |
| host_workspace, | |
| workspace_, | |
| stream_); | |
| if (status != cutlass::Status::kSuccess) { | |
| return status; | |
| } | |
| // Run the operator | |
| GemmArguments arguments{ | |
| ptr_A, | |
| ptr_B, | |
| ptr_C, | |
| ptr_D, | |
| alpha, | |
| beta, | |
| scalar_pointer_mode_ | |
| }; | |
| return operation->run(&arguments, host_workspace, workspace_, stream_); | |
| } | |
| /////////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Executes a GEMM computation: D <= alpha * A*B + beta * C. | |
| // | |
| // Supports batched-strided, batched array or split-K serial or split-K parallel. | |
| // | |
| Status Handle::gemm_universal( | |
| GemmUniversalMode mode, /// indicates the mode in which the kUniversal GEMM is launched | |
| int M, /// GEMM M dimension | |
| int N, /// GEMM N dimension | |
| int K, /// GEMM K dimension | |
| NumericTypeID element_compute, /// Data type of internal accumulation | |
| NumericTypeID element_scalar, /// Data type of alpha/beta scalars | |
| void const *alpha, /// Pointer to alpha scalar | |
| NumericTypeID element_A, /// Data type of A matrix elements | |
| LayoutTypeID layout_A, /// Layout of A matrix | |
| ComplexTransform transform_A, /// Complex transformation applied to A matrix - ignored for real-valued matrices | |
| void const * ptr_A, /// Pointer to A matrix in Global Memory | |
| int64_t lda, /// Leading dimension of A matrix | |
| NumericTypeID element_B, /// Data type of B matrix elements | |
| LayoutTypeID layout_B, /// Layout of B matrix | |
| ComplexTransform transform_B, /// Complex transformation applied to B matrix - ignored for real-valued matrices | |
| void const * ptr_B, /// Pointer to B matrix in Global Memory | |
| int64_t ldb, /// Leading dimension of B matrix | |
| void const * beta, /// Pointer to beta scalar | |
| NumericTypeID element_C, /// Data type of C matrix | |
| LayoutTypeID layout_C, /// Layout of D matrix | |
| void const * ptr_C, /// Pointer to C matrix | |
| int64_t ldc, /// Leading dimension of C matrix | |
| NumericTypeID element_D, /// Data type of D matrix | |
| LayoutTypeID layout_D, /// Layout of D matrix | |
| void * ptr_D, /// Pointer to D matrix | |
| int64_t ldd, /// Leading dimension of D matrix | |
| int batch_count, /// Batch count or number of split-K slices | |
| int64_t batch_stride_A, /// Batch stride of A operand | |
| int64_t batch_stride_B, /// Batch stride of B operand | |
| int64_t batch_stride_C, /// Batch stride of C operand | |
| int64_t batch_stride_D /// Batch stride of D operand | |
| ) { | |
| // | |
| // Find the operation | |
| // | |
| GemmFunctionalKey key( | |
| provider_, | |
| GemmKind::kUniversal, | |
| element_compute, | |
| element_scalar, | |
| element_A, | |
| layout_A, | |
| transform_A, | |
| element_B, | |
| layout_B, | |
| transform_B, | |
| element_C, | |
| layout_C, | |
| element_D, | |
| layout_D | |
| ); | |
| auto operators_it = Singleton::get().operation_table.gemm_operations.find(key); | |
| if (operators_it == Singleton::get().operation_table.gemm_operations.end()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| if (operators_it->second.empty()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // | |
| // Compute the largest alignment restriction the kernel can satisfy. | |
| // | |
| // Maximum alignment expectation among all kernels (in units of bytes) | |
| int const kMaximumAlignmentSize = 16; | |
| void const *ptr_A_check = ptr_A; | |
| void const *ptr_B_check = ptr_B; | |
| void const *ptr_C_check = ptr_C; | |
| void * ptr_D_check = ptr_D; | |
| // Ignore alignment of pointers to pointers. We can't check this from the host, | |
| // as each batch index has its own pointer in device memory. | |
| if (mode == GemmUniversalMode::kArray) { | |
| ptr_A_check = nullptr; | |
| ptr_B_check = nullptr; | |
| ptr_C_check = nullptr; | |
| ptr_D_check = nullptr; | |
| } | |
| int alignment = gemm_problem_alignment( | |
| M, N, K, | |
| element_A, ptr_A_check, lda, 0, | |
| element_B, ptr_B_check, ldb, 0, | |
| element_C, ptr_C_check, ldc, 0, | |
| ptr_D_check, ldd, 0, kMaximumAlignmentSize | |
| ); | |
| // | |
| // Find the best kernel in descending order of preference. | |
| // | |
| GemmPreferenceKey preference_key(compute_capability(), alignment); | |
| Operation const *operation = find_gemm_operation(operators_it, preference_key); | |
| if (!operation) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| last_operation_ = operation; | |
| // | |
| // Configure operation | |
| // | |
| GemmUniversalConfiguration configuration{ | |
| mode, | |
| {M, N, K}, | |
| batch_count, | |
| lda, | |
| ldb, | |
| ldc, | |
| ldd | |
| }; | |
| // Query host work space size | |
| uint64_t host_workspace_size_needed = operation->get_host_workspace_size(&configuration); | |
| if (uint64_t(kHostWorkspaceSize) < host_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| char host_workspace[kHostWorkspaceSize]; | |
| GemmUniversalArguments arguments{ | |
| {M, N, K}, | |
| batch_count, | |
| ptr_A, | |
| ptr_B, | |
| ptr_C, | |
| ptr_D, | |
| alpha, | |
| beta, | |
| scalar_pointer_mode_, | |
| lda, | |
| ldb, | |
| ldc, | |
| ldd, | |
| batch_stride_A, | |
| batch_stride_B, | |
| batch_stride_C, | |
| batch_stride_D | |
| }; | |
| // Query device workspace size | |
| uint64_t device_workspace_size_needed = operation->get_device_workspace_size(&configuration, &arguments); | |
| if (uint64_t(workspace_size_) < device_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // Initialize host and device workspaces | |
| Status status = operation->initialize( | |
| &configuration, | |
| host_workspace, | |
| workspace_, | |
| stream_); | |
| if (status != cutlass::Status::kSuccess) { | |
| return status; | |
| } | |
| // Run the operator | |
| return operation->run(&arguments, host_workspace, workspace_, stream_); | |
| } | |
| /////////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Planar complex GEMM | |
| Status Handle::gemm_planar_complex( | |
| int M, /// GEMM M dimension | |
| int N, /// GEMM N dimension | |
| int K, /// GEMM K dimension | |
| NumericTypeID element_compute, /// Data type of internal accumulation | |
| NumericTypeID element_scalar, /// Data type of alpha/beta scalars | |
| void const *alpha, /// Pointer to alpha scalar | |
| NumericTypeID element_A, /// Data type of A matrix elements | |
| LayoutTypeID layout_A, /// Layout of A matrix | |
| ComplexTransform transform_A, /// Complex transformation applied to A matrix | |
| void const * ptr_A_real, /// Pointer to real part of A matrix | |
| void const * ptr_A_imag, /// Pointer to imaginary part of A matrix | |
| int64_t lda_real, /// Leading dimension of real part of A matrix | |
| int64_t lda_imag, /// Leading dimension of imaginary part of A matrix | |
| NumericTypeID element_B, /// Data type of B matrix elements | |
| LayoutTypeID layout_B, /// Layout of B matrix | |
| ComplexTransform transform_B, /// Complex transformation applied to B matrix | |
| void const * ptr_B_real, /// Pointer to real part of B matrix | |
| void const * ptr_B_imag, /// Pointer to imaginary part of B matrix | |
| int64_t ldb_real, /// Leading dimension of real part of B matrix | |
| int64_t ldb_imag, /// Leading dimension of imaginary part of B matrix | |
| void const * beta, /// Pointer to beta scalar | |
| NumericTypeID element_C, /// Data type of C and D matrix | |
| void const * ptr_C_real, /// Pointer to real part of C matrix | |
| void const * ptr_C_imag, /// Pointer to imaginary part of C matrix | |
| int64_t ldc_real, /// Leading dimension of real part of C matrix | |
| int64_t ldc_imag, /// Leading dimension of imaginary part of C matrix | |
| void * ptr_D_real, /// Pointer to real part of D matrix | |
| void * ptr_D_imag, /// Pointer to imaginary part of D matrix | |
| int64_t ldd_real, /// Leading dimension of real part of D matrix | |
| int64_t ldd_imag, /// Leading dimension of imaginary part of D matrix | |
| int batch_count, /// Number of batched GEMMs to execute | |
| int64_t batch_stride_A_real, | |
| int64_t batch_stride_A_imag, | |
| int64_t batch_stride_B_real, | |
| int64_t batch_stride_B_imag, | |
| int64_t batch_stride_C_real, | |
| int64_t batch_stride_C_imag, | |
| int64_t batch_stride_D_real, | |
| int64_t batch_stride_D_imag | |
| ) { | |
| // | |
| // Find the operation | |
| // | |
| GemmFunctionalKey key( | |
| provider_, | |
| GemmKind::kPlanarComplex, | |
| element_compute, | |
| element_scalar, | |
| element_A, | |
| layout_A, | |
| transform_A, | |
| element_B, | |
| layout_B, | |
| transform_B, | |
| element_C, // C/D are same type | |
| LayoutTypeID::kColumnMajor, | |
| element_C, | |
| LayoutTypeID::kColumnMajor | |
| ); | |
| auto operators_it = Singleton::get().operation_table.gemm_operations.find(key); | |
| if (operators_it == Singleton::get().operation_table.gemm_operations.end()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| if (operators_it->second.empty()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // | |
| // Compute the largest alignment restriction the kernel can satisfy. | |
| // | |
| // Maximum alignment expectation among all kernels (in units of bytes) | |
| int const kMaximumAlignmentSize = 16; | |
| int alignment = std::max( | |
| gemm_problem_alignment( | |
| M, N, K, | |
| element_A, ptr_A_real, lda_real, batch_stride_A_real, | |
| element_B, ptr_B_real, ldb_real, batch_stride_B_real, | |
| element_C, ptr_C_real, ldc_real, batch_stride_C_real, | |
| ptr_D_real, ldd_real, batch_stride_D_real, kMaximumAlignmentSize | |
| ), | |
| gemm_problem_alignment( | |
| M, N, K, | |
| element_A, ptr_A_imag, lda_imag, batch_stride_A_imag, | |
| element_B, ptr_B_imag, ldb_imag, batch_stride_B_imag, | |
| element_C, ptr_C_imag, ldc_imag, batch_stride_C_imag, | |
| ptr_D_imag, ldd_imag, batch_stride_D_imag, kMaximumAlignmentSize | |
| ) | |
| ); | |
| // | |
| // Find the best kernel in descending order of preference. | |
| // | |
| GemmPreferenceKey preference_key(compute_capability(), alignment); | |
| Operation const *operation = find_gemm_operation(operators_it, preference_key); | |
| if (!operation) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| last_operation_ = operation; | |
| // | |
| // Configure operation | |
| // | |
| GemmPlanarComplexConfiguration configuration{ | |
| GemmUniversalMode::kBatched, | |
| {M, N, K}, | |
| batch_count, | |
| lda_real, | |
| lda_imag, | |
| ldb_real, | |
| ldb_imag, | |
| ldc_real, | |
| ldc_imag, | |
| ldd_real, | |
| ldd_imag | |
| }; | |
| // Query host work space size | |
| uint64_t host_workspace_size_needed = operation->get_host_workspace_size(&configuration); | |
| if (uint64_t(kHostWorkspaceSize) < host_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| char host_workspace[kHostWorkspaceSize]; | |
| // Query device workspace size | |
| uint64_t device_workspace_size_needed = operation->get_device_workspace_size(&configuration); | |
| if (uint64_t(workspace_size_) < device_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // Initialize host and device workspaces | |
| Status status = operation->initialize( | |
| &configuration, | |
| host_workspace, | |
| workspace_, | |
| stream_); | |
| if (status != cutlass::Status::kSuccess) { | |
| return status; | |
| } | |
| // Run the operator | |
| GemmPlanarComplexArguments arguments{ | |
| ptr_A_real, | |
| ptr_A_imag, | |
| ptr_B_real, | |
| ptr_B_imag, | |
| ptr_C_real, | |
| ptr_C_imag, | |
| ptr_D_real, | |
| ptr_D_imag, | |
| alpha, | |
| beta, | |
| scalar_pointer_mode_, | |
| batch_stride_A_real, | |
| batch_stride_A_imag, | |
| batch_stride_B_real, | |
| batch_stride_B_imag, | |
| batch_stride_C_real, | |
| batch_stride_C_imag, | |
| batch_stride_D_real, | |
| batch_stride_D_imag | |
| }; | |
| return operation->run(&arguments, host_workspace, workspace_, stream_); | |
| } | |
| ///////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Planar complex batched GEMM loading pointers from arrays in global memory | |
| Status Handle::gemm_planar_complex_array( | |
| int expected_M, /// Expected GEMM M dimension (used for sizing CUDA grid) | |
| int expected_N, /// Expected GEMM N dimension (used for sizing CUDA grid) | |
| int expected_K, /// Expected GEMM K dimension | |
| int batch_count, /// Number of independent GEMM computations to execute | |
| int const *M, /// Array containing the GEMM M dimension for each batch index | |
| int const *N, /// Array containing the GEMM N dimension for each batch index | |
| int const *K, /// Array containing the GEMM K dimension for each batch index | |
| NumericTypeID element_compute, /// Data type of internal accumulation | |
| NumericTypeID element_scalar, /// Data type of alpha/beta scalars | |
| void const *alpha, /// Pointer to alpha scalar | |
| NumericTypeID element_A, /// Data type of A matrix elements | |
| LayoutTypeID layout_A, /// Layout of A matrix | |
| ComplexTransform transform_A, /// Complex transformation applied to A matrix | |
| void const * const * ptr_A_real, /// Pointer to array containing pointers to real part of A matrices | |
| void const * const * ptr_A_imag, /// Pointer to array containing pointers to imaginary part of A matrices | |
| int64_t lda_real, /// Leading dimension of real part of A matrix | |
| int64_t lda_imag, /// Leading dimension of imaginary part of A matrix | |
| NumericTypeID element_B, /// Data type of B matrix elements | |
| LayoutTypeID layout_B, /// Layout of B matrix | |
| ComplexTransform transform_B, /// Complex transformation applied to B matrix | |
| void const * const * ptr_B_real, /// Pointer to array containing pointers to real part of B matrices | |
| void const * const * ptr_B_imag, /// Pointer to array containing pointers to imaginary part of B matrices | |
| int64_t ldb_real, /// Leading dimension of real part of B matrix | |
| int64_t ldb_imag, /// Leading dimension of imaginary part of B matrix | |
| void const * beta, /// Pointer to beta scalar | |
| NumericTypeID element_C, /// Data type of C and D matrix | |
| void const * const * ptr_C_real, /// Pointer to array containing pointers to real part of C matrices | |
| void const * const * ptr_C_imag, /// Pointer to array containing pointers to imaginary part of C matrices | |
| int64_t ldc_real, /// Leading dimension of real part of C matrix | |
| int64_t ldc_imag, /// Leading dimension of imaginary part of C matrix | |
| void * const * ptr_D_real, /// Pointer to array containing pointers to real part of D matrices | |
| void * const * ptr_D_imag, /// Pointer to array containing pointers to imaginary part of D matrices | |
| int64_t ldd_real, /// Leading dimension of real part of D matrix | |
| int64_t ldd_imag /// Leading dimension of imaginary part of D matrix | |
| ) { | |
| // | |
| // Find the operation | |
| // | |
| GemmFunctionalKey key( | |
| provider_, | |
| GemmKind::kPlanarComplexArray, | |
| element_compute, | |
| element_scalar, | |
| element_A, | |
| layout_A, | |
| transform_A, | |
| element_B, | |
| layout_B, | |
| transform_B, | |
| element_C, // C/D are same type | |
| LayoutTypeID::kColumnMajor, | |
| element_C, | |
| LayoutTypeID::kColumnMajor | |
| ); | |
| auto operators_it = Singleton::get().operation_table.gemm_operations.find(key); | |
| if (operators_it == Singleton::get().operation_table.gemm_operations.end()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| if (operators_it->second.empty()) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // | |
| // Compute the largest alignment restriction the kernel can satisfy. | |
| // | |
| // Maximum alignment expectation among all kernels (in units of bytes) | |
| int const kMaximumAlignmentSize = 16; | |
| int alignment = std::max( | |
| gemm_problem_alignment( | |
| expected_M, expected_N, expected_K, | |
| element_A, nullptr, lda_real, 0, | |
| element_B, nullptr, ldb_real, 0, | |
| element_C, nullptr, ldc_real, 0, | |
| nullptr, ldd_real, 0, kMaximumAlignmentSize | |
| ), | |
| gemm_problem_alignment( | |
| expected_M, expected_N, expected_K, | |
| element_A, nullptr, lda_imag, 0, | |
| element_B, nullptr, ldb_imag, 0, | |
| element_C, nullptr, ldc_imag, 0, | |
| nullptr, ldd_imag, 0, kMaximumAlignmentSize | |
| ) | |
| ); | |
| // | |
| // Find the best kernel in descending order of preference. | |
| // | |
| GemmPreferenceKey preference_key(compute_capability(), alignment); | |
| Operation const *operation = find_gemm_operation(operators_it, preference_key); | |
| if (!operation) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| last_operation_ = operation; | |
| // | |
| // Configure operation | |
| // | |
| GemmPlanarComplexArrayConfiguration configuration{ | |
| {expected_M, expected_N, expected_K}, | |
| batch_count, | |
| lda_real, | |
| lda_imag, | |
| ldb_real, | |
| ldb_imag, | |
| ldc_real, | |
| ldc_imag, | |
| ldd_real, | |
| ldd_imag | |
| }; | |
| // Query host work space size | |
| uint64_t host_workspace_size_needed = operation->get_host_workspace_size(&configuration); | |
| if (uint64_t(kHostWorkspaceSize) < host_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| char host_workspace[kHostWorkspaceSize]; | |
| // Query device workspace size | |
| uint64_t device_workspace_size_needed = operation->get_device_workspace_size(&configuration); | |
| if (uint64_t(workspace_size_) < device_workspace_size_needed) { | |
| return cutlass::Status::kErrorNotSupported; | |
| } | |
| // Initialize host and device workspaces | |
| Status status = operation->initialize( | |
| &configuration, | |
| host_workspace, | |
| workspace_, | |
| stream_); | |
| if (status != cutlass::Status::kSuccess) { | |
| return status; | |
| } | |
| // Run the operator | |
| GemmPlanarComplexArrayArguments arguments{ | |
| M, N, K, | |
| ptr_A_real, | |
| ptr_A_imag, | |
| ptr_B_real, | |
| ptr_B_imag, | |
| ptr_C_real, | |
| ptr_C_imag, | |
| ptr_D_real, | |
| ptr_D_imag, | |
| alpha, | |
| beta, | |
| scalar_pointer_mode_ | |
| }; | |
| return operation->run(&arguments, host_workspace, workspace_, stream_); | |
| } | |
| ///////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Finds conv operation instances with Conv::ElementC = Reduction::ElementWorkspace | |
| Operation const* find_conv_operation_for_parallel_reduction(Operation const *operation) { | |
| ConvDescription const &conv_desc = | |
| static_cast<ConvDescription const &>(operation->description()); | |
| // if the curren conv operation accumulator and output data type match return operation | |
| if(conv_desc.tile_description.math_instruction.element_accumulator == conv_desc.C.element) { | |
| return operation; | |
| } | |
| // find conv operation to match conv output and reduction workspace data type | |
| ConvFunctionalKey key( | |
| library::Provider::kCUTLASS, | |
| conv_desc.conv_kind, | |
| conv_desc.A.element, | |
| conv_desc.A.layout, | |
| conv_desc.B.element, | |
| conv_desc.B.layout, | |
| conv_desc.tile_description.math_instruction.element_accumulator, | |
| conv_desc.C.layout, | |
| conv_desc.tile_description.math_instruction.element_accumulator, | |
| conv_desc.element_epilogue); | |
| // conv operation table for conv2d or conv3d | |
| auto conv_operations = (conv_desc.kind == OperationKind::kConv2d) ? | |
| Singleton::get().operation_table.conv2d_operations : | |
| Singleton::get().operation_table.conv3d_operations; | |
| // find ConvFunctionalKey in convolution operation table | |
| auto operators_it = conv_operations.find(key); | |
| if (operators_it == conv_operations.end()) { | |
| return nullptr; | |
| } | |
| if (operators_it->second.empty()) { | |
| return nullptr; | |
| } | |
| // conv operation for same compute capability and iterator algorithm | |
| ConvPreferenceKey preference_key( | |
| conv_desc.tile_description.minimum_compute_capability, | |
| conv_desc.iterator_algorithm); | |
| auto it = operators_it->second.find(preference_key); | |
| if(it == operators_it->second.end()) { | |
| return nullptr; | |
| } | |
| // return matching conv opertion (same tile sizes and instruction) | |
| for (auto op : it->second) { | |
| if (op->description().tile_description == operation->description().tile_description) { | |
| return op; | |
| } | |
| } | |
| return nullptr; | |
| } | |
| ///////////////////////////////////////////////////////////////////////////////////////////////// | |
| /// Finds gemm operation instances with Gemm::ElementC = Reduction::ElementWorkspace | |
| Operation const* find_gemm_operation_for_parallel_reduction(Operation const *operation) { | |
| GemmDescription const &gemm_desc = | |
| static_cast<GemmDescription const &>(operation->description()); | |
| // if the curren gemm operation accumulator and output data type match return operation | |
| if(gemm_desc.tile_description.math_instruction.element_accumulator == gemm_desc.D.element) { | |
| return operation; | |
| } | |
| // find gemm operation to match gemm output and reduction workspace data type | |
| GemmFunctionalKey key( | |
| library::Provider::kCUTLASS, | |
| gemm_desc.gemm_kind, | |
| gemm_desc.tile_description.math_instruction.element_accumulator, | |
| gemm_desc.element_epilogue, | |
| gemm_desc.A.element, | |
| gemm_desc.A.layout, | |
| gemm_desc.transform_A, | |
| gemm_desc.B.element, | |
| gemm_desc.B.layout, | |
| gemm_desc.transform_B, | |
| gemm_desc.tile_description.math_instruction.element_accumulator, // C/D are same type | |
| LayoutTypeID::kColumnMajor, | |
| gemm_desc.tile_description.math_instruction.element_accumulator, | |
| LayoutTypeID::kColumnMajor); | |
| // gemm operation table | |
| auto gemm_operations = Singleton::get().operation_table.gemm_operations; | |
| // find ConvFunctionalKey in gemm operation table | |
| auto operators_it = gemm_operations.find(key); | |
| if (operators_it == gemm_operations.end()) { | |
| return nullptr; | |
| } | |
| if (operators_it->second.empty()) { | |
| return nullptr; | |
| } | |
| // gemm operation for same compute capability and max operand alignment | |
| int alignment = std::max( | |
| gemm_desc.A.alignment, | |
| gemm_desc.B.alignment); | |
| GemmPreferenceKey preference_key( | |
| gemm_desc.tile_description.minimum_compute_capability, | |
| alignment); | |
| auto it = operators_it->second.find(preference_key); | |
| if(it == operators_it->second.end()) { | |
| return nullptr; | |
| } | |
| // return matching gemm opertion (same tile shape, stages, warp count, and instruction) | |
| for (auto op : it->second) { | |
| if (op->description().tile_description == operation->description().tile_description) { | |
| return op; | |
| } | |
| } | |
| // return nullptr if no matching gemm operation found for parallel split-k reduction | |
| return nullptr; | |
| } | |
| ///////////////////////////////////////////////////////////////////////////////////////////////// | |
| } // namespace library | |
| } // namespace cutlass | |
| ///////////////////////////////////////////////////////////////////////////////////////////////// | |