以下の内容はhttps://lpha-z.hatenablog.com/entry/2025/02/09/231500より取得しました。

cutlassの使い方の勉強（その1）

NVIDIAが開発している、cutlassというライブラリがあります。

github.com

cutlassは、GPUで線形代数演算をしたいときに役立つ外枠のライブラリです。 どう役に立つのか、もっとも単純な行列積を例にとって説明します。

行列積プログラムを書くことは非常に簡単（三重のfor文を書けばいいだけ）ですが、性能を追求しようと思うと途端に面倒になります。 例えば、キャッシュブロッキングや、GPUであればメモリアクセスがコアレスアクセスになるようにするなど、種々の技法があり、一から書くのはかなり面倒です。

ところで、それらの技法は、最内周のコードによらず、考え方は全く同じです。 ここで最内周とは、普通の単精度行列積ならc = fmaf( a, b, c )のことです。 例えば、整数の行列積（最内周はc = a * b + c）やトロピカル半環で考えた行列積（最内周はc = min( a + b, c )）であっても、外側を変える必要はありません。 もちろん、「最内周のコードによらず」と言っても、実際には、演算量やaやbのバイト数などによって適当にパラメータを変える必要はあるかもしれませんが、本質的なコードは一緒になります。

そこで、外枠だけをテンプレート化することに意義があります。 それらを抽象化したヘッダオンリーライブラリがcutlassです。

cutlassを勉強するにあたって、とりあえず、cuBLASを動かすと使われるカーネルであるcutlass_80_simt_sgemm_256x128_8x4_nn_align1を動かしてみることを目標にします。

解読

まず、cutlass_80_simt_sgemm_256x128_8x4_nn_align1に似た記述を含むファイルが存在しません。 しかしながら、一度ビルドすると出現するので、まずはビルドしてみます。

ビルド手順

cutlassは、ビルド手順難読化ツールのCMakeを使っているため、ビルドが困難です。

まず、nvccの場所を指定します。 環境を自動検出するのならパスが通っているものを使えばいいのに、なぜか/usr/bin/nvccを使おうとするので、CMakeにnvccの場所を教える必要があります。 公式ドキュメントにはexport CUDACXX=${CUDA_INSTALL_PATH}/bin/nvccを使えばよいと書かれていますが、これは大嘘で、何の効果もありません。 正しくは、cmake -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvccのように、cmakeのコマンドライン引数で渡します。

生成されたファイル

cutlass/build/tools/library/generated/gemm/80/sgemm/以下に、cutlass_simt_sgemm_256x128_8x4_nn_align1.cuというファイルが生成されます。

その中身は、以下のようになっています。

/*
  Generated by gemm_operation.py - Do not edit.
*/

///////////////////////////////////////////////////////////////////////////////////////////////////

#include "cutlass/cutlass.h"
#include "cutlass/library/library.h"
#include "cutlass/library/manifest.h"
#include "library_internal.h"
#include "gemm_operation.h"
#include "gemm_operation_3x.hpp"
#include "sparse_gemm_operation_3x.hpp"
#include "cutlass/arch/wmma.h"
#include "cutlass/numeric_types.h"
#include "cutlass/arch/arch.h"
#include "cutlass/arch/mma.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/gemm/device/gemm.h"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/default_gemm_universal.h"

///////////////////////////////////////////////////////////////////////////////////////////////////


// Gemm operator cutlass_simt_sgemm_256x128_8x4_nn_align1
using cutlass_simt_sgemm_256x128_8x4_nn_align1_base =
  typename cutlass::gemm::kernel::DefaultGemmUniversal<
    float, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1,    // transposed B operand
    float, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1,    // transposed A operand
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassSimt,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<256, 128, 8>,
    cutlass::gemm::GemmShape<64, 64, 8>,
    cutlass::gemm::GemmShape<1, 1, 1>,

    cutlass::epilogue::thread::LinearCombination<
      float,
      1,
      float,
      float
    >
,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>,
    4,
    cutlass::arch::OpMultiplyAdd
>::GemmKernel;

// Define named type
struct cutlass_simt_sgemm_256x128_8x4_nn_align1 :
  public cutlass_simt_sgemm_256x128_8x4_nn_align1_base { };


///////////////////////////////////////////////////////////////////////////////////////////////////

