wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:40:42 +08:00
commit e25996e7db
15472 changed files with 3536181 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/paddle/apy/matmul_pass/matmul/cutlass_patch/backend.h
+24
View File
@@ -0,0 +1,24 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#ifdef __NVCC__
#include "cutlass/cutlass.h"
#define GPUStream_t cudaStream_t
#elif defined(__HIPCC__)
#include "hytlass/hytlass.h"
namespace cutlass = hytlass;
#define GPUStream_t hipStream_t
#endif
python/paddle/apy/matmul_pass/matmul/cutlass_patch/batched_matrix_coord.h
+33
View File
@@ -0,0 +1,33 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "cutlass_patch/backend.h"
namespace cutlass_patch {
struct BatchedMatrixCoord {
int batch;
int row;
int column;
CUTLASS_HOST_DEVICE
BatchedMatrixCoord() : batch(0), row(0), column(0) {}
CUTLASS_HOST_DEVICE
BatchedMatrixCoord(int b, int r, int c) : batch(b), row(r), column(c) {}
};
}; // namespace cutlass_patch
python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/all_tuning_configs.h
+496
View File
@@ -0,0 +1,496 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// auto-generated by generate_configs.py
#pragma once
#include "cutlass/gemm_coord.h"
namespace ap {
constexpr int kNumConfigsHalf = 23;
constexpr int kNumConfigsFloat = 13;
template <int SwizzleFactor, bool Batched>
struct SwizzleWrapper {
using Type =
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
};
// template <int SwizzleFactor>
// struct SwizzleWrapper<SwizzleFactor, true> {
// using Type =
// cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
// };
#define AP_AUTOTUNE(func, stream_ptr, count, ...) \
{ \
using FuncType = decltype(func<0>); \
static int selected_config_id = -1; \
static std::vector<std::function<FuncType>> matmul_functions = \
[]<std::size_t... Is>(std::index_sequence<Is...>) { \
return std::vector<std::function<FuncType>>{func<Is>...}; \
} \
(std::make_index_sequence<count>()); \
\
if (selected_config_id == -1) { \
selected_config_id = \
ap::ProfileBestConfig(matmul_functions, stream_ptr, ##__VA_ARGS__); \
} \
\
matmul_functions[selected_config_id](__VA_ARGS__); \
}
#define AP_AUTOTUNE_half(func, stream_ptr, ...) \
AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsHalf, __VA_ARGS__)
#define AP_AUTOTUNE_float(func, stream_ptr, ...) \
AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsFloat, __VA_ARGS__)
#define AP_AUTOTUNE_bfloat16(func, stream_ptr, ...) \
AP_AUTOTUNE_half(func, stream_ptr, __VA_ARGS__)
template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
struct GemmTuningConfigs {
using TShape = cutlass::gemm::GemmShape<256, 128, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 2;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = Id;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 1> {
using TShape = cutlass::gemm::GemmShape<64, 128, 64>;
using WShape = cutlass::gemm::GemmShape<32, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 1;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 2> {
using TShape = cutlass::gemm::GemmShape<64, 128, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 2;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 3> {
using TShape = cutlass::gemm::GemmShape<128, 64, 64>;
using WShape = cutlass::gemm::GemmShape<64, 32, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 3;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 4> {
using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 4;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 5> {
using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 5;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 6> {
using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 6;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 7> {
using TShape = cutlass::gemm::GemmShape<256, 64, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 7;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 8> {
using TShape = cutlass::gemm::GemmShape<256, 128, 32>;
using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 8;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 9> {
using TShape = cutlass::gemm::GemmShape<256, 128, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 9;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 10> {
using TShape = cutlass::gemm::GemmShape<128, 32, 64>;
using WShape = cutlass::gemm::GemmShape<32, 32, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 10;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 11> {
using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 11;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 12> {
using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 12;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 13> {
using TShape = cutlass::gemm::GemmShape<256, 64, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 13;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 14> {
using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 14;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 15> {
using TShape = cutlass::gemm::GemmShape<32, 64, 64>;
using WShape = cutlass::gemm::GemmShape<16, 32, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 15;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 16> {
using TShape = cutlass::gemm::GemmShape<64, 64, 64>;
using WShape = cutlass::gemm::GemmShape<32, 32, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 16;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 17> {
using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 17;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 18> {
using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 18;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 19> {
using TShape = cutlass::gemm::GemmShape<64, 128, 32>;
using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 6;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 19;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 20> {
using TShape = cutlass::gemm::GemmShape<128, 64, 32>;
using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 6;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 20;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 21> {
using TShape = cutlass::gemm::GemmShape<128, 32, 32>;
using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 7;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 21;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 22> {
using TShape = cutlass::gemm::GemmShape<64, 64, 32>;
using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
static constexpr int kNumStages = 10;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 22;
};
// Specialization for float
template <int SwizzleFactor, bool Batched, int Id>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
using TShape = cutlass::gemm::GemmShape<64, 64, 16>;
using WShape = cutlass::gemm::GemmShape<32, 32, 16>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = Id;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 1> {
using TShape = cutlass::gemm::GemmShape<64, 64, 32>;
using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 1;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 2> {
using TShape = cutlass::gemm::GemmShape<64, 128, 32>;
using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 2;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 3> {
using TShape = cutlass::gemm::GemmShape<64, 256, 16>;
using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 3;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 4> {
using TShape = cutlass::gemm::GemmShape<64, 256, 32>;
using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 4;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 5> {
using TShape = cutlass::gemm::GemmShape<128, 64, 32>;
using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 5;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 6> {
using TShape = cutlass::gemm::GemmShape<128, 128, 16>;
using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 6;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 7> {
using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 7;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 8> {
using TShape = cutlass::gemm::GemmShape<256, 64, 16>;
using WShape = cutlass::gemm::GemmShape<64, 32, 16>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 8;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 9> {
using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 9;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 10> {
using TShape = cutlass::gemm::GemmShape<64, 128, 16>;
using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 10;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 11> {
using TShape = cutlass::gemm::GemmShape<128, 64, 16>;
using WShape = cutlass::gemm::GemmShape<64, 32, 16>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 11;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 12> {
using TShape = cutlass::gemm::GemmShape<128, 128, 16>;
using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 12;
};
} // namespace ap
python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/cutlass_matmul.cuh
+273
View File
@@ -0,0 +1,273 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include "cutlass/cutlass.h"
#include "cutlass/gemm_coord.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h"
#include "cutlass/util/device_memory.h"
#include "cutlass/gemm/device/gemm_universal.h"
#include "cutlass/gemm/device/gemm_universal_with_broadcast.h"
#include "cutlass_patch/batched_matrix_coord.h"
#include "cutlass_patch/cuda/default_config_id.h"
#include "cutlass_patch/epilogue/thread/linear_combination_unary.h"
#include "cutlass_patch/epilogue/thread/linear_combination_variadic.h"
#include "cutlass_patch/gemm/device/gemm_universal_with_variadic.h"
#include "params.h" // NOLINT
#define CHECK_CUTLASS(status) \
{ \
cutlass::Status error = status; \
if (error != cutlass::Status::kSuccess) { \
std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) \
<< " at: " << __LINE__ << std::endl; \
exit(EXIT_FAILURE); \
} \
}
namespace ap {
using bfloat16 = nv_bfloat16;
template <typename T, int N>
using Array = cutlass::Array<T, N>;
using MatrixCoord = cutlass_patch::BatchedMatrixCoord;
// Convert CUDA data type to cutlass data type
template <typename T>
struct CutlassDataType {
using Type = T;
};
template <>
struct CutlassDataType<half> {
using Type = cutlass::half_t;
};
template <>
struct CutlassDataType<__nv_bfloat16> {
using Type = cutlass::bfloat16_t;
};
// Convert to cutlass layout
template <bool Transposed>
struct MatrixLayout {
using Type = cutlass::layout::RowMajor;
};
template <>
struct MatrixLayout<true> {
using Type = cutlass::layout::ColumnMajor;
};
// Operation performed by GEMM
template <typename ElementT>
struct GemmOperation {
using Type = cutlass::arch::OpMultiplyAdd;
};
template <>
struct GemmOperation<float> {
using Type = cutlass::arch::OpMultiplyAddFastF32;
};
static cutlass::gemm::GemmUniversalMode GetGemmMode(int batch_count) {
return batch_count > 1 ? cutlass::gemm::GemmUniversalMode::kBatched
: cutlass::gemm::GemmUniversalMode::kGemm;
}
static void *GetWorkspace(size_t workspace_size) {
static cutlass::device_memory::allocation<uint8_t> workspace;
if (workspace.size() < workspace_size) {
workspace.reset(workspace_size);
}
return workspace.get();
}
template <typename GemmFunc>
cutlass::Status SetMaxDynamicSharedMemorySize() {
cudaError_t cudart_result;
// If requires more than 48KB: configure for extended, dynamic shared memory
if constexpr (GemmFunc::kSharedStorageSize >= (48 << 10)) {
cudart_result =
cudaFuncSetAttribute(cutlass::Kernel2<typename GemmFunc::GemmKernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
GemmFunc::kSharedStorageSize);
if (cudart_result != cudaSuccess) {
CUTLASS_TRACE_HOST("cudaFuncSetAttribute() returned error "
<< cudaGetErrorString(cudart_result));
return cutlass::Status::kErrorInternal;
}
}
#if AP_ENABLE_DEBUG
// Update SM occupancy member
int sm_occupancy = -1;
cudart_result = cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
&sm_occupancy,
cutlass::Kernel2<typename GemmFunc::GemmKernel>,
GemmFunc::GemmKernel::kThreadCount,
GemmFunc::kSharedStorageSize,
cudaOccupancyDisableCachingOverride);
if (cudart_result != cudaSuccess) {
CUTLASS_TRACE_HOST(
"cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned "
"error "
<< cudaGetErrorString(cudart_result));
return cutlass::Status::kErrorInternal;
}
CUTLASS_TRACE_HOST("sm_occupancy: (" << sm_occupancy
<< ") "
"smem_size: ("
<< GemmFunc::kSharedStorageSize
<< ") "
"GemmKernel::kThreadCount: ("
<< GemmFunc::GemmKernel::kThreadCount
<< ")");
#endif
return cutlass::Status::kSuccess;
}
template <typename ElementT,
typename ElementComputeT,
template <typename T>
class VariadicFunctor,
int AlignA = 128 / cutlass::sizeof_bits<ElementT>::value,
int AlignB = 128 / cutlass::sizeof_bits<ElementT>::value,
int ConfigId = DefaultConfig::kConfigId,
int SwizzleFactor = DefaultConfig::kSwizzleFactor,
bool Batched = DefaultConfig::kBatched>
void MatmulAddVariadic(
const GemmEpilogueParams &params,
const typename VariadicFunctor<ElementComputeT>::Arguments &variadic_args) {
using ElementAccumulator =
typename CutlassDataType<ElementComputeT>::Type; // <- data type of
// accumulator
using ElementComputeEpilogue =
ElementAccumulator; // <- data type of epilogue operations
using ElementInputA =
typename CutlassDataType<ElementT>::Type; // <- data type of elements in
// input matrix A
using ElementInputB =
typename CutlassDataType<ElementT>::Type; // <- data type of elements in
// input matrix B
using ElementOutput =
typename CutlassDataType<ElementT>::Type; // <- data type of elements in
// output matrix D
constexpr int AlignC = AlignB;
// Epilogue operation as LinearCombination:
// alpha * accumulator + beta * source
using EpilogueOutputOp =
cutlass_patch::epilogue::thread::LinearCombinationVariadic<
VariadicFunctor,
ElementOutput,
AlignC,
ElementAccumulator,
ElementComputeEpilogue,
cutlass::epilogue::thread::ScaleType::NoBetaScaling>; // <- alpha x
// AB + bias
using GemmFunc = cutlass_patch::gemm::device::GemmUniversalWithVariadic<
ElementInputA,
cutlass::layout::RowMajor,
ElementInputB,
cutlass::layout::RowMajor,
ElementOutput,
cutlass::layout::RowMajor,
ElementAccumulator,
cutlass::arch::OpClassTensorOp,
cutlass::arch::Sm80,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
TShape,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
WShape,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
IShape,
EpilogueOutputOp,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
SwizzleThreadBlock,
GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::kNumStages,
AlignA,
AlignB,
typename GemmOperation<ElementT>::Type>;
CHECK_CUTLASS(SetMaxDynamicSharedMemorySize<GemmFunc>());
/// Arguments
cutlass::gemm::GemmCoord problem_size{params.m, params.n, params.k};
const ElementInputA *input =
reinterpret_cast<const ElementInputA *>(params.input);
const ElementInputB *weight =
reinterpret_cast<const ElementInputB *>(params.weight);
const ElementOutput *bias =
reinterpret_cast<const ElementOutput *>(params.bias);
ElementOutput *output = reinterpret_cast<ElementOutput *>(params.output);
ElementComputeEpilogue alpha = static_cast<ElementComputeEpilogue>(1);
ElementComputeEpilogue beta = bias ? static_cast<ElementComputeEpilogue>(1)
: static_cast<ElementComputeEpilogue>(0);
typename GemmFunc::Arguments arguments{
GetGemmMode(params.batch_count),
problem_size, // <- problem size of matrix multiplication
params.batch_count, // <- batch_count or k-dimension split factor
{alpha, beta, variadic_args}, // <- epilogue params, alpha, beta
input, // <- input, ptr_A, A, shape={M, K}
weight, // <- input, ptr_B, B, shape={K, N}
bias, // <- input, ptr_C, shape={M, N} or {1, N}
output, // <- output, ptr_D, Z, shape={M, N}
params.shape_args.batch_stride_A,
params.shape_args.batch_stride_B,
params.shape_args.batch_stride_C,
params.shape_args.batch_stride_D,
params.shape_args.lda,
params.shape_args.ldb,
params.shape_args.ldc_bias,
params.shape_args.ldd};
size_t workspace_size = GemmFunc::get_workspace_size(arguments);
void *workspace = workspace_size > 0 ? GetWorkspace(workspace_size) : nullptr;
GemmFunc device_gemm;
cudaStream_t *stream_ptr =
reinterpret_cast<cudaStream_t *>(params.stream_ptr);
CHECK_CUTLASS(device_gemm.can_implement(arguments));
CHECK_CUTLASS(device_gemm.initialize(arguments, workspace, *stream_ptr));
//
// Run the GEMM
//
CHECK_CUTLASS(device_gemm(*stream_ptr));
#if AP_ENABLE_DEBUG
CHECK_CUDA(cudaStreamSynchronize(*stream_ptr));
#endif
}
} // namespace ap
python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/default_config_id.h
+27
View File
@@ -0,0 +1,27 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "all_tuning_configs.h" // NOLINT
namespace ap {
struct DefaultConfig {
static constexpr int kConfigId = 0;
static constexpr int kSwizzleFactor = 1;
static constexpr bool kBatched = false;
};
} // namespace ap
python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/thread/linear_combination_unary.h
+302
View File
@@ -0,0 +1,302 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*! \file
\brief Functor performing linear combination operations used by epilogues.
It is modified from LinearCombinationGeneric.
*/
#pragma once
#include "cutlass_patch/backend.h"
#ifdef __NVCC__
#include "cutlass/array.h"
#include "cutlass/epilogue/thread/scale_type.h"
#include "cutlass/functional.h"
#include "cutlass/numeric_conversion.h"
#include "cutlass/numeric_types.h"
#elif defined(__HIPCC__)
#include "hytlass/array.h"
#include "hytlass/epilogue/thread/scale_type.h"
#include "hytlass/functional.h"
#include "hytlass/numeric_conversion.h"
#include "hytlass/numeric_types.h"
#endif
namespace cutlass_patch {
namespace epilogue {
namespace thread {
template <class UnaryOp, class = void>
struct GenericUnaryTraits {
static constexpr bool IsArgumentsNeeded = false;
struct Arguments {};
};
template <class UnaryOp>
struct GenericUnaryTraits<UnaryOp,
decltype(typename UnaryOp::Arguments(), void())> {
static constexpr bool IsArgumentsNeeded = true;
using Arguments = typename UnaryOp::Arguments;
};
/// Applies a linear combination operator followed by an unary function to an
/// array of elements.
///
/// D = unary_op(alpha * accumulator + beta * source)
///
template <
template <typename T>
class UnaryOp,
typename ElementOutput_, ///< Data type used to load and store tensors
int ElementsPerAccess, ///< Number of elements computed per operation
///< Usually it is 128/sizeof_bits<ElementOutput_>,
///< but we use 64 or 32 sometimes when there are
///< not enough data to store
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
typename ElementCompute_ =
ElementOutput_, ///< Data type used to compute linear combination
cutlass::epilogue::thread::ScaleType::Kind Scale =
cutlass::epilogue::thread::ScaleType::Default, ///< Control Alpha and
///< Beta scaling
cutlass::FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest,
bool IsHeavy = false>
class LinearCombinationUnary {
public:
using ElementOutput = ElementOutput_;
using ElementAccumulator = ElementAccumulator_;
using ElementCompute = ElementCompute_;
using UnaryArguments =
typename GenericUnaryTraits<UnaryOp<ElementCompute>>::Arguments;
static bool const kIsHeavy = IsHeavy;
static int const kElementsPerAccess = ElementsPerAccess;
static int const kCount = ElementsPerAccess;
static const cutlass::epilogue::thread::ScaleType::Kind kScale = Scale;
using FragmentOutput = cutlass::Array<ElementOutput, kElementsPerAccess>;
using FragmentAccumulator =
cutlass::Array<ElementAccumulator, kElementsPerAccess>;
using FragmentSource = cutlass::Array<ElementOutput, kElementsPerAccess>;
using FragmentCompute = cutlass::Array<ElementCompute, kElementsPerAccess>;
static cutlass::FloatRoundStyle const kRound = Round;
/// Host-constructable parameters structure
struct Params {
ElementCompute alpha; ///< scales accumulators
ElementCompute beta; ///< scales source tensor
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if
///< not null, loads it from memory
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not
///< null, loads it from memory
UnaryArguments unary_args;
//
// Methods
//
CUTLASS_HOST_DEVICE
Params()
: alpha(ElementCompute(1)),
beta(ElementCompute(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) {}
CUTLASS_HOST_DEVICE
Params(ElementCompute alpha,
ElementCompute beta = ElementCompute(0),
UnaryArguments unary_args_ = UnaryArguments{})
: alpha(alpha),
beta(beta),
alpha_ptr(nullptr),
beta_ptr(nullptr),
unary_args(unary_args_) {}
CUTLASS_HOST_DEVICE
Params(ElementCompute const *alpha_ptr,
ElementCompute const *beta_ptr = nullptr)
: alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {}
};
private:
//
// Data members
//
Params params_;
bool skip_elementwise_;
public:
/// Constructs the function object, possibly loading from pointers in host
/// memory
CUTLASS_HOST_DEVICE
explicit LinearCombinationUnary(Params const &params) {
params_ = params;
params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta);
skip_elementwise_ = false;
}
/// Returns true if source is needed
CUTLASS_HOST_DEVICE
bool is_source_needed() const {
if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling)
return params_.beta != ElementCompute(0);
if (Scale == cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling)
return false;
if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) return false;
return params_.beta != ElementCompute(0);
}
/// Functionally required for serial reduction in the epilogue
CUTLASS_HOST_DEVICE
void set_k_partition(int k_partition, int k_partition_count) {
if (k_partition) {
params_.beta = ElementCompute(1);
}
if (k_partition != k_partition_count - 1) {
skip_elementwise_ = true;
}
}
/// Computes linear scaling: D = alpha * accumulator + beta * source
CUTLASS_HOST_DEVICE
FragmentOutput operator()(FragmentAccumulator const &accumulator,
FragmentOutput const &source) const {
// Convert source to internal compute numeric type
cutlass::NumericArrayConverter<ElementCompute,
ElementOutput,
kElementsPerAccess,
Round>
source_converter;
cutlass::NumericArrayConverter<ElementCompute,
ElementAccumulator,
kElementsPerAccess,
Round>
accumulator_converter;
FragmentCompute converted_source = source_converter(source);
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
// Perform binary operations
FragmentCompute intermediate;
cutlass::multiplies<FragmentCompute> mul_add_source;
cutlass::multiply_add<FragmentCompute> mul_add_accumulator;
UnaryOp<ElementCompute> unary_op;
if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling) {
intermediate = converted_source;
// D = alpha * Accum + X
intermediate = mul_add_accumulator(
params_.alpha, converted_accumulator, intermediate);
} else if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
intermediate = converted_accumulator;
} else {
// X = beta * C + uniform
intermediate = mul_add_source(params_.beta, converted_source);
// D = alpha * Accum + X
intermediate = mul_add_accumulator(
params_.alpha, converted_accumulator, intermediate);
}
if constexpr (GenericUnaryTraits<
UnaryOp<ElementCompute>>::IsArgumentsNeeded) {
if (!skip_elementwise_) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = unary_op(intermediate[i], params_.unary_args);
}
}
} else {
if (!skip_elementwise_) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = unary_op(intermediate[i]);
}
}
}
// Convert to destination numeric type
cutlass::NumericArrayConverter<ElementOutput,
ElementCompute,
kElementsPerAccess,
Round>
destination_converter;
return destination_converter(intermediate);
}
/// Computes linear scaling: D = alpha * accumulator
CUTLASS_HOST_DEVICE
FragmentOutput operator()(FragmentAccumulator const &accumulator) const {
// Convert source to internal compute numeric type
cutlass::NumericArrayConverter<ElementCompute,
ElementAccumulator,
kElementsPerAccess,
Round>
accumulator_converter;
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
// Perform binary operations
FragmentCompute intermediate;
cutlass::multiplies<FragmentCompute> mul_add_accumulator;
UnaryOp<ElementCompute> unary_op;
if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
intermediate = converted_accumulator;
} else {
// D = alpha * Accum
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator);
}
if constexpr (GenericUnaryTraits<
UnaryOp<FragmentCompute>>::IsArgumentsNeeded) {
if (!skip_elementwise_) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = unary_op(intermediate[i], params_.unary_args);
}
}
} else {
if (!skip_elementwise_) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = unary_op(intermediate[i]);
}
}
}
// Convert to destination numeric type
cutlass::NumericArrayConverter<ElementOutput,
ElementCompute,
kElementsPerAccess,
Round>
destination_converter;
return destination_converter(intermediate);
}
};
} // namespace thread
} // namespace epilogue
} // namespace cutlass_patch
python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/thread/linear_combination_variadic.h
+337
View File
@@ -0,0 +1,337 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*! \file
\brief Functor performing linear combination operations used by epilogues.
*/
#pragma once
#include "cutlass_patch/backend.h"
#ifdef __NVCC__
#include "cutlass/array.h"
#include "cutlass/epilogue/thread/scale_type.h"
#include "cutlass/functional.h"
#include "cutlass/numeric_conversion.h"
#include "cutlass/numeric_types.h"
#elif defined(__HIPCC__)
#include "hytlass/array.h"
#include "hytlass/epilogue/thread/scale_type.h"
#include "hytlass/functional.h"
#include "hytlass/numeric_conversion.h"
#include "hytlass/numeric_types.h"
#endif
#include "cutlass_patch/batched_matrix_coord.h"
#include "cutlass_patch/trace_device.h"
namespace cutlass_patch {
namespace epilogue {
namespace thread {
template <class VariadicOp, class = void>
struct GenericVariadicTraits {
static constexpr bool IsArgumentsNeeded = false;
struct Arguments {};
};
template <class VariadicOp>
struct GenericVariadicTraits<VariadicOp,
decltype(typename VariadicOp::Arguments(),
void())> {
static constexpr bool IsArgumentsNeeded = true;
using Arguments = typename VariadicOp::Arguments;
};
/// Applies a linear combination operator to an array of elements.
///
/// D = VariadicOp(alpha * accumulator + beta * source)
///
template <
template <typename T>
class VariadicOp,
typename ElementOutput_, ///< Data type used to load and store tensors
int ElementsPerAccess, ///< Number of elements computed per operation.
///< Usually it is 128/sizeof_bits<ElementOutput_>,
///< but we use 64 or 32 sometimes when there are
///< not enough data to store
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
typename ElementCompute_ =
ElementOutput_, ///< Data type used to compute linear combination
cutlass::epilogue::thread::ScaleType::Kind Scale =
cutlass::epilogue::thread::ScaleType::Default, ///< Control Alpha and
///< Beta scaling
cutlass::FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest,
bool IsHeavy = false>
class LinearCombinationVariadic {
public:
using ElementOutput = ElementOutput_;
using ElementAccumulator = ElementAccumulator_;
using ElementCompute = ElementCompute_;
using VariadicArguments =
typename GenericVariadicTraits<VariadicOp<ElementCompute>>::Arguments;
static bool const kIsHeavy = IsHeavy;
static int const kElementsPerAccess = ElementsPerAccess;
static int const kCount = ElementsPerAccess;
static const cutlass::epilogue::thread::ScaleType::Kind kScale = Scale;
using FragmentOutput = cutlass::Array<ElementOutput, kElementsPerAccess>;
using FragmentAccumulator =
cutlass::Array<ElementAccumulator, kElementsPerAccess>;
using FragmentSource = cutlass::Array<ElementOutput, kElementsPerAccess>;
using FragmentCompute = cutlass::Array<ElementCompute, kElementsPerAccess>;
static cutlass::FloatRoundStyle const kRound = Round;
/// Host-constructable parameters structure
struct Params {
ElementCompute alpha; ///< scales accumulators
ElementCompute beta; ///< scales source tensor
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if
///< not null, loads it from memory
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not
///< null, loads it from memory
VariadicArguments variadic_args;
CUTLASS_HOST_DEVICE
Params()
: alpha(ElementCompute(1)),
beta(ElementCompute(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) {}
CUTLASS_HOST_DEVICE
Params(ElementCompute alpha,
ElementCompute beta,
VariadicArguments variadic_args_ = VariadicArguments{})
: alpha(alpha),
beta(beta),
alpha_ptr(nullptr),
beta_ptr(nullptr),
variadic_args(variadic_args_) {}
};
private:
//
// Data members
//
Params params_;
bool skip_elementwise_;
public:
/// Constructs the function object, possibly loading from pointers in host
/// memory
CUTLASS_HOST_DEVICE
LinearCombinationVariadic(Params const &params) {
params_ = params;
params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta);
skip_elementwise_ = false;
}
/// Returns true if source is needed
CUTLASS_HOST_DEVICE
bool is_source_needed() const {
if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling)
return params_.beta != ElementCompute(0);
if (Scale == cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling)
return false;
if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) return false;
return params_.beta != ElementCompute(0);
}
/// Functionally required for serial reduction in the epilogue
CUTLASS_HOST_DEVICE
void set_k_partition(int k_partition, int k_partition_count) {
if (k_partition) {
params_.beta = ElementCompute(1);
}
if (k_partition != k_partition_count - 1) {
skip_elementwise_ = true;
}
}
/// Computes linear scaling with source: D = alpha * accumulator + beta *
/// source
CUTLASS_HOST_DEVICE
FragmentOutput operator()(FragmentAccumulator const &accumulator,
FragmentSource const &source,
int row_offset,
int column_offset) const {
CUTLASS_TRACE_DEVICE(
"kElementsPerAccess: %d, row_offset: %d, column_offset: %d",
kElementsPerAccess,
row_offset,
column_offset);
// Convert source to internal compute numeric type
cutlass::NumericArrayConverter<ElementCompute,
ElementOutput,
kElementsPerAccess,
Round>
source_converter;
cutlass::NumericArrayConverter<ElementCompute,
ElementAccumulator,
kElementsPerAccess,
Round>
accumulator_converter;
FragmentCompute converted_source = source_converter(source);
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
// Perform binary operations
FragmentCompute intermediate;
cutlass::multiplies<FragmentCompute> mul_add_source;
cutlass::multiply_add<FragmentCompute> mul_add_accumulator;
VariadicOp<ElementCompute> variadic_op;
if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling) {
intermediate = converted_source;
// D = alpha * Accum + X
intermediate = mul_add_accumulator(
params_.alpha, converted_accumulator, intermediate);
} else if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
intermediate = converted_accumulator;
} else {
// X = beta * C + uniform
intermediate = mul_add_source(params_.beta, converted_source);
// D = alpha * Accum + X
intermediate = mul_add_accumulator(
params_.alpha, converted_accumulator, intermediate);
}
if constexpr (GenericVariadicTraits<
VariadicOp<ElementCompute>>::IsArgumentsNeeded) {
if (!skip_elementwise_) {
#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
intermediate = variadic_op.Compute<kElementsPerAccess>(
intermediate,
params_.variadic_args,
BatchedMatrixCoord(blockIdx.z, row_offset, column_offset));
#else
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = variadic_op(
intermediate[i],
params_.variadic_args,
BatchedMatrixCoord(blockIdx.z, row_offset, column_offset + i));
}
#endif
}
} else {
if (!skip_elementwise_) {
#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
intermediate = variadic_op.Compute<kElementsPerAccess>(intermediate);
#else
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = variadic_op(intermediate[i]);
}
#endif
}
}
// Convert to destination numeric type
cutlass::NumericArrayConverter<ElementOutput,
ElementCompute,
kElementsPerAccess,
Round>
destination_converter;
return destination_converter(intermediate);
}
/// Computes linear scaling: D = alpha * accumulator
CUTLASS_HOST_DEVICE
FragmentOutput operator()(FragmentAccumulator const &accumulator,
int row_offset,
int column_offset) const {
CUTLASS_TRACE_DEVICE(
"kElementsPerAccess: %d, row_offset: %d, column_offset: %d",
kElementsPerAccess,
row_offset,
column_offset);
// Convert source to internal compute numeric type
cutlass::NumericArrayConverter<ElementCompute,
ElementAccumulator,
kElementsPerAccess,
Round>
accumulator_converter;
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
// Perform binary operations
FragmentCompute intermediate;
cutlass::multiplies<FragmentCompute> mul_accumulator;
VariadicOp<ElementCompute> variadic_op;
if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
intermediate = converted_accumulator;
} else {
// D = alpha * Accum
intermediate = mul_accumulator(params_.alpha, converted_accumulator);
}
if constexpr (GenericVariadicTraits<
VariadicOp<FragmentCompute>>::IsArgumentsNeeded) {
if (!skip_elementwise_) {
#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
intermediate = variadic_op.Compute<kElementsPerAccess>(
intermediate,
params_.variadic_args,
BatchedMatrixCoord(blockIdx.z, row_offset, column_offset));
#else
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = variadic_op(
intermediate[i],
params_.variadic_args,
BatchedMatrixCoord(blockIdx.z, row_offset, column_offset + i));
}
#endif
}
} else {
if (!skip_elementwise_) {
#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
intermediate = variadic_op.Compute<kElementsPerAccess>(intermediate);
#else
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kElementsPerAccess; ++i) {
intermediate[i] = variadic_op(intermediate[i]);
}
#endif
}
}
// Convert to destination numeric type
cutlass::NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
destination_converter;
return destination_converter(intermediate);
}
};
} // namespace thread
} // namespace epilogue
} // namespace cutlass_patch
python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/threadblock/default_epilogue_with_variadic.h
+243
View File
@@ -0,0 +1,243 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*! \file
\brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
The epilogue rearranges the result of a matrix product through shared memory
to match canonical tensor layouts in global memory. Epilogues support
conversion and reduction operations.
*/
#pragma once
#include "cutlass_patch/backend.h"
#ifdef __NVCC__
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
#include "cutlass/epilogue/threadblock/epilogue.h"
#include "cutlass/layout/permute.h"
#elif defined(__HIPCC__)
#include "hytlass/array.h"
#include "hytlass/numeric_types.h"
#include "hytlass/gemm/gemm.h"
#include "hytlass/epilogue/threadblock/default_epilogue_tensor_op.h"
#include "hytlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
#include "hytlass/epilogue/threadblock/epilogue.h"
#include "hytlass/layout/permute.h"
#endif
#include "cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h"
// #include "cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h"
namespace cutlass_patch {
namespace epilogue {
namespace threadblock {
/// Defines sensible defaults for epilogues for SimtOps.
template <typename Shape,
typename WarpMmaSimt,
typename ElementOutput,
typename OutputOp,
int ElementsPerAccess,
bool ScatterD = false,
typename PermuteDLayout = cutlass::layout::NoPermute,
cutlass::conv::StrideSupport StrideSupport =
cutlass::conv::StrideSupport::kUnity,
int Rank = 4>
struct DefaultEpilogueWithVariadicSimt {
static cutlass::conv::StrideSupport const kStrideSupport = StrideSupport;
static int const kRank = Rank;
static bool const UseCUDAStore =
cutlass::platform::is_same<ElementOutput, double>::value;
/// Use defaults related to the existing epilogue
using Base = cutlass::epilogue::threadblock::
DefaultEpilogueSimt<Shape, WarpMmaSimt, OutputOp, ElementsPerAccess>;
using PackedOutputTileIterator =
cutlass::epilogue::threadblock::PredicatedTileIterator<
typename Base::OutputTileThreadMap,
ElementOutput,
ScatterD,
PermuteDLayout,
UseCUDAStore>;
using StridedOutputTileIterator =
cutlass::epilogue::threadblock::PredicatedTileIteratorConv<
typename Base::OutputTileThreadMap,
ElementOutput,
ScatterD,
PermuteDLayout,
UseCUDAStore,
kRank>;
//
// Stores the result z = (y = GEMM(A, B, C), variadic)
//
using OutputTileIterator = typename cutlass::platform::conditional<
StrideSupport == cutlass::conv::StrideSupport::kUnity,
PackedOutputTileIterator,
StridedOutputTileIterator>::type;
//
// Define the epilogue
//
using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
Shape,
WarpMmaSimt,
Base::kPartitionsK,
OutputTileIterator,
typename Base::AccumulatorFragmentIterator,
typename Base::WarpTileIterator,
typename Base::SharedLoadIterator,
OutputOp,
typename Base::Padding>;
};
/// Defines sensible defaults for strided dgrad epilogues for SimtOps.
template <typename Shape,
typename WarpMmaSimt,
typename ElementOutput,
typename OutputOp,
int ElementsPerAccess,
bool ScatterD = false,
typename PermuteDLayout = cutlass::layout::NoPermute>
struct DefaultEpilogueWithVariadicSimtStridedDgrad {
/// Use defaults related to the existing epilogue
using Base = cutlass::epilogue::threadblock::DefaultEpilogueSimtStridedDgrad<
Shape,
WarpMmaSimt,
OutputOp,
ElementsPerAccess>;
//
// Stores the result z = (y = GEMM(A, B, C), variadic)
//
using OutputTileIterator =
cutlass::epilogue::threadblock::PredicatedTileIteratorStridedDgrad<
typename Base::OutputTileThreadMap,
ElementOutput>;
//
// Define the epilogue
//
using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
Shape,
WarpMmaSimt,
Base::kPartitionsK,
OutputTileIterator,
typename Base::AccumulatorFragmentIterator,
typename Base::WarpTileIterator,
typename Base::SharedLoadIterator,
OutputOp,
typename Base::Padding>;
};
/// Defines sensible defaults for epilogues for TensorOps.
template <typename Shape,
typename WarpMmaTensorOp,
int PartitionsK,
typename ElementOutput,
typename OutputOp,
int ElementsPerAccess,
bool ScatterD = false,
typename PermuteDLayout = cutlass::layout::NoPermute>
struct DefaultEpilogueWithVariadicTensorOp {
/// Use defaults related to the existing epilogue
using Base = cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
Shape,
WarpMmaTensorOp,
PartitionsK,
OutputOp,
ElementsPerAccess>;
//
// Stores the result z = (y = GEMM(A, B, C), variadic)
//
using OutputTileIterator =
cutlass::epilogue::threadblock::PredicatedTileIterator<
typename Base::OutputTileThreadMap,
ElementOutput,
ScatterD,
PermuteDLayout>;
//
// Define the epilogue
//
using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
Shape,
WarpMmaTensorOp,
PartitionsK,
OutputTileIterator,
typename Base::AccumulatorFragmentIterator,
typename Base::WarpTileIterator,
typename Base::SharedLoadIterator,
OutputOp,
typename Base::Padding,
Base::kFragmentsPerIteration>;
};
/// Defines sensible defaults for epilogues for VoltaTensorOps.
template <typename Shape,
typename WarpMmaTensorOp,
int PartitionsK,
typename ElementOutput,
typename OutputOp,
int ElementsPerAccess>
struct DefaultEpilogueWithVariadicVoltaTensorOp {
/// Use defaults related to the existing epilogue
using Base = cutlass::epilogue::threadblock::DefaultEpilogueVoltaTensorOp<
Shape,
WarpMmaTensorOp,
PartitionsK,
OutputOp,
ElementsPerAccess>;
//
// Stores the result z = (y = GEMM(A, B, C), variadic)
//
using OutputTileIterator = cutlass::epilogue::threadblock::
PredicatedTileIterator<typename Base::OutputTileThreadMap, ElementOutput>;
//
// Define the epilogue
//
using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
Shape,
WarpMmaTensorOp,
PartitionsK,
OutputTileIterator,
typename Base::AccumulatorFragmentIterator,
typename Base::WarpTileIterator,
typename Base::SharedLoadIterator,
OutputOp,
typename Base::Padding>;
};
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass_patch
python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h
+666
View File
@@ -0,0 +1,666 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*! \file
\brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
The epilogue rearranges the result of a matrix product through shared memory
to match canonical tensor layouts in global memory. Epilogues support
conversion and reduction operations.
The shared memory resource is time-sliced across warps.
*/
#pragma once
#include "cutlass_patch/backend.h"
#ifdef __NVCC__
#include <cuda/std/cassert>
#include "cutlass/aligned_buffer.h"
#include "cutlass/array.h"
#include "cutlass/functional.h"
#include "cutlass/layout/tensor.h"
#include "cutlass/layout/vector.h"
#include "cutlass/numeric_types.h"
#include "cutlass/tensor_coord.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/transform/pitch_linear_thread_map.h"
#include "cutlass/transform/threadblock/regular_tile_iterator.h"
#include "cutlass/epilogue/threadblock/epilogue_base.h"
#include "cutlass/epilogue/threadblock/epilogue_base_streamk.h"
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
#elif defined(__HIPCC__)
#include "hytlass/aligned_buffer.h"
#include "hytlass/array.h"
#include "hytlass/functional.h"
#include "hytlass/layout/tensor.h"
#include "hytlass/layout/vector.h"
#include "hytlass/numeric_types.h"
#include "hytlass/tensor_coord.h"
#include "hytlass/gemm/gemm.h"
#include "hytlass/transform/pitch_linear_thread_map.h"
#include "hytlass/transform/threadblock/regular_tile_iterator.h"
#include "hytlass/epilogue/threadblock/epilogue_base.h"
#include "hytlass/epilogue/threadblock/epilogue_base_streamk.h"
#include "hytlass/epilogue/threadblock/predicated_tile_iterator.h"
#endif
#include "cutlass_patch/trace_device.h"
namespace cutlass_patch {
namespace epilogue {
namespace threadblock {
/// Epilogue operator
template <
typename Shape_, ///< Shape of threadblock tile (concept: GemmShape)
typename WarpMmaOperator_, ///< Warp-level MMA operator (concept:
///< gemm::warp::MmaTensorOp)
int PartitionsK, ///< Number of partitions of the K dimension
typename OutputTileIterator_, ///< Tile iterator reading and writing
///< output tensors
typename AccumulatorFragmentIterator_, ///< Fragment iterator
///< selecting accumulators
typename WarpTileIterator_, ///< Warp-scoped tile iterator writing
///< accumulators to SMEM
typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator
///< loading from SMEM
typename OutputOp_, ///< Output operator
typename Padding_, ///< Padding added to SMEM allocation to avoid
///< bank conflicts (concept: MatrixShape)
int FragmentsPerPartition =
1, ///< Used to coarsten the epilogue granularity
int IterationsUnroll = ///< Used to reduce binary size when epilogue
///< op is large
(!cutlass::epilogue::threadblock::IsEpilogueFunctorHeavy<OutputOp_>::value)>
class EpilogueWithVariadic
: public cutlass::epilogue::threadblock::EpilogueBase<
Shape_,
typename WarpMmaOperator_::Shape,
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
Padding_,
FragmentsPerPartition>,
public cutlass::epilogue::threadblock::EpilogueBaseStreamK<
Shape_,
PartitionsK,
WarpMmaOperator_,
AccumulatorFragmentIterator_> {
public:
using Base = cutlass::epilogue::threadblock::EpilogueBase<
Shape_,
typename WarpMmaOperator_::Shape,
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
Padding_,
FragmentsPerPartition>;
using BaseStreamK = cutlass::epilogue::threadblock::EpilogueBaseStreamK<
Shape_,
PartitionsK,
WarpMmaOperator_,
AccumulatorFragmentIterator_>;
using Shape = Shape_;
using WarpMmaOperator = WarpMmaOperator_;
static int const kPartitionsK = PartitionsK;
using OutputTileIterator = OutputTileIterator_;
using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
using WarpTileIterator = WarpTileIterator_;
using SharedLoadIterator = SharedLoadIterator_;
using OutputOp = OutputOp_;
using Padding = Padding_;
using Layout = cutlass::layout::RowMajor;
using LongIndex = typename Layout::LongIndex;
/// Number of warps per block
using WarpCount = typename Base::WarpCount;
/// Number of threads per block
static int const kBlockThreads = 32 * WarpCount::kCount;
/// Per-thread accumulator tile type
using AccumulatorTile = typename Base::AccumulatorTile;
/// Numerical accumulation element type
using ElementAccumulator = typename WarpMmaOperator::ElementC;
/// Fragment type used by the accumulator tile's fragment iterator
using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment;
/// Output element
using ElementOutput = typename OutputTileIterator::Element;
/// Output access size
static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
/// Tensor reference to destination tensor
using TensorRef = typename OutputTileIterator::TensorRef;
/// Tensor reference to sync tensor
using SyncTensorRef =
typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
/// Const tensor reference to source tensor
using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
/// Vector type used by the global output iterator
using OutputAccessType =
cutlass::Array<typename OutputTileIterator::Element,
OutputTileIterator::kElementsPerAccess>;
/// Vector type used by the shared output iterator
using AccumulatorAccessType =
cutlass::Array<typename WarpTileIterator::Element,
OutputTileIterator::kElementsPerAccess>;
static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1
? Base::kFragmentsPerIteration
: kPartitionsK;
static int constexpr kSmemPointerOffset =
Base::SharedStorage::StorageShape::kCount / kSmemTiles;
public:
static_assert(
SharedLoadIterator::Fragment::kElements ==
OutputTileIterator::Fragment::kElements,
"Mismatch between shared load iterator and output tile iterator.");
static_assert(OutputTileIterator::kElementsPerAccess,
"OutputTileIterator::kElementsPerAccess must not be zero.");
static_assert(!(OutputTileIterator::Fragment::kElements %
OutputTileIterator::kElementsPerAccess),
"Divisibility");
static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1,
"One of these must be exactly 1.");
public:
/// Aspect for when epilogue source is not needed
struct SourceAspectNotNeeded {
/// Constructor
CUTLASS_DEVICE
SourceAspectNotNeeded() {}
// No-op
CUTLASS_DEVICE
void load() {}
/// Invoke the output functor over each vector of output
CUTLASS_DEVICE
void apply_output_operator(
const OutputTileIterator &output_iterator,
typename OutputTileIterator::Fragment &output_fragment, // NOLINT
OutputOp const &output_op,
typename SharedLoadIterator::Fragment const &aligned_accum_fragment) {
CUTLASS_TRACE_DEVICE("");
OutputAccessType *output_frag_ptr =
reinterpret_cast<OutputAccessType *>(&output_fragment);
AccumulatorAccessType const *compute_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(
&aligned_accum_fragment);
const int32_t thread_start_row = output_iterator.thread_start_row();
const int32_t thread_start_column = output_iterator.thread_start_column();
const typename OutputTileIterator::Index extent_row =
output_iterator.extent_row();
const typename OutputTileIterator::Index extent_column =
output_iterator.extent_column();
using ThreadMap = typename OutputTileIterator::ThreadMap;
typename OutputTileIterator::Mask mask;
output_iterator.get_mask(mask);
CUTLASS_PRAGMA_UNROLL
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
++cluster) {
CUTLASS_PRAGMA_UNROLL
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
CUTLASS_PRAGMA_UNROLL
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
int frag_row_idx =
(row + ThreadMap::Iterations::kRow *
(group + ThreadMap::Iterations::kGroup * cluster));
int row_offset = thread_start_row + row * ThreadMap::Delta::kRow +
group * ThreadMap::Delta::kGroup +
cluster * ThreadMap::Delta::kCluster;
bool row_guard = row_offset < extent_row;
CUTLASS_PRAGMA_UNROLL
for (int column = 0; column < ThreadMap::Iterations::kColumn;
++column) {
bool guard = row_guard && mask.predicates[column];
if (!guard) {
continue;
}
int column_offset =
thread_start_column + column * ThreadMap::Delta::kColumn;
int frag_offset =
frag_row_idx * ThreadMap::Iterations::kColumn + column;
output_frag_ptr[frag_offset] = output_op(
compute_frag_ptr[frag_offset], row_offset, column_offset);
}
}
}
}
}
};
/// Aspect for when epilogue source is needed
struct SourceAspectNeeded {
OutputTileIterator source_iterator;
typename OutputTileIterator::Fragment source_fragment;
/// Invoke the output functor over each vector of output
CUTLASS_DEVICE
static void apply_output_operator(
const OutputTileIterator &output_iterator,
typename OutputTileIterator::Fragment &output_fragment, // NOLINT
OutputOp const &output_op,
typename SharedLoadIterator::Fragment const &aligned_accum_fragment,
typename OutputTileIterator::Fragment const &source_fragment) {
CUTLASS_TRACE_DEVICE("");
OutputAccessType *output_frag_ptr =
reinterpret_cast<OutputAccessType *>(&output_fragment);
AccumulatorAccessType const *compute_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(
&aligned_accum_fragment);
OutputAccessType const *source_frag_ptr =
reinterpret_cast<OutputAccessType const *>(&source_fragment);
typename OutputTileIterator::Element const *source_ptr =
reinterpret_cast<typename OutputTileIterator::Element const *>(
&source_fragment);
const int32_t thread_start_row = output_iterator.thread_start_row();
const int32_t thread_start_column = output_iterator.thread_start_column();
const typename OutputTileIterator::Index extent_row =
output_iterator.extent_row();
const typename OutputTileIterator::Index extent_column =
output_iterator.extent_column();
using ThreadMap = typename OutputTileIterator::ThreadMap;
typename OutputTileIterator::Mask mask;
output_iterator.get_mask(mask);
CUTLASS_PRAGMA_UNROLL
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
++cluster) {
CUTLASS_PRAGMA_UNROLL
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
CUTLASS_PRAGMA_UNROLL
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
int frag_row_idx =
(row + ThreadMap::Iterations::kRow *
(group + ThreadMap::Iterations::kGroup * cluster));
int row_offset = thread_start_row + row * ThreadMap::Delta::kRow +
group * ThreadMap::Delta::kGroup +
cluster * ThreadMap::Delta::kCluster;
bool row_guard = row_offset < extent_row;
CUTLASS_PRAGMA_UNROLL
for (int column = 0; column < ThreadMap::Iterations::kColumn;
++column) {
bool guard = row_guard && mask.predicates[column];
if (!guard) {
continue;
}
int column_offset =
thread_start_column + column * ThreadMap::Delta::kColumn;
int frag_offset =
frag_row_idx * ThreadMap::Iterations::kColumn + column;
output_frag_ptr[frag_offset] =
output_op(compute_frag_ptr[frag_offset],
source_frag_ptr[frag_offset],
row_offset,
column_offset);
}
}
}
}
}
/// Constructor
CUTLASS_DEVICE
explicit SourceAspectNeeded(OutputTileIterator source_iterator)
: source_iterator(source_iterator) {
source_fragment.clear();
}
// Load addend source fragment from global memory
CUTLASS_DEVICE
void load() {
source_iterator.load(source_fragment);
++source_iterator;
}
/// Invoke the output functor over each vector of output
CUTLASS_DEVICE
void apply_output_operator(
const OutputTileIterator &output_iterator,
typename OutputTileIterator::Fragment &output_fragment, // NOLINT
OutputOp const &output_op,
typename SharedLoadIterator::Fragment const &aligned_accum_fragment) {
apply_output_operator(output_iterator,
output_fragment,
output_op,
aligned_accum_fragment,
source_fragment);
}
};
private:
/// Loads fragment from shared memory aligned with output tensor
SharedLoadIterator shared_load_iterator_;
/// Thread index in the threadblock
int thread_idx;
/// Warp index in the threadblock
int warp_idx;
public:
/// Constructor
CUTLASS_DEVICE
EpilogueWithVariadic(
typename Base::SharedStorage
&shared_storage, // NOLINT ///< Shared storage object
int thread_idx, ///< ID of a thread within the threadblock
int warp_idx, ///< ID of warp within threadblock
int lane_idx) ///< Id of thread within warp
: Base(shared_storage, thread_idx, warp_idx, lane_idx),
BaseStreamK(thread_idx),
shared_load_iterator_(shared_storage.reference(), thread_idx),
thread_idx(thread_idx),
warp_idx(warp_idx) {}
/// Aggregates the accumulator sets shared by peer blocks in the global
/// workspace, performing epilogue computations, writing to output
CUTLASS_DEVICE
void reduce(int peer_idx_begin,
int peer_idx_end,
int reduce_fragment_idx,
void *element_workspace,
OutputOp const &output_op, ///< Output operator
OutputTileIterator
destination_iterator, ///< Tile iterator for destination
OutputTileIterator
source_iterator) { ///< Threadblock tile coordinate in GEMM
///< (in units of threadblock tiles)
CUTLASS_TRACE_DEVICE("");
// Reduce peer accumulator fragments into one fragment
AccumulatorFragment accum_fragment;
BaseStreamK::reduce(accum_fragment,
peer_idx_begin,
peer_idx_end,
reduce_fragment_idx,
element_workspace);
// Store fragment to shared memory
this->warp_tile_iterator_.store(accum_fragment);
__syncthreads();
// Initialize/load source-fragment data
typename OutputTileIterator::Fragment source_fragment;
source_fragment.clear();
if (output_op.is_source_needed()) {
source_iterator += reduce_fragment_idx;
source_iterator.load(source_fragment);
}
// Load fragment from shared memory
typename SharedLoadIterator::Fragment aligned_accum_fragment;
shared_load_iterator_.load(aligned_accum_fragment);
// Add fragments shared by other k partitions
if (kPartitionsK > 1) {
cutlass::plus<typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for (int i = 1; i < kPartitionsK; ++i) {
typename SharedLoadIterator::Fragment aligned_addend_fragment;
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_addend_fragment);
aligned_accum_fragment =
add_fragments(aligned_accum_fragment, aligned_addend_fragment);
}
}
// Compute the output result
typename OutputTileIterator::Fragment output_fragment;
// Apply the output operator
SourceAspectNeeded::apply_output_operator(
output_fragment, output_op, aligned_accum_fragment, source_fragment);
// Store the final result
destination_iterator += reduce_fragment_idx;
destination_iterator.store(output_fragment);
}
/// Perform the epilogue computations and stream the result to global memory.
CUTLASS_DEVICE
void operator()(OutputOp const &output_op, ///< Output operator
OutputTileIterator
destination_iterator, ///< Tile iterator for destination
AccumulatorTile const &
accumulators) { ///< Complete warp-level accumulator tile
CUTLASS_TRACE_DEVICE("");
operator()(
output_op, destination_iterator, accumulators, SourceAspectNotNeeded());
}
/// Perform the epilogue computations and stream the result to global memory.
/// Implements two alternative codepaths, depending on whether the output op
/// requires addend data to be loaded.
CUTLASS_DEVICE
void operator()(OutputOp const &output_op, ///< Output operator
OutputTileIterator
destination_iterator, ///< Tile iterator for destination
AccumulatorTile const
&accumulators, ///< Complete warp-level accumulator tile
OutputTileIterator
source_iterator) { ///< Tile iterator for addend source
CUTLASS_TRACE_DEVICE("");
if (output_op.is_source_needed()) {
operator()(output_op,
destination_iterator,
accumulators,
SourceAspectNeeded(source_iterator));
} else {
operator()(output_op,
destination_iterator,
accumulators,
SourceAspectNotNeeded());
}
}
/// Perform the epilogue computations and stream the result to global memory.
/// Implements a single codepath, regardless of whether the output op requires
/// addend data to be loaded
CUTLASS_DEVICE
void unified(OutputOp const &output_op, ///< Output operator
OutputTileIterator
destination_iterator, ///< Tile iterator for destination
AccumulatorTile const
&accumulators, ///< Complete warp-level accumulator tile
OutputTileIterator
source_iterator) { ///< Tile iterator for addend source
CUTLASS_TRACE_DEVICE("");
if (!output_op.is_source_needed()) {
source_iterator.clear_mask();
__syncthreads(); // Dummy (CUDA 11.0)
}
operator()(output_op,
destination_iterator,
accumulators,
SourceAspectNeeded(source_iterator));
}
template <class Seq>
struct acc2smem;
template <size_t... Seq>
struct acc2smem<cutlass::index_sequence<Seq...>> {
template <int Advance>
CUTLASS_DEVICE static void helper(
AccumulatorFragmentIterator accum_fragment_iterator,
WarpTileIterator &warp_tile_iterator) { // NOLINT
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Advance; i++) {
++accum_fragment_iterator;
}
typename AccumulatorFragmentIterator::Fragment accum_fragment;
accum_fragment_iterator.load(accum_fragment);
++accum_fragment_iterator;
warp_tile_iterator.store(accum_fragment);
}
CUTLASS_DEVICE
static void push(size_t pos,
AccumulatorFragmentIterator const &iterator_begin,
WarpTileIterator &warp_tile_iterator) { // NOLINT
int dummy[] = {(pos == Seq) &&
(helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
}
};
/// Streams the result to global memory
template <typename SourceAspect>
CUTLASS_DEVICE void operator()(
OutputOp const &output_op, ///< Output operator
OutputTileIterator
destination_iterator, ///< Tile iterator for destination
AccumulatorTile const
&accumulators, ///< Complete warp-level accumulator tile
SourceAspect source) {
CUTLASS_TRACE_DEVICE("");
// Iterator over warp-level accumulator fragment
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
//
// Iterate over accumulator tile
//
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wcuda-compat"
// Turn off clangs warning about loop unroll argument using parens.
#endif
#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1)
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
//
// Load the source
//
source.load();
//
// Convert and store fragment
//
__syncthreads();
acc2smem<cutlass::make_index_sequence<OutputTileIterator::kIterations>>::
push(iter, accum_fragment_iterator, this->warp_tile_iterator_);
__syncthreads();
//
// Load fragments from shared memory
//
typename SharedLoadIterator::Fragment
aligned_accum_fragment[kPartitionsK];
shared_load_iterator_.load(aligned_accum_fragment[0]);
if (kPartitionsK > 1) {
cutlass::plus<typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for (int i = 1; i < kPartitionsK; ++i) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_accum_fragment[i]);
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0],
aligned_accum_fragment[i]);
}
shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) *
kSmemPointerOffset);
}
//
// Compute the output result
//
typename OutputTileIterator::Fragment output_fragment;
source.apply_output_operator(destination_iterator,
output_fragment,
output_op,
aligned_accum_fragment[0]);
//
// Store the final result
//
destination_iterator.store(output_fragment);
++destination_iterator;
}
#ifdef __clang__
#pragma clang diagnostic pop
#endif
}
};
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass_patch
python/paddle/apy/matmul_pass/matmul/cutlass_patch/gemm/device/gemm_universal_with_variadic.h
+416
View File
@@ -0,0 +1,416 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*! \file
\brief
*/
#pragma once
#include "cutlass_patch/backend.h"
#ifdef __NVCC__
#include "cutlass/arch/arch.h"
#include "cutlass/arch/mma.h"
#include "cutlass/device_kernel.h"
#include "cutlass/numeric_types.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/kernel/gemm_universal.h"
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
#include "cutlass/gemm/device/default_gemm_configuration.h"
#include "cutlass/gemm/device/gemm_universal_base.h"
#include "cutlass/gemm/kernel/default_gemm_universal.h"
#include "cutlass/layout/permute.h"
#elif defined(__HIPCC__)
#include "hytlass/arch/arch.h"
#include "hytlass/arch/mma.h"
#include "hytlass/device_kernel.h"
#include "hytlass/numeric_types.h"
#include "hytlass/gemm/gemm.h"
#include "hytlass/gemm/kernel/gemm_universal.h"
#include "hytlass/gemm/threadblock/threadblock_swizzle.h"
#include "hytlass/gemm/device/default_gemm_configuration.h"
#include "hytlass/gemm/device/gemm_universal_base.h"
#include "hytlass/gemm/kernel/default_gemm_universal.h"
#include "hytlass/layout/permute.h"
#endif
#include "cutlass_patch/gemm/kernel/default_gemm_with_variadic.h"
namespace cutlass_patch {
namespace gemm {
namespace device {
/*!
GemmUniversal with variadic epilogues.
*/
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Element type for C and D matrix operands
typename ElementC_,
/// Layout type for C and D matrix operands
typename LayoutC_,
/// Element type for internal accumulation
typename ElementAccumulator_ = ElementC_,
/// Operator class tag
typename OperatorClass_ = cutlass::arch::OpClassSimt,
/// Tag indicating architecture to tune for. This is the minimum SM that
/// supports the intended feature. The device kernel can be built
/// targeting any SM larger than this number.
typename ArchTag_ = cutlass::arch::Sm70,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape_ = typename cutlass::gemm::device::
DefaultGemmConfiguration<OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape_ = typename cutlass::gemm::device::
DefaultGemmConfiguration<OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::WarpShape,
/// Instruction-level tile size (concept: GemmShape)
typename InstructionShape_ = typename cutlass::gemm::device::
DefaultGemmConfiguration<OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::InstructionShape,
/// Epilogue output operator
typename EpilogueOutputOp_ = typename cutlass::gemm::device::
DefaultGemmConfiguration<OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle_ =
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
/// Number of stages used in the pipelined mainloop
int Stages = cutlass::gemm::device::DefaultGemmConfiguration<
OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::kStages,
/// Access granularity of A matrix in units of elements
int AlignmentA = cutlass::gemm::device::DefaultGemmConfiguration<
OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::kAlignmentA,
/// Access granularity of B matrix in units of elements
int AlignmentB = cutlass::gemm::device::DefaultGemmConfiguration<
OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::kAlignmentB,
/// Operation performed by GEMM
typename Operator_ = typename cutlass::gemm::device::
DefaultGemmConfiguration<OperatorClass_,
ArchTag_,
ElementA_,
ElementB_,
ElementC_,
ElementAccumulator_>::Operator,
/// Complex elementwise transformation on A operand
cutlass::ComplexTransform TransformA = cutlass::ComplexTransform::kNone,
/// Complex elementwise transformation on B operand
cutlass::ComplexTransform TransformB = cutlass::ComplexTransform::kNone>
class GemmUniversalWithVariadic
: public cutlass::gemm::device::GemmUniversalBase<
typename cutlass_patch::gemm::kernel::DefaultGemmWithVariadic<
ElementA_,
LayoutA_,
TransformA,
AlignmentA,
ElementB_,
LayoutB_,
TransformB,
AlignmentB,
ElementC_,
LayoutC_,
ElementAccumulator_,
OperatorClass_,
ArchTag_,
ThreadblockShape_,
WarpShape_,
InstructionShape_,
EpilogueOutputOp_,
ThreadblockSwizzle_,
Stages,
Operator_>::GemmKernel> {
public:
using ElementAccumulator = ElementAccumulator_;
using OperatorClass = OperatorClass_;
using ArchTag = ArchTag_;
using ThreadblockShape = ThreadblockShape_;
using WarpShape = WarpShape_;
using InstructionShape = InstructionShape_;
using EpilogueOutputOp = EpilogueOutputOp_;
using ThreadblockSwizzle = ThreadblockSwizzle_;
using Operator = Operator_;
static int const kStages = Stages;
static int const kAlignmentA = AlignmentA;
static int const kAlignmentB = AlignmentB;
static int const kAlignmentC = EpilogueOutputOp::kCount;
static cutlass::ComplexTransform const kTransformA = TransformA;
static cutlass::ComplexTransform const kTransformB = TransformB;
using Base = cutlass::gemm::device::GemmUniversalBase<
typename cutlass_patch::gemm::kernel::DefaultGemmWithVariadic<
ElementA_,
LayoutA_,
TransformA,
AlignmentA,
ElementB_,
LayoutB_,
TransformB,
AlignmentB,
ElementC_,
LayoutC_,
ElementAccumulator_,
OperatorClass_,
ArchTag_,
ThreadblockShape_,
WarpShape_,
InstructionShape_,
EpilogueOutputOp_,
ThreadblockSwizzle_,
Stages,
Operator_>::GemmKernel>;
using Arguments = typename Base::Arguments;
using GemmKernel = typename Base::GemmKernel;
};
/// Partial specialization for column-major output exchanges problem size and
/// operand.
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Element type for C and D matrix operands
typename ElementC_,
/// Element type for internal accumulation
typename ElementAccumulator_,
/// Operator class tag
typename OperatorClass_,
/// Tag indicating architecture to tune for. This is the minimum SM that
/// supports the intended feature. The device kernel can be built
/// targeting any SM larger than this number.
typename ArchTag_,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape_,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape_,
/// Instruction-level tile size (concept: GemmShape)
typename InstructionShape_,
/// Epilogue output operator
typename EpilogueOutputOp_,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle_,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Access granularity of A matrix in units of elements
int AlignmentA,
/// Access granularity of B matrix in units of elements
int AlignmentB,
/// Operation performed by GEMM
typename Operator_,
/// Complex elementwise transformation on A operand
cutlass::ComplexTransform TransformA,
/// Complex elementwise transformation on B operand
cutlass::ComplexTransform TransformB>
class GemmUniversalWithVariadic<ElementA_,
LayoutA_,
ElementB_,
LayoutB_,
ElementC_,
cutlass::layout::ColumnMajor, // partially
// specialized on
// LayoutC
ElementAccumulator_,
OperatorClass_,
ArchTag_,
ThreadblockShape_,
WarpShape_,
InstructionShape_,
EpilogueOutputOp_,
ThreadblockSwizzle_,
Stages,
AlignmentA,
AlignmentB,
Operator_,
TransformA,
TransformB> {
public:
using ElementA = ElementA_;
using LayoutA = LayoutA_;
using TensorRefA = cutlass::TensorRef<ElementA const, LayoutA>;
using ElementB = ElementB_;
using LayoutB = LayoutB_;
using TensorRefB = cutlass::TensorRef<ElementB const, LayoutB>;
using ElementC = ElementC_;
using LayoutC = cutlass::layout::ColumnMajor;
using TensorRefC = cutlass::TensorRef<ElementC const, LayoutC>;
using TensorRefD = cutlass::TensorRef<ElementC, LayoutC>;
using ElementAccumulator = ElementAccumulator_;
using OperatorClass = OperatorClass_;
using ArchTag = ArchTag_;
using ThreadblockShape = ThreadblockShape_;
using WarpShape = WarpShape_;
using InstructionShape = InstructionShape_;
using EpilogueOutputOp = EpilogueOutputOp_;
using ThreadblockSwizzle = ThreadblockSwizzle_;
using Operator = Operator_;
static int const kStages = Stages;
static int const kAlignmentA = AlignmentA;
static int const kAlignmentB = AlignmentB;
static cutlass::ComplexTransform const kTransformA = TransformA;
static cutlass::ComplexTransform const kTransformB = TransformB;
using UnderlyingOperator = typename GemmUniversalWithVariadic<
ElementB,
typename cutlass::layout::LayoutTranspose<LayoutB>::type,
ElementA,
typename cutlass::layout::LayoutTranspose<LayoutA>::type,
ElementC,
cutlass::layout::RowMajor,
ElementAccumulator,
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
kAlignmentB,
kAlignmentA,
Operator,
kTransformB,
kTransformA>::Base;
using GemmKernel = typename UnderlyingOperator::GemmKernel;
static int const kAlignmentC = EpilogueOutputOp::kCount;
/// Argument structure
using Arguments = typename UnderlyingOperator::Arguments;
private:
UnderlyingOperator underlying_operator_;
public:
/// Constructs the GEMM.
GemmUniversalWithVariadic() {}
/// Helper to construct a transposed equivalent for the underlying GEMM
/// operator
static Arguments to_underlying_arguments(Arguments const &args) {
return args.transposed_problem();
}
/// Determines whether the GEMM can execute the given problem.
static cutlass::Status can_implement(Arguments const &args) {
return UnderlyingOperator::can_implement(to_underlying_arguments(args));
}
/// Gets the workspace size
static size_t get_workspace_size(Arguments const &args) {
return UnderlyingOperator::get_workspace_size(
to_underlying_arguments(args));
}
/// Computes the grid shape
static dim3 get_grid_shape(Arguments const &args) {
return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args));
}
/// Computes the maximum number of active blocks per multiprocessor
static int maximum_active_blocks(int smem_capacity = -1) {
return UnderlyingOperator::maximum_active_blocks(smem_capacity);
}
/// Initializes GEMM state from arguments.
cutlass::Status initialize(Arguments const &args,
void *workspace = nullptr,
GPUStream_t stream = nullptr) {
return underlying_operator_.initialize(
to_underlying_arguments(args), workspace, stream);
}
/// Lightweight update given a subset of arguments
cutlass::Status update(Arguments const &args, void *workspace = nullptr) {
return underlying_operator_.update(to_underlying_arguments(args),
workspace);
}
/// Runs the kernel using initialized state.
cutlass::Status run(GPUStream_t stream = nullptr) {
return underlying_operator_.run(stream);
}
/// Runs the kernel using initialized state.
cutlass::Status operator()(GPUStream_t stream = nullptr) {
return run(stream);
}
/// Runs the kernel using initialized state.
cutlass::Status operator()(Arguments const &args,
void *workspace = nullptr,
GPUStream_t stream = nullptr) {
cutlass::Status status = initialize(args, workspace, stream);
if (status == cutlass::Status::kSuccess) {
status = run(stream);
}
return status;
}
};
} // namespace device
} // namespace gemm
} // namespace cutlass_patch
python/paddle/apy/matmul_pass/matmul/cutlass_patch/gemm/kernel/default_gemm_with_variadic.h
+227
View File
@@ -0,0 +1,227 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*! \file
\brief
Defines a GEMM with Reduction based on an existing UniversalGemm kernel.
*/
#pragma once
#include "cutlass_patch/backend.h"
#ifdef __NVCC__
#include "cutlass/gemm/kernel/default_gemm_universal.h"
#include "cutlass/gemm/kernel/gemm_universal.h"
#elif defined(__HIPCC__)
#include "hytlass/gemm/kernel/default_gemm_universal.h"
#include "hytlass/gemm/kernel/gemm_universal.h"
#endif
#include "cutlass_patch/epilogue/threadblock/default_epilogue_with_variadic.h"
#include "cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h"
namespace cutlass_patch {
namespace gemm {
namespace kernel {
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Complex elementwise transformation on A operand
cutlass::ComplexTransform TransformA,
/// Access granularity of A matrix in units of elements
int kAlignmentA,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Complex elementwise transformation on B operand
cutlass::ComplexTransform TransformB,
/// Access granularity of B matrix in units of elements
int kAlignmentB,
/// Element type for C and D matrix operands
typename ElementC_,
/// Layout type for C and D matrix operands
typename LayoutC_,
/// Element type for internal accumulation
typename ElementAccumulator,
/// Operator class tag
typename OperatorClass,
/// Tag indicating architecture to tune for
typename ArchTag,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape,
/// Warp-level tile size (concept: GemmShape)
typename InstructionShape,
/// Epilogue output operator - must satisfy concept of
/// 'EpilogueWithVariadicOp'
typename EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Operation performed by GEMM
typename Operator,
///
typename Enable = void>
struct DefaultGemmWithVariadic {
using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<
ElementA_,
LayoutA_,
TransformA,
kAlignmentA,
ElementB_,
LayoutB_,
TransformB,
kAlignmentB,
ElementC_,
LayoutC_,
ElementAccumulator,
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
Operator>::GemmKernel;
// Define epilogue
using Epilogue = typename cutlass_patch::epilogue::threadblock::
DefaultEpilogueWithVariadicTensorOp<
typename GemmBase::Epilogue::Shape,
typename GemmBase::Epilogue::WarpMmaOperator,
GemmBase::Epilogue::kPartitionsK,
ElementC_,
EpilogueOutputOp,
GemmBase::Epilogue::kElementsPerAccess>::Epilogue;
// Compose the GEMM kernel
using GemmKernel = cutlass::gemm::kernel::
GemmUniversal<typename GemmBase::Mma, Epilogue, ThreadblockSwizzle>;
};
/// Partial specialization: ArchTag = cutlass::arch::Sm70
///
///
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Complex elementwise transformation on A operand
cutlass::ComplexTransform TransformA,
/// Access granularity of A matrix in units of elements
int kAlignmentA,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Complex elementwise transformation on B operand
cutlass::ComplexTransform TransformB,
/// Access granularity of B matrix in units of elements
int kAlignmentB,
/// Element type for C and D matrix operands
typename ElementC_,
/// Layout type for C and D matrix operands
typename LayoutC_,
/// Element type for internal accumulation
typename ElementAccumulator,
/// Operator class tag
typename OperatorClass,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape,
/// Warp-level tile size (concept: GemmShape)
typename InstructionShape,
/// Epilogue output operator - must satisfy concept of
/// 'EpilogueWithVariadicOp'
typename EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Operation performed by GEMM
typename Operator,
///
typename Enable>
struct DefaultGemmWithVariadic<ElementA_,
LayoutA_,
TransformA,
kAlignmentA,
ElementB_,
LayoutB_,
TransformB,
kAlignmentB,
ElementC_,
LayoutC_,
ElementAccumulator,
OperatorClass,
cutlass::arch::Sm70,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
Operator,
Enable> {
using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<
ElementA_,
LayoutA_,
TransformA,
kAlignmentA,
ElementB_,
LayoutB_,
TransformB,
kAlignmentB,
ElementC_,
LayoutC_,
ElementAccumulator,
OperatorClass,
cutlass::arch::Sm70,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
Operator>::GemmKernel;
// Define epilogue
using Epilogue = typename cutlass_patch::epilogue::threadblock::
DefaultEpilogueWithVariadicVoltaTensorOp<
typename GemmBase::Epilogue::Shape,
typename GemmBase::Epilogue::WarpMmaOperator,
GemmBase::Epilogue::kPartitionsK,
ElementC_,
EpilogueOutputOp,
GemmBase::Epilogue::kElementsPerAccess>::Epilogue;
// Compose the GEMM kernel
using GemmKernel = cutlass::gemm::kernel::
GemmUniversal<typename GemmBase::Mma, Epilogue, ThreadblockSwizzle>;
};
} // namespace kernel
} // namespace gemm
} // namespace cutlass_patch
python/paddle/apy/matmul_pass/matmul/cutlass_patch/hip/all_tuning_configs.h
+552
View File
@@ -0,0 +1,552 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "hytlass/gemm_coord.h"
namespace ap {
constexpr int kNumConfigsHalf = 28;
constexpr int kNumConfigsFloat = 13;
template <int SwizzleFactor, bool Batched>
struct SwizzleWrapper {
using Type =
hytlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
};
#define AP_AUTOTUNE(func, stream_ptr, count, ...) \
{ \
using FuncType = decltype(func<0>); \
static int selected_config_id = -1; \
static std::vector<std::function<FuncType>> matmul_functions = \
[]<std::size_t... Is>(std::index_sequence<Is...>) { \
return std::vector<std::function<FuncType>>{func<Is>...}; \
} \
(std::make_index_sequence<count>()); \
if (selected_config_id == -1) { \
selected_config_id = \
ap::ProfileBestConfig(matmul_functions, stream_ptr, ##__VA_ARGS__); \
} \
matmul_functions[selected_config_id](__VA_ARGS__); \
}
#define AP_AUTOTUNE_half(func, stream_ptr, ...) \
AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsHalf, __VA_ARGS__)
#define AP_AUTOTUNE_float(func, stream_ptr, ...) \
AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsFloat, __VA_ARGS__)
#define AP_AUTOTUNE_bfloat16(func, stream_ptr, ...) \
AP_AUTOTUNE_half(func, stream_ptr, __VA_ARGS__)
template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
struct GemmTuningConfigs {
using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 2;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = Id;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 1> {
using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
using WShape = hytlass::gemm::GemmShape<32, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 1;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 2> {
using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 2;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 3> {
using TShape = hytlass::gemm::GemmShape<128, 64, 64>;
using WShape = hytlass::gemm::GemmShape<64, 32, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 3;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 4> {
using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 4;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 5> {
using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 5;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 6> {
using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 6;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 7> {
using TShape = hytlass::gemm::GemmShape<256, 64, 64>;
using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 7;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 8> {
using TShape = hytlass::gemm::GemmShape<256, 128, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 8;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 9> {
using TShape = hytlass::gemm::GemmShape<256, 128, 64>;
using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 9;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 10> {
using TShape = hytlass::gemm::GemmShape<128, 32, 64>;
using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 10;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 11> {
using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 11;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 12> {
using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 12;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 13> {
using TShape = hytlass::gemm::GemmShape<256, 64, 64>;
using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 13;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 14> {
using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 14;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 15> {
using TShape = hytlass::gemm::GemmShape<32, 64, 64>;
using WShape = hytlass::gemm::GemmShape<16, 32, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 15;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 16> {
using TShape = hytlass::gemm::GemmShape<64, 64, 64>;
using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 16;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 17> {
using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 17;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 18> {
using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 5;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 18;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 19> {
using TShape = hytlass::gemm::GemmShape<64, 128, 32>;
using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 6;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 19;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 20> {
using TShape = hytlass::gemm::GemmShape<128, 64, 32>;
using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 6;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 20;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 21> {
using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 10;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 21;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 22> {
using TShape = hytlass::gemm::GemmShape<128, 256, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 2;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 22;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 23> {
using TShape = hytlass::gemm::GemmShape<128, 256, 32>;
using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 23;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 24> {
using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 24;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 25> {
using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 25;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 26> {
using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 26;
};
template <typename ElementT, int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 27> {
using TShape = hytlass::gemm::GemmShape<128, 64, 64>;
using WShape = hytlass::gemm::GemmShape<64, 32, 64>;
using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 27;
};
// Specialization for float
template <int SwizzleFactor, bool Batched, int Id>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
using TShape = hytlass::gemm::GemmShape<64, 64, 16>;
using WShape = hytlass::gemm::GemmShape<32, 32, 16>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = Id;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 1> {
using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 1;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 2> {
using TShape = hytlass::gemm::GemmShape<64, 128, 32>;
using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 2;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 3> {
using TShape = hytlass::gemm::GemmShape<64, 256, 16>;
using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 3;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 4> {
using TShape = hytlass::gemm::GemmShape<64, 256, 32>;
using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 4;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 5> {
using TShape = hytlass::gemm::GemmShape<128, 64, 32>;
using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 5;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 6> {
using TShape = hytlass::gemm::GemmShape<128, 128, 16>;
using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 6;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 7> {
using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 7;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 8> {
using TShape = hytlass::gemm::GemmShape<256, 64, 16>;
using WShape = hytlass::gemm::GemmShape<64, 32, 16>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 8;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 9> {
using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 3;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 9;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 10> {
using TShape = hytlass::gemm::GemmShape<64, 128, 16>;
using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 10;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 11> {
using TShape = hytlass::gemm::GemmShape<128, 64, 16>;
using WShape = hytlass::gemm::GemmShape<64, 32, 16>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 11;
};
template <int SwizzleFactor, bool Batched>
struct GemmTuningConfigs<float, SwizzleFactor, Batched, 12> {
using TShape = hytlass::gemm::GemmShape<128, 128, 16>;
using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
static constexpr int kNumStages = 4;
using SwizzleThreadBlock =
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
static constexpr int kId = 12;
};
struct DefaultConfig {
static constexpr int kConfigId = 0;
static constexpr int kSwizzleFactor = 1;
static constexpr bool kBatched = false;
};
} // namespace ap
python/paddle/apy/matmul_pass/matmul/cutlass_patch/hip/hytlass_matmul.h
+254
View File
@@ -0,0 +1,254 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "hytlass/epilogue/thread/linear_combination_bias_elementwise.h"
#include "hytlass/gemm/device/gemm_universal.h"
#include "hytlass/gemm/device/gemm_universal_with_broadcast.h"
#include "hytlass/util/device_memory.h"
#include "cutlass_patch/batched_matrix_coord.h"
#include "cutlass_patch/epilogue/thread/linear_combination_unary.h"
#include "cutlass_patch/epilogue/thread/linear_combination_variadic.h"
#include "cutlass_patch/gemm/device/gemm_universal_with_variadic.h"
#include "cutlass_patch/hip/all_tuning_configs.h"
#include "params.h" // NOLINT
#define CHECK_HYTLASS(status) \
{ \
auto error = status; \
if (error != hytlass::Status::kSuccess) { \
std::cerr << "HYTLASS error = " << int(error) << " (" \
<< hytlassGetStatusString(error) << ")" \
<< " at line " << __LINE__ << std::endl; \
std::abort(); \
} \
}
namespace ap {
using MatrixCoord = cutlass_patch::BatchedMatrixCoord;
using bfloat16 = __hip_bfloat16;
// Operation performed by GEMM
template <typename ElementT>
struct GemmOperation {
using Type = hytlass::arch::OpMultiplyAdd;
};
template <>
struct GemmOperation<float> {
using Type = hytlass::arch::OpMultiplyAddFastF32;
};
static hytlass::gemm::GemmUniversalMode GetGemmMode(int batch_count) {
return batch_count > 1 ? hytlass::gemm::GemmUniversalMode::kBatched
: hytlass::gemm::GemmUniversalMode::kGemm;
}
static void *GetWorkspace(size_t workspace_size) {
static hytlass::device_memory::allocation<uint8_t> workspace;
if (workspace.size() < workspace_size) {
workspace.reset(workspace_size);
}
return workspace.get();
}
template <typename GemmFunc>
hytlass::Status SetMaxDynamicSharedMemorySize() {
hipError_t hiprt_result;
// If requires more than 48KB: configure for extended, dynamic shared memory
if constexpr (GemmFunc::kSharedStorageSize >= (48 << 10)) {
hiprt_result = hipFuncSetAttribute(
(const void *)hytlass::Kernel2<typename GemmFunc::GemmKernel>,
hipFuncAttributeMaxDynamicSharedMemorySize,
GemmFunc::kSharedStorageSize);
if (hiprt_result != hipSuccess) {
HYTLASS_TRACE_HOST("hipFuncSetAttribute() returned error "
<< hipGetErrorString(hiprt_result));
return hytlass::Status::kErrorInternal;
}
}
#if AP_ENABLE_DEBUG
// Update SM occupancy member
int sm_occupancy = -1;
hiprt_result = hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
&sm_occupancy,
hytlass::Kernel2<typename GemmFunc::GemmKernel>,
GemmFunc::GemmKernel::kThreadCount,
GemmFunc::kSharedStorageSize,
hipOccupancyDisableCachingOverride);
if (hiprt_result != hipSuccess) {
HYTLASS_TRACE_HOST(
"hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned "
"error "
<< hipGetErrorString(hiprt_result));
return hytlass::Status::kErrorInternal;
}
HYTLASS_TRACE_HOST("sm_occupancy: (" << sm_occupancy
<< ") "
"smem_size: ("
<< GemmFunc::kSharedStorageSize
<< ") "
"GemmKernel::kThreadCount: ("
<< GemmFunc::GemmKernel::kThreadCount
<< ")");
#endif
return hytlass::Status::kSuccess;
}
// Convert HIP data type to hytlass data type
template <typename T>
struct HytlassDataType {
using Type = T;
};
template <>
struct HytlassDataType<half> {
using Type = hytlass::half_t;
};
template <>
struct HytlassDataType<__hip_bfloat16> {
using Type = hytlass::bfloat16_t;
};
// Convert to hytlass layout
template <bool Transposed>
struct MatrixLayout {
using Type = hytlass::layout::RowMajor;
};
template <>
struct MatrixLayout<true> {
using Type = hytlass::layout::ColumnMajor;
};
template <typename ElementT,
typename ElementComputeT,
template <typename T>
class VariadicFunctor,
int AlignA = 128 / hytlass::sizeof_bits<ElementT>::value,
int AlignB = 128 / hytlass::sizeof_bits<ElementT>::value,
int ConfigId = DefaultConfig::kConfigId,
int SwizzleFactor = DefaultConfig::kSwizzleFactor,
bool Batched = DefaultConfig::kBatched>
void MatmulAddVariadic(
const GemmEpilogueParams &params,
const typename VariadicFunctor<ElementComputeT>::Arguments &variadic_args) {
// <- data type of accumulator
using ElementAccumulator = typename HytlassDataType<ElementComputeT>::Type;
// <- data type of epilogue operations
using ElementComputeEpilogue = ElementAccumulator;
// <- data type of elements in input matrix A
using ElementInputA = typename HytlassDataType<ElementT>::Type;
// <- data type of elements in input matrix B
using ElementInputB = typename HytlassDataType<ElementT>::Type;
// <- data type of elements in output matrix D
using ElementOutput = typename HytlassDataType<ElementT>::Type;
constexpr int AlignC = AlignB;
// Epilogue operation as LinearCombination:
// alpha * accumulator + beta * source
using EpilogueOutputOp =
cutlass_patch::epilogue::thread::LinearCombinationVariadic<
VariadicFunctor,
ElementOutput,
AlignC,
ElementAccumulator,
ElementComputeEpilogue,
hytlass::epilogue::thread::ScaleType::NoBetaScaling>;
using GemmFunc = cutlass_patch::gemm::device::GemmUniversalWithVariadic<
ElementInputA,
hytlass::layout::RowMajor,
ElementInputB,
hytlass::layout::RowMajor,
ElementOutput,
hytlass::layout::RowMajor,
ElementAccumulator,
hytlass::arch::OpClassTensorOp,
hytlass::arch::Gfx928,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
TShape,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
WShape,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
IShape,
EpilogueOutputOp,
typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
SwizzleThreadBlock,
GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::kNumStages,
AlignA,
AlignB,
typename GemmOperation<ElementT>::Type>;
CHECK_HYTLASS(SetMaxDynamicSharedMemorySize<GemmFunc>());
/// Arguments
hytlass::gemm::GemmCoord problem_size{params.m, params.n, params.k};
const ElementInputA *input =
reinterpret_cast<const ElementInputA *>(params.input);
const ElementInputB *weight =
reinterpret_cast<const ElementInputB *>(params.weight);
const ElementOutput *bias =
reinterpret_cast<const ElementOutput *>(params.bias);
ElementOutput *output = reinterpret_cast<ElementOutput *>(params.output);
ElementComputeEpilogue alpha = static_cast<ElementComputeEpilogue>(1);
ElementComputeEpilogue beta = bias ? static_cast<ElementComputeEpilogue>(1)
: static_cast<ElementComputeEpilogue>(0);
typename GemmFunc::Arguments arguments{
GetGemmMode(params.batch_count),
problem_size, // <- problem size of matrix multiplication
params.batch_count, // <- batch_count or k-dimension split factor
{alpha, beta, variadic_args}, // <- epilogue params, alpha, beta
input, // <- input, ptr_A, A, shape={M, K}
weight, // <- input, ptr_B, B, shape={K, N}
bias, // <- input, ptr_C, shape={M, N} or {1, N}
output, // <- output, ptr_D, Z, shape={M, N}
params.shape_args.batch_stride_A,
params.shape_args.batch_stride_B,
params.shape_args.batch_stride_C,
params.shape_args.batch_stride_D,
params.shape_args.lda,
params.shape_args.ldb,
params.shape_args.ldc_bias,
params.shape_args.ldd};
size_t workspace_size = GemmFunc::get_workspace_size(arguments);
void *workspace = workspace_size > 0 ? GetWorkspace(workspace_size) : nullptr;
GemmFunc device_gemm;
hipStream_t *stream_ptr = reinterpret_cast<hipStream_t *>(params.stream_ptr);
CHECK_HYTLASS(device_gemm.can_implement(arguments));
CHECK_HYTLASS(device_gemm.initialize(arguments, workspace, *stream_ptr));
// Run the GEMM
CHECK_HYTLASS(device_gemm(*stream_ptr));
#if AP_ENABLE_DEBUG
CHECK_HIP(hipStreamSynchronize(*stream_ptr));
#endif
}
} // namespace ap
python/paddle/apy/matmul_pass/matmul/cutlass_patch/trace_device.h
+73
View File
@@ -0,0 +1,73 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#if CUTLASS_DEBUG_TRACE_LEVEL
#ifndef CUTLASS_TRACE_DEVICE
#define CUTLASS_TRACE_DEVICE(format, ...) \
{ \
if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && \
threadIdx.x == 0 && threadIdx.y == 0) { \
printf("[DEVICE][%s:%d, %s]" format "\n", \
__FILE__, \
__LINE__, \
__FUNCTION__, \
##__VA_ARGS__); \
} \
}
#endif
#ifndef CUTLASS_TRACE_DEVICE_TID_DETAIL
#define CUTLASS_TRACE_DEVICE_TID_DETAIL(bidz, bidx, tidx, format, ...) \
{ \
if (blockIdx.x == bidx && blockIdx.y == 0 && blockIdx.z == bidz && \
threadIdx.x == tidx && threadIdx.y == 0) { \
printf("[DEVICE][%s:%d, %s][bid={%d,%d,%d}, tid={%d,%d,%d}]" format \
"\n", \
__FILE__, \
__LINE__, \
__FUNCTION__, \
blockIdx.x, \
blockIdx.y, \
blockIdx.z, \
threadIdx.x, \
threadIdx.y, \
threadIdx.z, \
##__VA_ARGS__); \
} \
}
#endif
#ifndef CUTLASS_TRACE_DEVICE_TID
#define CUTLASS_TRACE_DEVICE_TID(format, ...) \
{ \
CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 0, 0, format, ##__VA_ARGS__) \
CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 0, 1, format, ##__VA_ARGS__) \
CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 1, 0, format, ##__VA_ARGS__) \
}
#endif
#else
#ifndef CUTLASS_TRACE_DEVICE
#define CUTLASS_TRACE_DEVICE(format, ...)
#endif
#ifndef CUTLASS_TRACE_DEVICE_TID
#define CUTLASS_TRACE_DEVICE_TID(format, ...)
#endif
#endif