namespace cutlass {
namespace library {

///////////////////////////////////////////////////////////////////////////////////////////////////

void initialize_cutlass_simt_sgemm_256x128_8x4_nn_align1(Manifest &manifest) {



  manifest.append(new GemmUniversalOperation<
      cutlass::gemm::device::GemmUniversalAdapter<cutlass_simt_sgemm_256x128_8x4_nn_align1>
    >("cutlass_simt_sgemm_256x128_8x4_nn_align1"));



}

///////////////////////////////////////////////////////////////////////////////////////////////////

} // namespace library
} // namespace cutlass

///////////////////////////////////////////////////////////////////////////////////////////////////

※上記のコードは、NVIDIA CORPORATION & AFFILIATESがBSD-3-Clauseでライセンスしているcutlassライブラリに含まれるgemm_operation.pyが出力したコードです。そのライセンス全文は以下です。

Copyright (c) 2017 - 2025 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.

cutlass_simt_sgemm_256x128_8x4_nn_align1を使ってみたい

cutlass_simt_sgemm_256x128_8x4_nn_align1のありかは突き止めましたが、使い方がまるでわかりません。 引数をvoid*で受け取りつつ、さんざん引き回した先でstatic_castして取り出す、などの難読化が施されているようです。 CMakeといいvoid*といい、なぜそんなに難読化したいのか、理解に苦しみます。 テンプレート引数の制約がドキュメントされていないことも相まって、せっかくのテンプレートライブラリなのに、引数の型を特定することが著しく困難です。 なお、cutlass/media/docs/quickstart.md at main · NVIDIA/cutlass · GitHubにいくつか使い方の例が書かれていますが、DefaultGemmUniversalの使い方は書かれていません（他にもドキュメントはありますが、検索する限り、どこにも書かれていないようです）。

CUTLASS Profilerはヘッダオンリーではなく、ビルド方法は難読化されているため、必要な部分だけを取り出すということも難しいです。 幸い、CUTLASS Profilerは小さいので、ヘッダオンリーでない部分を手作業で外すことができました。 DeviceAllocationがcutlass::device_memory::DeviceAllocationではなくcutlass::profiler::DeviceAllocationを意味している、などの罠も仕掛けられていましたが、なんとかライブラリでの使い方をまねして動かせるコードを作ることができました。

以下のようなコードを作ることで、動かすことができるようです。

#include "gemm_operation.h"

// Gemm operator cutlass_simt_sgemm_256x128_8x4_nn_align1
using cutlass_simt_sgemm_256x128_8x4_nn_align1 =
  typename cutlass::gemm::kernel::DefaultGemmUniversal<
    float, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1,    // transposed B operand
    float, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1,    // transposed A operand
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassSimt,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<256, 128, 8>,
    cutlass::gemm::GemmShape<64, 64, 8>,
    cutlass::gemm::GemmShape<1, 1, 1>,

    cutlass::epilogue::thread::LinearCombination<
      float,
      1,
      float,
      float
    >
,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>,
    4,
    cutlass::arch::OpMultiplyAdd
>::GemmKernel;

/*
以上のコードは、NVIDIA CORPORATION & AFFILIATESがBSD-3-Clauseで
ライセンスしているcutlassライブラリに含まれるgemm_operation.pyが
出力したコードを改変したものです。そのライセンス全文は以下です。

Copyright (c) 2017 - 2025 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.
*/

#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/default_gemm_universal.h"

constexpr int N = 2;
float A[N*N] = { 1.0f, 2.5f, 6.0f, 7.0f };
float B[N*N] = { 1.0f, 2.0f, 3.0f, 4.0f };
float C[N*N] = { 1.0f, 2.0f, 3.0f, 4.0f };
float D[N*N] = { 100.f, 100.f, 100.f, 100.f };

int main() {
    cutlass::library::GemmUniversalOperation<cutlass::gemm::device::GemmUniversalAdapter<cutlass_simt_sgemm_256x128_8x4_nn_align1>> operation;

    cutlass::library::GemmUniversalConfiguration configuration;
    configuration.mode = cutlass::library::GemmUniversalMode::kGemm;
    configuration.problem_size.m() = N;
    configuration.problem_size.n() = N;
    configuration.problem_size.k() = N;
    configuration.lda = N;
    configuration.ldb = N;
    configuration.ldc = N;
    configuration.ldd = N;
    configuration.device_count = 1;

    std::vector<char> host_workspace(operation.get_host_workspace_size(&configuration)); // 368 Bytes

    float alpha = 1.0f;
    float beta = 1.0f;
    float* dA; cudaMalloc( &dA, N * N * sizeof(float) );
    float* dB; cudaMalloc( &dB, N * N * sizeof(float) );
    float* dC; cudaMalloc( &dC, N * N * sizeof(float) );
    float* dD; cudaMalloc( &dD, N * N * sizeof(float) );

    cutlass::library::GemmUniversalArguments arguments;    
    cudaMemcpy( dA, A, N * N * sizeof(float), cudaMemcpyHostToDevice );
    cudaMemcpy( dB, B, N * N * sizeof(float), cudaMemcpyHostToDevice );
    cudaMemcpy( dC, C, N * N * sizeof(float), cudaMemcpyHostToDevice );
    cudaMemcpy( dD, D, N * N * sizeof(float), cudaMemcpyHostToDevice );
    arguments.problem_size = cutlass::gemm::GemmCoord(N, N, N);
    arguments.batch_count = 1;
    arguments.A = dA;
    arguments.B = dB;
    arguments.C = dC;
    arguments.D = dD;
    arguments.alpha = &alpha;
    arguments.beta = &beta;
    arguments.pointer_mode = cutlass::library::ScalarPointerMode::kHost;
    arguments.lda = N;
    arguments.ldb = N;
    arguments.ldc = N;
    arguments.ldd = N;
    arguments.batch_stride_A = 0;
    arguments.batch_stride_B = 0;
    arguments.batch_stride_C = 0;
    arguments.batch_stride_D = 0;
    arguments.sm_count = 128;
    arguments.swizzle_size = 1;
    arguments.device_index = 0;

    std::vector<char> device_workspace(operation.get_device_workspace_size(&configuration, &arguments)); // 0 Bytes

    operation.initialize( &configuration, host_workspace.data(), device_workspace.data() );
    operation.run( &arguments, host_workspace.data() );

    cudaMemcpy( A, dA, N * N * sizeof(float), cudaMemcpyDeviceToHost );
    cudaMemcpy( B, dB, N * N * sizeof(float), cudaMemcpyDeviceToHost );
    cudaMemcpy( C, dC, N * N * sizeof(float), cudaMemcpyDeviceToHost );
    cudaMemcpy( D, dD, N * N * sizeof(float), cudaMemcpyDeviceToHost );

    // C^T += B^T A^T
    for( int j = 0; j < N; ++j )
        for( int k = 0; k < N; ++k )
            for( int i = 0; i < N; ++i )
                C[j*N+i] += B[j*N+k] * A[k*N+i];

    for( int i = 0; i < N; ++i )
        for( int j = 0; j < N; ++j )
            assert( C[i*N+j] == D[i*N+j] );
}

コンパイルオプションは、nvcc main.cu -std=c++17 --expt-relaxed-constexpr -arch=compute_89 -code=sm_89 -I /path/to/cutlass/tools/library/src/ -I /path/to/cutlass/include/ -I /path/to/cutlass/tools/library/include/です。

cuBLASとの比較

これをN=4096で実行すると、55.2 TFLOPSくらいの性能となり、cuBLASが60 TFLOPS程度であることを考えれば、悪くない数字です。 なんだかグローバルメモリからシェアードメモリに転送するところでパディングを入れつつ転置している（cutlass/gemm/threadblock/default_mma_core_simt.h 370～371行目）みたいで、プロファイラで見てみるとここでシェアードメモリのバンクコンフリクトが起こっているようですがいいのでしょうか。

cuBLASとの差は、仕事の分割を行っているかにありそうです。 cuBLASは256スレッドのブロックを256×2×3個生成するのに対し、このコードは256スレッドのブロックを128×4個生成します。 つまり、cuBLASは仕事を3分割しています。 cutlass/media/docs/efficient_gemm.md at main · NVIDIA/cutlass · GitHubを見ると、"Split K"や"Sliced K"という分割手法があると書かれているため、これを使っているのでしょう。 例えば、DefaultGemmSplitKParallelというテンプレートがライブラリに用意されています。

これの使い方もまたドキュメントが全くなくてよくわからないですが、また今度挑戦してみようと思います。

まとめ

cutlassの難読化を解除し、cutlass_simt_sgemm_256x128_8x4_nn_align1を動かすことのできる短いコードを作ることができました。 単純にそれを動かすだけで55.2 TFLOPSの性能が出るなど、cutlassの威力を知ることができました。

cutlass_simt_sgemm_256x128_8x4_nn_align1の中で何が行われているかや、cuBLASとの違いと思われるSplit Kの導入方法はよくわからなかったので、また今度挑戦してみます。



以上の内容はhttps://lpha-z.hatenablog.com/entry/2025/02/09/231500より取得しました。
このページはhttp://font.textar.tv/のウェブフォントを使用してます
不具合報告/要望等はこちらへお願いします。
モバイルやる夫Viewer Ver0.14