releaase 2.11 (#703)
This commit is contained in:
Aditya Atluri
63
examples/44_multi_gemm_ir_and_codegen/README.md
Normal file
63
examples/44_multi_gemm_ir_and_codegen/README.md
Normal file
@@ -0,0 +1,63 @@
|
||||
This example provides utilities for generating back-to-back (B2B) GEMMs using CUTLASS.
|
||||
|
||||
## Quick start
|
||||
A configuration file containing the GEMMs to be fused together is located in [config.json](config.json). Edit
|
||||
this to change the configuration that you would like to run.
|
||||
```shell
|
||||
cd ir_gen
|
||||
|
||||
# Set up basic variables
|
||||
out_dir=directory_to_emit_files
|
||||
cutlass_dir=$(pwd)/../../..
|
||||
config_file=$(pwd)/../config.json
|
||||
|
||||
# Generate code for GEMMs described in `config_file`
|
||||
./generate.sh $config_file $out_dir $cutlass_dir
|
||||
|
||||
# Build the generated code
|
||||
cd $out_dir
|
||||
mkdir build && cd build
|
||||
cmake .. -DGPU_ARCHS="75;80"
|
||||
make -j
|
||||
|
||||
# Run the generated code with M=1024 K0=32 and Batch=1
|
||||
./sample 1024 32 1
|
||||
```
|
||||
|
||||
## Current restrictions
|
||||
This experimental example has the following restrictions:
|
||||
1. N tile should not exceed 256, or register spilling will occur.
|
||||
2. Only FP16 is supported currently
|
||||
3. Matrix A must be row major, matrix B must be column major, matrices C and D must be row major.
|
||||
|
||||
## Copyright
|
||||
|
||||
Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
```
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright notice, this
|
||||
list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
this list of conditions and the following disclaimer in the documentation
|
||||
and/or other materials provided with the distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived from
|
||||
this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
```
|
||||
32
examples/44_multi_gemm_ir_and_codegen/config.json
Normal file
32
examples/44_multi_gemm_ir_and_codegen/config.json
Normal file
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"0": {
|
||||
"A_tp": "fp16", "B_tp": "fp16", "C_tp": "fp16", "Acc_tp": "fp16",
|
||||
"A_format": "Row", "B_format": "Col", "C_format": "Row",
|
||||
"mnk": [15000, 256, 32],
|
||||
"epilogue": {
|
||||
"tp": "LeakyRelu",
|
||||
"bias": {"addbias": false, "bias_tp": "mat"},
|
||||
"args": [["float", "leaky_alpha", 1.3]]
|
||||
}
|
||||
},
|
||||
"1": {
|
||||
"A_tp": "fp16", "B_tp": "fp16", "C_tp": "fp16", "Acc_tp": "fp16",
|
||||
"A_format": "Row", "B_format": "Col", "C_format": "Row",
|
||||
"mnk": [15000, 128, 256],
|
||||
"epilogue": {
|
||||
"tp": "LeakyRelu",
|
||||
"bias": {"addbias": false, "bias_tp": "mat"},
|
||||
"args": [["float", "leaky_alpha", 1.3]]
|
||||
}
|
||||
},
|
||||
"2": {
|
||||
"A_tp": "fp16", "B_tp": "fp16", "C_tp": "fp16", "Acc_tp": "fp16",
|
||||
"A_format": "Row", "B_format": "Col", "C_format": "Row",
|
||||
"mnk": [15000, 64, 128],
|
||||
"epilogue": {
|
||||
"tp": "LeakyRelu",
|
||||
"bias": {"addbias": false, "bias_tp": "mat"},
|
||||
"args": [["float", "leaky_alpha", 1.3]]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief 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/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_complex_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator_mixed.h"
|
||||
|
||||
// #include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
#include "cutlass/epilogue/threadblock/interleaved_epilogue.h"
|
||||
|
||||
#include "fused_bias_act_epilogue.h"
|
||||
#include "../warp/fused_bias_act_fragment_iterator_tensor_op.h"
|
||||
#include "output_tile_thread_map_for_fused_bias.h"
|
||||
#include "default_thread_map_tensor_op_for_fused_bias.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for TensorOps.
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaTensorOp_,
|
||||
int PartitionsK,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultFusedBiasActEpilogueTensorOp {
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaTensorOp = WarpMmaTensorOp_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaTensorOp::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOpForFusedBias<
|
||||
Shape,
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
kPartitionsK,
|
||||
ElementOutput,
|
||||
kElementsPerAccess
|
||||
>::Type;
|
||||
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
OutputTileThreadMap,
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = typename std::conditional<is_complex<ElementOutput>::value,
|
||||
cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC>,
|
||||
cutlass::epilogue::warp::FusedBiasActFragmentIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC> >::type;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::FusedBiasActEpilogue<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
kPartitionsK,
|
||||
OutputTileIterator,
|
||||
AccumulatorFragmentIterator,
|
||||
OutputOp
|
||||
>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,113 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for TensorOp accumulator layouts
|
||||
template <
|
||||
typename ThreadblockShape_,
|
||||
typename WarpShape_,
|
||||
int PartitionsK,
|
||||
typename Element_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultThreadMapTensorOpForFusedBias {
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using Element = Element_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
/// Tensor Operations fundamentally perform operations on 8 rows
|
||||
static int const kTensorOpRows = 8;
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static_assert(
|
||||
!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM), "Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount = gemm::GemmShape<
|
||||
ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN,
|
||||
kPartitionsK
|
||||
>;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap
|
||||
using Type = OutputTileOptimalThreadMapBiasAct <
|
||||
OutputTileShape<ThreadblockShape::kN, Detail::kTensorOpRows, Detail::WarpCount::kM, 1, 1>,
|
||||
OutputTileShape<1, WarpShape::kM / Detail::kTensorOpRows, 1, 1, WarpShape::kM / Detail::kTensorOpRows>,
|
||||
Detail::kThreads,
|
||||
kElementsPerAccess,
|
||||
sizeof_bits<Element>::value
|
||||
>;
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,222 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief 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
|
||||
|
||||
#if defined(__CUDACC_RTC__)
|
||||
#include <cuda/std/cassert>
|
||||
#else
|
||||
#include <assert.h>
|
||||
#endif
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/vector.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/aligned_buffer.h"
|
||||
#include "cutlass/functional.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/predicated_tile_iterator.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Epilogue operator without splitk
|
||||
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 OutputOp_ ///< Output operator
|
||||
>
|
||||
class FusedBiasActEpilogue {
|
||||
|
||||
public:
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaOperator = WarpMmaOperator_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputTileIterator = OutputTileIterator_;
|
||||
using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
|
||||
using OutputOp = OutputOp_;
|
||||
|
||||
/// Output layout is always row-major
|
||||
using Layout = layout::RowMajor;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
/// The complete warp-level accumulator tile
|
||||
using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile;
|
||||
|
||||
/// Output element
|
||||
using ElementOutput = typename OutputTileIterator::Element;
|
||||
|
||||
/// Output access size
|
||||
static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
|
||||
|
||||
|
||||
public:
|
||||
|
||||
|
||||
static_assert(OutputTileIterator::kElementsPerAccess, "OutputTileIterator::kElementsPerAccess must not be zero.");
|
||||
|
||||
static_assert(!(OutputTileIterator::Fragment::kElements % OutputTileIterator::kElementsPerAccess),
|
||||
"Divisibility");
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
FusedBiasActEpilogue(
|
||||
){ }
|
||||
|
||||
/// Streams the result to global memory
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile
|
||||
AccumulatorTile & fused_bias_act_accumlators,
|
||||
OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
bool need_bias = output_op.is_source_needed();
|
||||
|
||||
if (need_bias)
|
||||
compute_source_needed_(output_op, accumulators, fused_bias_act_accumlators, source_iterator);
|
||||
else
|
||||
compute_source_no_needed_(output_op, accumulators, fused_bias_act_accumlators);
|
||||
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile
|
||||
AccumulatorTile & fused_bias_act_accumlators) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
compute_source_no_needed_(output_op, accumulators, fused_bias_act_accumlators);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void compute_source_needed_(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile
|
||||
AccumulatorTile & fused_bias_act_accumlators,
|
||||
OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
typename OutputTileIterator::Fragment source_fragment;
|
||||
|
||||
|
||||
source_fragment.clear();
|
||||
|
||||
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
|
||||
AccumulatorFragmentIterator fused_bias_act_fragment_iterator(fused_bias_act_accumlators);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
|
||||
|
||||
source_iterator.load(source_fragment);
|
||||
++source_iterator;
|
||||
|
||||
typename AccumulatorFragmentIterator::Fragment accum_fragment;
|
||||
|
||||
accum_fragment_iterator.load(accum_fragment);
|
||||
++accum_fragment_iterator;
|
||||
|
||||
typename AccumulatorFragmentIterator::Fragment fused_bias_act_fragment;
|
||||
fused_bias_act_fragment = output_op(accum_fragment, source_fragment);
|
||||
|
||||
fused_bias_act_fragment_iterator.store(fused_bias_act_fragment);
|
||||
++fused_bias_act_fragment_iterator;
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void compute_source_no_needed_(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile
|
||||
AccumulatorTile & fused_bias_act_accumlators) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
|
||||
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
|
||||
AccumulatorFragmentIterator fused_bias_act_fragment_iterator(fused_bias_act_accumlators);
|
||||
|
||||
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int iter = 0; iter < AccumulatorFragmentIterator::kIterations; ++iter) {
|
||||
|
||||
typename AccumulatorFragmentIterator::Fragment accum_fragment;
|
||||
|
||||
accum_fragment_iterator.load(accum_fragment);
|
||||
++accum_fragment_iterator;
|
||||
|
||||
typename AccumulatorFragmentIterator::Fragment fused_bias_act_fragment;
|
||||
fused_bias_act_fragment = output_op(accum_fragment);
|
||||
|
||||
fused_bias_act_fragment_iterator.store(fused_bias_act_fragment);
|
||||
++fused_bias_act_fragment_iterator;
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,311 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Metaprogram for determining the mapping of output elements to threads for epilogue tiles.
|
||||
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
/// RowArrangement determines how one or more warps cover a region of consecutive rows.
|
||||
template <
|
||||
typename Shape,
|
||||
int WarpsRemaining,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize,
|
||||
bool Is2dTile
|
||||
>
|
||||
struct RowArrangementBiasAct;
|
||||
|
||||
/// RowArrangement in which each warp's access is a 1D tiled arrangement.
|
||||
template <
|
||||
typename Shape,
|
||||
int WarpsRemaining,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize
|
||||
>
|
||||
struct RowArrangementBiasAct<Shape, WarpsRemaining, ElementsPerAccess, ElementSize, false> {
|
||||
static int const kWarpSize = 32;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
static int const kIterationsRow = 1;
|
||||
static int const kDeltaRow = 1;
|
||||
static int const kIterationsColumn = Shape::kColumn / kElementsPerAccess / kWarpSize;
|
||||
static int const kDeltaColumn = kWarpSize * kElementsPerAccess;
|
||||
|
||||
static int const kAccessWidth = kWarpSize;
|
||||
static int const kAccessRows = 1;
|
||||
static int const kWarpPartitionsRow = 1;
|
||||
static int const kWarpPartitionsColumn = WarpsRemaining;
|
||||
};
|
||||
|
||||
/// RowArrangement in which each warp's access is a 2D tiled arrangement.
|
||||
template <
|
||||
typename Shape,
|
||||
int WarpsRemaining,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize
|
||||
>
|
||||
struct RowArrangementBiasAct<Shape, WarpsRemaining, ElementsPerAccess, ElementSize, true> {
|
||||
|
||||
static int const kMemoryAccessSize = 4;//128;
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
struct Detail {
|
||||
static int const kShapeRow = Shape::kRow / WarpsRemaining;
|
||||
static int const kShapeWidth = Shape::kColumn / kElementsPerAccess;
|
||||
|
||||
static int const kTargetMemoryAccessWidth =
|
||||
kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8);
|
||||
|
||||
static int const kTargetAccessRows = kWarpSize / kTargetMemoryAccessWidth;
|
||||
};
|
||||
|
||||
static int const kAccessWidth =
|
||||
(Detail::kTargetAccessRows > Detail::kShapeRow ?
|
||||
kWarpSize / Detail::kShapeRow
|
||||
: const_min(
|
||||
Detail::kShapeWidth,
|
||||
const_min(kWarpSize, kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8))
|
||||
));
|
||||
|
||||
static int const kAccessRows =
|
||||
(Detail::kTargetAccessRows > Detail::kShapeRow ?
|
||||
Detail::kShapeRow
|
||||
: const_min(Shape::kRow, kWarpSize / kAccessWidth));
|
||||
|
||||
static int const kIterationsRow = Detail::kShapeRow / kAccessRows;
|
||||
static int const kDeltaRow = kAccessRows;
|
||||
|
||||
static int const kIterationsColumn = Detail::kShapeWidth / kAccessWidth;
|
||||
static int const kDeltaColumn = kAccessWidth * kElementsPerAccess;
|
||||
|
||||
static_assert( kAccessWidth * kElementsPerAccess <= Shape::kColumn, "Accessing too many elements per access");
|
||||
static_assert( kIterationsColumn > 0, "Iteration Count Column must be > 0" );
|
||||
static_assert( kIterationsRow > 0, "Iteration Count Row must be > 0" );
|
||||
|
||||
static int const kWarpPartitionsRow = 1;
|
||||
static int const kWarpPartitionsColumn = 1;
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Template metaprogram for partitioning a 4D space across warps to achieve several performance
|
||||
/// objectives:
|
||||
///
|
||||
/// - coalesced memory accesses in units of 16 Byte lines
|
||||
/// - minimal address arithmetic
|
||||
/// - minimal predicate calculations
|
||||
///
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Count_,
|
||||
int Threads,
|
||||
int ElementsPerAccess,
|
||||
int ElementSize
|
||||
>
|
||||
struct OutputTileOptimalThreadMapBiasAct {
|
||||
|
||||
using Shape = Shape_;
|
||||
using Count = Count_;
|
||||
|
||||
static int const kWarpSize = 32;
|
||||
static int const kThreads = Threads;
|
||||
static int const kWarpCount = kThreads / kWarpSize;
|
||||
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
//
|
||||
// Metaprogram computation
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
|
||||
// Clusters
|
||||
static int const kIterationsCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
Shape::kCluster / kWarpCount
|
||||
: 1);
|
||||
|
||||
static int const kDeltaCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup * Shape::kCluster / kIterationsCluster
|
||||
: 1);
|
||||
|
||||
static int const kCompactedDeltaCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
Shape::kRow * Shape::kGroup * Shape::kCluster / kIterationsCluster
|
||||
: 1);
|
||||
|
||||
static int const kWarpPartitionsCluster =
|
||||
((Shape::kCluster > kWarpCount) ?
|
||||
kWarpCount
|
||||
: kWarpCount / Shape::kCluster);
|
||||
|
||||
static int const kWarpsRemainingForGroups =
|
||||
((Shape::kCluster > kWarpCount) ? 1 : kWarpCount / Shape::kCluster);
|
||||
|
||||
// Groups
|
||||
static int const kIterationsGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
Shape::kGroup / kWarpsRemainingForGroups
|
||||
: 1);
|
||||
|
||||
static int const kDeltaGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
Shape::kRow * Count::kRow * Shape::kGroup / kIterationsGroup
|
||||
: 1);
|
||||
|
||||
static int const kCompactedDeltaGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
Shape::kRow * Shape::kGroup / kIterationsGroup
|
||||
: 1);
|
||||
|
||||
static int const kWarpPartitionsGroup =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
1
|
||||
: kWarpsRemainingForGroups / Shape::kGroup);
|
||||
|
||||
static int const kWarpsRemainingForRows =
|
||||
((Shape::kGroup > kWarpsRemainingForGroups) ?
|
||||
1
|
||||
: kWarpsRemainingForGroups / Shape::kGroup);
|
||||
|
||||
// Rows
|
||||
using RowArrangement = detail::RowArrangementBiasAct<
|
||||
Shape,
|
||||
kWarpsRemainingForRows,
|
||||
kElementsPerAccess,
|
||||
kElementSize,
|
||||
(Shape::kRow > kWarpsRemainingForRows)
|
||||
>;
|
||||
|
||||
// Warp partitions
|
||||
using WarpPartitions = OutputTileShape<
|
||||
RowArrangement::kWarpPartitionsColumn,
|
||||
RowArrangement::kWarpPartitionsRow,
|
||||
kWarpPartitionsGroup,
|
||||
kWarpPartitionsCluster,
|
||||
1>;
|
||||
|
||||
static int const kAccessWidth = RowArrangement::kAccessWidth;
|
||||
static int const kAccessRows = RowArrangement::kAccessRows;
|
||||
};
|
||||
|
||||
//
|
||||
// Output
|
||||
//
|
||||
|
||||
using Iterations = OutputTileShape<
|
||||
Detail::RowArrangement::kIterationsColumn,
|
||||
Detail::RowArrangement::kIterationsRow,
|
||||
Detail::kIterationsGroup,
|
||||
Detail::kIterationsCluster,
|
||||
1>;
|
||||
|
||||
using Delta = OutputTileShape<
|
||||
Detail::RowArrangement::kDeltaColumn,
|
||||
Detail::RowArrangement::kDeltaRow,
|
||||
Detail::kDeltaGroup,
|
||||
Detail::kDeltaCluster,
|
||||
1>;
|
||||
|
||||
/// Initial offset function
|
||||
CUTLASS_HOST_DEVICE
|
||||
static MatrixCoord initial_offset(int thread_idx) {
|
||||
|
||||
int warp_idx = thread_idx / kWarpSize;
|
||||
int lane_idx = thread_idx % kWarpSize;
|
||||
|
||||
// Compute warp location
|
||||
int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster;
|
||||
int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster;
|
||||
|
||||
int group_idx = residual_cluster / Detail::WarpPartitions::kGroup;
|
||||
int residual_group = residual_cluster % Detail::WarpPartitions::kGroup;
|
||||
|
||||
int row_idx = residual_group / Detail::WarpPartitions::kRow;
|
||||
int col_idx = residual_group % Detail::WarpPartitions::kRow;
|
||||
|
||||
// Compute per-lane offset
|
||||
int lane_row_offset = lane_idx / Detail::kAccessWidth;
|
||||
int lane_col_offset = lane_idx % Detail::kAccessWidth;
|
||||
|
||||
// Compute coordinate in output space
|
||||
int cluster_offset = cluster_idx * Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup;
|
||||
int group_offset = group_idx * Shape::kRow * Count::kRow;
|
||||
int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows;
|
||||
int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess;
|
||||
|
||||
return MatrixCoord(
|
||||
cluster_offset + group_offset + row_offset + lane_row_offset,
|
||||
(column_offset + lane_col_offset) * kElementsPerAccess
|
||||
);
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
@@ -0,0 +1,189 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
|
||||
that participate in one warp-level store operation.
|
||||
|
||||
Typically, the accumulator tile is the largest single block of register-backed storage
|
||||
within the kernel. Storing it to memory is best accomplished by partitioning it into
|
||||
smaller tiles and storing these sequentially.
|
||||
|
||||
Round trips through shared memory during the Epilogue phase require partitioning, as
|
||||
shared memory capacity is typically insufficient for a threadblock's total accumulator
|
||||
size.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/tensor_op_policy.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace warp {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
///
|
||||
template <
|
||||
typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape)
|
||||
typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape)
|
||||
typename OperatorElementC, ///< matrix multiply operation data type (concept: data type)
|
||||
typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array)
|
||||
typename Layout ///< target shared memory layout
|
||||
>
|
||||
class FusedBiasActFragmentIteratorTensorOp;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for row-major shared memory
|
||||
template <
|
||||
typename WarpShape_, ///< shape of the warp-level GEMM tile
|
||||
typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape)
|
||||
typename OperatorElementC_, ///< matrix multiply operation data type (concept: data type)
|
||||
typename OperatorFragmentC_ ///< matrix multiply operation fragment (concept: Array)
|
||||
>
|
||||
class FusedBiasActFragmentIteratorTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_, layout::RowMajor> {
|
||||
public:
|
||||
|
||||
using WarpShape = WarpShape_;
|
||||
using OperatorShape = OperatorShape_;
|
||||
using OperatorElementC = OperatorElementC_;
|
||||
using OperatorFragmentC = OperatorFragmentC_;
|
||||
using Layout = layout::RowMajor;
|
||||
|
||||
using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
|
||||
|
||||
/// This is the fragment size produced by one access of the iterator.
|
||||
using Fragment = Array<
|
||||
OperatorElementC,
|
||||
Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
|
||||
|
||||
/// This is the complete warp-level accumulator tile.
|
||||
using AccumulatorTile = Array<
|
||||
OperatorElementC,
|
||||
OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn>;
|
||||
|
||||
using OutputAccumulatorTile = AccumulatorTile;
|
||||
|
||||
/// Number of times this iterator can be incremented
|
||||
static int const kIterations = Policy::kIterations;
|
||||
|
||||
private:
|
||||
|
||||
/// Internal access type
|
||||
using AccessType = Array<OperatorElementC, Policy::kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Accumulator tile
|
||||
AccessType *accumulators_;
|
||||
|
||||
/// Internal index
|
||||
int index_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs an iterator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FusedBiasActFragmentIteratorTensorOp(AccumulatorTile &accum):
|
||||
accumulators_(reinterpret_cast<AccessType *>(&accum)),
|
||||
index_(0) {
|
||||
}
|
||||
|
||||
/// Increments
|
||||
CUTLASS_HOST_DEVICE
|
||||
FusedBiasActFragmentIteratorTensorOp &operator++() {
|
||||
++index_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Decrements
|
||||
CUTLASS_HOST_DEVICE
|
||||
FusedBiasActFragmentIteratorTensorOp &operator--() {
|
||||
--index_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Loads a fragment from the referenced part of the accumulator tile
|
||||
CUTLASS_HOST_DEVICE
|
||||
void load(Fragment &frag, int index_offset = 0) const {
|
||||
|
||||
int index = index_ + index_offset;
|
||||
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
|
||||
|
||||
int accumulator_access_offset =
|
||||
index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess;
|
||||
|
||||
frag_ptr[n] = accumulators_[accumulator_access_offset];
|
||||
}
|
||||
}
|
||||
/// Stores a fragment from the referenced part of the accumulator tile
|
||||
CUTLASS_HOST_DEVICE
|
||||
void store(Fragment &frag, int index_offset = 0) const {
|
||||
|
||||
int index = index_ + index_offset;
|
||||
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
|
||||
|
||||
int accumulator_access_offset =
|
||||
index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess;
|
||||
|
||||
accumulators_[accumulator_access_offset] = frag_ptr[n];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,427 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace warp {
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
/// Size of the matrix to load (concept: MatrixShape)
|
||||
typename Shape_,
|
||||
/// Size of the accumulation tile shape (concept: MatrixShape)
|
||||
typename AccumulatorShape_,
|
||||
/// KBlocks columns to compute residual
|
||||
int KBlocksColumn_,
|
||||
/// Accumulator Element type
|
||||
typename ElementAccumulator_,
|
||||
/// Element type
|
||||
typename Element_,
|
||||
/// Layout of operand in memory
|
||||
typename Layout_,
|
||||
/// Shape of one matrix product operation (concept: MatrixShape)
|
||||
typename InstructionShape_,
|
||||
/// Whether beta is zero
|
||||
bool IsBetaZero_ >
|
||||
class MmaTensorOpPureFragmentIterator;
|
||||
|
||||
|
||||
// Partial specialization for col-major accumulator tile
|
||||
// And Element type is the same as Accumulator Element type
|
||||
|
||||
template <
|
||||
/// Shape of warp tile to load (concept: MatrixShape)
|
||||
typename Shape_,
|
||||
/// Shape of the warp accumulation tile (concept: MatrixShape)
|
||||
typename AccumulatorShape_,
|
||||
/// KBlocks columns to compute residual
|
||||
int KBlocksColumn_,
|
||||
/// Element type
|
||||
typename Element_,
|
||||
/// Shape of one matrix product operation (concept: MatrixShape)
|
||||
typename InstructionShape_>
|
||||
class MmaTensorOpPureFragmentIterator<Shape_, AccumulatorShape_, KBlocksColumn_, Element_, Element_,
|
||||
cutlass::layout::ColumnMajor,
|
||||
InstructionShape_, true> {
|
||||
public:
|
||||
|
||||
/// Shape of warp tile to load (concept: MatrixShape)
|
||||
using Shape = Shape_;
|
||||
|
||||
/// Shape of the warp accumulation tile (concept: MatrixShape)
|
||||
using AccumulatorShape = AccumulatorShape_;
|
||||
|
||||
/// KBlocks columns to compute residual
|
||||
static int const kKBlockColumn = KBlocksColumn_;
|
||||
|
||||
/// Element type
|
||||
using Element = Element_;
|
||||
|
||||
/// Layout of source tile
|
||||
using Layout = cutlass::layout::ColumnMajor;
|
||||
|
||||
/// Shape of one matrix product operation (concept: MatrixShape)
|
||||
using InstructionShape = InstructionShape_;
|
||||
|
||||
/// Whether beta is zero
|
||||
static bool const IsBetaZero = true;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = 32;
|
||||
|
||||
/// Internal structure of iterator - made public to enable introspection
|
||||
struct Policy {
|
||||
static_assert(
|
||||
!(Shape::kRow % InstructionShape::kM) &&
|
||||
!(Shape::kColumn % InstructionShape::kN),
|
||||
"Shape of warp-level Mma must be divisible by operator shape.");
|
||||
static_assert(
|
||||
!(AccumulatorShape::kRow % Shape::kRow) &&
|
||||
!(AccumulatorShape::kColumn % Shape::kColumn),
|
||||
"Shape of Warp Accumulator must be divisible by warp shape.");
|
||||
static_assert(
|
||||
!(kKBlockColumn % Shape::kColumn),
|
||||
"KBlock size must be divisible by warp shape.");
|
||||
|
||||
/// Number of times this iterator can be incremented
|
||||
static int const kIterations = AccumulatorShape::kCount / Shape::kCount;
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
static int const kElementsPerAccess = InstructionShape::kM * InstructionShape::kN / kThreads;
|
||||
|
||||
/// Number of mma operations performed by a warp
|
||||
using MmaIterations = MatrixShape<Shape::kRow / InstructionShape::kM,
|
||||
Shape::kColumn / InstructionShape::kN>;
|
||||
/// Number of mma operations performed by the entire accumulator
|
||||
using AccumulatorIterations = MatrixShape<AccumulatorShape::kRow / InstructionShape::kM,
|
||||
AccumulatorShape::kColumn / InstructionShape::kN>;
|
||||
|
||||
/// Number of K iterations
|
||||
static int const kKBlockIterations = (AccumulatorShape::kColumn + kKBlockColumn - 1) / kKBlockColumn;
|
||||
static int const kResidualColumn = AccumulatorShape::kColumn - (kKBlockIterations - 1) * kKBlockColumn;
|
||||
static int const kKBlockColumnIterations = kKBlockColumn / Shape::kColumn
|
||||
* (AccumulatorShape::kRow / Shape::kRow);
|
||||
static int const kResidualIndex = kResidualColumn / Shape::kColumn
|
||||
* (AccumulatorShape::kRow / Shape::kRow);
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Derived quantities
|
||||
//
|
||||
|
||||
/// Fragment object holding a thread's part of a tile
|
||||
/// This is the fragment size produced by one access of the iterator.
|
||||
using Fragment = Array<Element, Shape::kCount / kThreads>;
|
||||
|
||||
/// Accumulator Fragment object
|
||||
using AccumulatorFragment = Array<Element, AccumulatorShape::kCount / kThreads>;
|
||||
|
||||
|
||||
private:
|
||||
|
||||
/// Internal access type
|
||||
using AccessType = Array<Element, kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Accumulator tile
|
||||
AccessType const *accumulators_;
|
||||
|
||||
/// Internal index
|
||||
int index_;
|
||||
|
||||
/// Used to access residual tile first
|
||||
bool is_residual_tile_;
|
||||
|
||||
public:
|
||||
/// Constructs an iterator
|
||||
CUTLASS_HOST_DEVICE
|
||||
MmaTensorOpPureFragmentIterator(AccumulatorFragment const &accum)
|
||||
: accumulators_(reinterpret_cast<AccessType const *>(&accum)),
|
||||
index_(0), is_residual_tile_(true) {}
|
||||
|
||||
/// Add offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_offset(int index_offset) {
|
||||
index_ += index_offset;
|
||||
if(is_residual_tile_ && index_ >= kKBlockColumnIterations) {
|
||||
index_ = index_ - kKBlockColumnIterations + kResidualIndex;
|
||||
is_residual_tile_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Increments
|
||||
CUTLASS_HOST_DEVICE
|
||||
MmaTensorOpPureFragmentIterator &operator++() {
|
||||
add_offset(1);
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Decrements
|
||||
CUTLASS_HOST_DEVICE
|
||||
MmaTensorOpPureFragmentIterator &operator--() {
|
||||
add_offset(-1);
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Loads a fragment from the referenced part of the accumulator tile
|
||||
CUTLASS_HOST_DEVICE
|
||||
void load(Fragment &frag) const {
|
||||
|
||||
AccessType src_fragment;
|
||||
src_fragment.clear();
|
||||
|
||||
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
int index_m = (index_ * MmaIterations::kRow) % AccumulatorIterations::kRow;
|
||||
int index_n = (index_ * MmaIterations::kRow) / AccumulatorIterations::kRow
|
||||
* MmaIterations::kColumn;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < MmaIterations::kColumn; n++) {
|
||||
for (int m = 0; m < MmaIterations::kRow; m++) {
|
||||
int accumulator_access_offset =
|
||||
(n + index_n) * AccumulatorIterations::kRow + m + index_m;
|
||||
|
||||
frag_ptr[n * MmaIterations::kRow + m].clear();
|
||||
if(!(is_residual_tile_ && index_ >= kResidualIndex))
|
||||
frag_ptr[n * MmaIterations::kRow + m] = accumulators_[accumulator_access_offset];
|
||||
// frag_ptr[n * MmaIterations::kRow + m] = output_op(accumulators_[accumulator_access_offset], src_fragment);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
// Partial specialization for row-major accumulator tile
|
||||
|
||||
template <
|
||||
/// Shape of warp tile to load (concept: MatrixShape)
|
||||
typename Shape_,
|
||||
/// Shape of the warp accumulation tile (concept: MatrixShape)
|
||||
typename AccumulatorShape_,
|
||||
/// KBlocks columns to compute residual
|
||||
int KBlocksColumn_,
|
||||
/// Accumulator Element type
|
||||
typename ElementAccumulator_,
|
||||
/// Element type
|
||||
typename Element_,
|
||||
/// Shape of one matrix product operation (concept: MatrixShape)
|
||||
typename InstructionShape_>
|
||||
class MmaTensorOpPureFragmentIterator<Shape_, AccumulatorShape_, KBlocksColumn_, ElementAccumulator_, Element_,
|
||||
cutlass::layout::RowMajor,
|
||||
InstructionShape_, true> {
|
||||
public:
|
||||
|
||||
/// Shape of warp tile to load (concept: MatrixShape)
|
||||
using Shape = Shape_;
|
||||
|
||||
/// Shape of the warp accumulation tile (concept: MatrixShape)
|
||||
using AccumulatorShape = AccumulatorShape_;
|
||||
|
||||
/// KBlocks columns to compute residual
|
||||
static int const kKBlockColumn = KBlocksColumn_;
|
||||
|
||||
/// Accumulator Element type
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
|
||||
/// Element type
|
||||
using Element = Element_;
|
||||
|
||||
/// Layout of source tile
|
||||
using Layout = cutlass::layout::RowMajor;
|
||||
|
||||
/// Shape of one matrix product operation (concept: MatrixShape)
|
||||
using InstructionShape = InstructionShape_;
|
||||
|
||||
/// Whether beta is zero
|
||||
static bool const IsBetaZero = true;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = 32;
|
||||
|
||||
/// Internal structure of iterator - made public to enable introspection
|
||||
struct Policy {
|
||||
static_assert(
|
||||
!(Shape::kRow % InstructionShape::kM) &&
|
||||
!(Shape::kColumn % InstructionShape::kN),
|
||||
"Shape of warp-level Mma must be divisible by operator shape.");
|
||||
static_assert(
|
||||
!(AccumulatorShape::kRow % Shape::kRow) &&
|
||||
!(AccumulatorShape::kColumn % Shape::kColumn),
|
||||
"Shape of Warp Accumulator must be divisible by warp shape.");
|
||||
static_assert(
|
||||
!(kKBlockColumn % Shape::kColumn),
|
||||
"KBlock size must be divisible by warp shape.");
|
||||
|
||||
/// Number of times this iterator can be incremented
|
||||
static int const kIterations = AccumulatorShape::kCount / Shape::kCount;
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
static int const kElementsPerAccess = InstructionShape::kM * InstructionShape::kN / kThreads;
|
||||
|
||||
/// Number of mma operations performed by a warp
|
||||
using MmaIterations = MatrixShape<Shape::kRow / InstructionShape::kM,
|
||||
Shape::kColumn / InstructionShape::kN>;
|
||||
/// Number of mma operations performed by the entire accumulator
|
||||
using AccumulatorIterations = MatrixShape<AccumulatorShape::kRow / InstructionShape::kM,
|
||||
AccumulatorShape::kColumn / InstructionShape::kN>;
|
||||
|
||||
/// Number of K iterations
|
||||
static int const kKBlockIterations = (AccumulatorShape::kColumn + kKBlockColumn - 1) / kKBlockColumn;
|
||||
static int const kResidualColumn = AccumulatorShape::kColumn - (kKBlockIterations - 1) * kKBlockColumn;
|
||||
static int const kKBlockColumnIterations = kKBlockColumn / Shape::kColumn
|
||||
* (AccumulatorShape::kRow / Shape::kRow);
|
||||
static int const kResidualIndex = kResidualColumn / Shape::kColumn
|
||||
* (AccumulatorShape::kRow / Shape::kRow);
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Derived quantities
|
||||
//
|
||||
|
||||
/// Fragment object holding a thread's part of a tile
|
||||
/// This is the fragment size produced by one access of the iterator.
|
||||
using Fragment = Array<Element, Shape::kCount / kThreads>;
|
||||
|
||||
/// Accumulator Fragment object
|
||||
using AccumulatorFragment = Array<ElementAccumulator, AccumulatorShape::kCount / kThreads>;
|
||||
|
||||
|
||||
private:
|
||||
|
||||
/// Internal access type
|
||||
using AccessType = Array<ElementAccumulator, kElementsPerAccess>;
|
||||
using FragmentAccessType = Array<Element, kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Accumulator tile
|
||||
AccessType const *accumulators_;
|
||||
|
||||
/// Internal index
|
||||
int index_;
|
||||
|
||||
/// Used to access residual tile first
|
||||
bool is_residual_tile_;
|
||||
|
||||
public:
|
||||
/// Constructs an iterator
|
||||
CUTLASS_HOST_DEVICE
|
||||
MmaTensorOpPureFragmentIterator(AccumulatorFragment const &accum)
|
||||
: accumulators_(reinterpret_cast<AccessType const *>(&accum)),
|
||||
index_(0), is_residual_tile_(true) {}
|
||||
|
||||
/// Add offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_offset(int index_offset) {
|
||||
index_ += index_offset;
|
||||
if(is_residual_tile_ && index_ >= kKBlockColumnIterations) {
|
||||
index_ = index_ - kKBlockColumnIterations + kResidualIndex;
|
||||
is_residual_tile_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Increments
|
||||
CUTLASS_HOST_DEVICE
|
||||
MmaTensorOpPureFragmentIterator &operator++() {
|
||||
add_offset(1);
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Decrements
|
||||
CUTLASS_HOST_DEVICE
|
||||
MmaTensorOpPureFragmentIterator &operator--() {
|
||||
add_offset(-1);
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Loads a fragment from the referenced part of the accumulator tile
|
||||
CUTLASS_HOST_DEVICE
|
||||
void load(Fragment &frag) const {
|
||||
|
||||
|
||||
FragmentAccessType src_fragment;
|
||||
src_fragment.clear();
|
||||
|
||||
FragmentAccessType *frag_ptr = reinterpret_cast<FragmentAccessType *>(&frag);
|
||||
|
||||
int index_m = (index_ * MmaIterations::kRow) % AccumulatorIterations::kRow;
|
||||
int index_n = (index_ * MmaIterations::kRow) / AccumulatorIterations::kRow
|
||||
* MmaIterations::kColumn;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < MmaIterations::kRow; m++) {
|
||||
for (int n = 0; n < MmaIterations::kColumn; n++) {
|
||||
int accumulator_access_offset =
|
||||
(m + index_m) * AccumulatorIterations::kColumn + n + index_n;
|
||||
|
||||
frag_ptr[m * MmaIterations::kColumn + n].clear();
|
||||
if(!(is_residual_tile_ && index_ >= kResidualIndex))
|
||||
frag_ptr[m * MmaIterations::kColumn + n] = (accumulators_[accumulator_access_offset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
129
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_all_code.py
Normal file
129
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_all_code.py
Normal file
@@ -0,0 +1,129 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import gen_turing_and_volta as api_generator
|
||||
import gen_sample as sample_creater
|
||||
import gen_cmake as cmake_creater
|
||||
import gen_verify as verify_creater
|
||||
import gen_device as b2b_fused_generator
|
||||
import replace_fix_impl_header
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description="Generates Fused Multi-GEMM CUTLASS Kernels")
|
||||
parser.add_argument("--config-file", default="config.json", help="JSON file containing configuration to generate")
|
||||
parser.add_argument("--gen-name", default="FusedMultiGemmForward", help="Specific the output name")
|
||||
parser.add_argument("--output-dir", default="", help="Specifies the output dir")
|
||||
parser.add_argument("--cutlass-dir", default="", help="Specifies the dependent CUTLASS repo dir")
|
||||
parser.add_argument("--gen-include-cutlass-dir", default="", help="Specifies the generated CUTLASS code include dir, if needed.")
|
||||
args = parser.parse_args()
|
||||
|
||||
gen_name = args.gen_name
|
||||
|
||||
cutlass_deps_dir = args.cutlass_dir
|
||||
|
||||
output_dir = args.output_dir
|
||||
output_dir += "/"
|
||||
|
||||
cutlass_deps_root = args.gen_include_cutlass_dir
|
||||
if cutlass_deps_root == '':
|
||||
cutlass_deps_root = cutlass_deps_dir + "/include/"
|
||||
cutlass_deps_root +='/'
|
||||
|
||||
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen")
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "fixed_impl"):
|
||||
os.mkdir(output_dir + "/" + "fixed_impl" )
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "sample"):
|
||||
os.mkdir(output_dir + "/" + "sample" )
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen" + "/" + "device"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen" + "/" + "device")
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen" + "/" + "kernel"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen" + "/" + "kernel")
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen" + "/" + "threadblock"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen" + "/" + "threadblock")
|
||||
|
||||
with open(args.config_file, 'r') as infile:
|
||||
gemm_info_dict = json.load(infile)
|
||||
|
||||
keys = sorted(gemm_info_dict.keys())
|
||||
fuse_gemm_info = [gemm_info_dict[k] for k in keys]
|
||||
|
||||
|
||||
for_cutlass_gen_user_include_header_file = [
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination_leaky_relu.h",
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination.h",
|
||||
]
|
||||
|
||||
for_fused_wrapper = [
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination_leaky_relu.h",
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination.h",
|
||||
"auto_gen/device/" + gen_name + ".h",
|
||||
cutlass_deps_root + "cutlass/gemm/device/gemm_batched.h",
|
||||
cutlass_deps_root + "cutlass/cutlass.h",
|
||||
]
|
||||
|
||||
# Copy fixed implementation to the output directory
|
||||
fix_impl = replace_fix_impl_header.replace_fix_impl("../fixed_impl/", output_dir +"/fixed_impl/", cutlass_deps_root)
|
||||
fix_impl.gen_code()
|
||||
|
||||
auto_gen_output_dir = output_dir + "/auto_gen/"
|
||||
project_root = ""
|
||||
turing_plus = b2b_fused_generator.gen_device(fuse_gemm_info, gen_name, for_cutlass_gen_user_include_header_file, cutlass_deps_root, project_root, auto_gen_output_dir)
|
||||
turing_plus.gen_code(75, 'hmma1688', False)
|
||||
|
||||
api = api_generator.gen_one_API(fuse_gemm_info, gen_name, for_fused_wrapper, output_dir)
|
||||
api.gen_code()
|
||||
|
||||
# Generate C++ sample
|
||||
os.system("cp ../leaky_bias.h " + output_dir + "/sample/")
|
||||
os.system("cp ../utils.h " + output_dir + "/sample/")
|
||||
|
||||
sample_dir = output_dir + "/sample/"
|
||||
sample = sample_creater.gen_test(fuse_gemm_info, gen_name, for_cutlass_gen_user_include_header_file, sample_dir)
|
||||
sample.gen_cpp_sample()
|
||||
|
||||
cmake_gen = cmake_creater.gen_build_sys(cutlass_deps_dir, output_dir)
|
||||
cmake_gen.gen_code()
|
||||
|
||||
verify = verify_creater.gen_verify(fuse_gemm_info, gen_name, for_fused_wrapper, output_dir)
|
||||
verify.gen_code()
|
||||
131
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_cmake.py
Normal file
131
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_cmake.py
Normal file
@@ -0,0 +1,131 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
class gen_build_sys:
|
||||
def __init__(self, cutlass_deps_dir, output_dir = "../"):
|
||||
self.output_dir = output_dir
|
||||
self.cutlass_deps_dir = cutlass_deps_dir
|
||||
|
||||
def gen_top(self):
|
||||
code = ""
|
||||
code += '''\
|
||||
# Auto Generated code - Do not edit.
|
||||
|
||||
cmake_minimum_required(VERSION 3.8)
|
||||
project(CUTLASS_MULTI_GEMMS LANGUAGES CXX CUDA)
|
||||
find_package(CUDAToolkit)
|
||||
set(CUDA_PATH ${{CUDA_TOOLKIT_ROOT_DIR}})
|
||||
set(CUTLASS_PATH \"{cutlass_deps_dir}/include\")
|
||||
set(CUTLASS_UTIL_PATH \"{cutlass_deps_dir}/tools/util/include\")
|
||||
list(APPEND CMAKE_MODULE_PATH ${{CUDAToolkit_LIBRARY_DIR}})
|
||||
'''.format(cutlass_deps_dir=self.cutlass_deps_dir)
|
||||
|
||||
code += '''\
|
||||
set(GPU_ARCHS \"\" CACHE STRING
|
||||
\"List of GPU architectures (semicolon-separated) to be compiled for.\")
|
||||
|
||||
if(\"${GPU_ARCHS}\" STREQUAL \"\")
|
||||
set(GPU_ARCHS \"70\")
|
||||
endif()
|
||||
|
||||
foreach(arch ${GPU_ARCHS})
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -gencode arch=compute_${arch},code=sm_${arch}\")
|
||||
if(SM STREQUAL 70 OR SM STREQUAL 75)
|
||||
set(CMAKE_C_FLAGS \"${CMAKE_C_FLAGS} -DWMMA\")
|
||||
set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS} -DWMMA\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -DWMMA\")
|
||||
endif()
|
||||
endforeach()
|
||||
|
||||
set(CMAKE_C_FLAGS \"${CMAKE_C_FLAGS}\")
|
||||
set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS}\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -Xcompiler -Wall\")
|
||||
|
||||
set(CMAKE_C_FLAGS_DEBUG \"${CMAKE_C_FLAGS_DEBUG} -Wall -O0\")
|
||||
set(CMAKE_CXX_FLAGS_DEBUG \"${CMAKE_CXX_FLAGS_DEBUG} -Wall -O0\")
|
||||
set(CMAKE_CUDA_FLAGS_DEBUG \"${CMAKE_CUDA_FLAGS_DEBUG} -O0 -G -Xcompiler -Wall\")
|
||||
|
||||
set(CMAKE_CXX_STANDARD 11)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
|
||||
if(CMAKE_CXX_STANDARD STREQUAL \"11\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} --expt-extended-lambda\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr\")
|
||||
endif()
|
||||
|
||||
set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS} -g -O3\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -Xcompiler -O3\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -Xcompiler=-fno-strict-aliasing\")
|
||||
|
||||
set(COMMON_HEADER_DIRS
|
||||
${PROJECT_SOURCE_DIR}
|
||||
${CUDAToolkit_INCLUDE_DIRS}
|
||||
)
|
||||
|
||||
set(COMMON_LIB_DIRS
|
||||
${CUDAToolkit_LIBRARY_DIR}
|
||||
)
|
||||
list(APPEND COMMON_HEADER_DIRS ${CUTLASS_PATH})
|
||||
list(APPEND COMMON_HEADER_DIRS ${CUTLASS_UTIL_PATH})
|
||||
'''
|
||||
code += '''\
|
||||
include_directories(
|
||||
${COMMON_HEADER_DIRS}
|
||||
)
|
||||
|
||||
link_directories(
|
||||
${COMMON_LIB_DIRS}
|
||||
)
|
||||
|
||||
add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0)
|
||||
add_definitions(-DGOOGLE_CUDA=1)
|
||||
|
||||
add_executable(sample
|
||||
sample/sample.cu
|
||||
one_api.cu
|
||||
)
|
||||
target_link_libraries(sample PRIVATE
|
||||
-lcudart
|
||||
-lnvToolsExt
|
||||
${CMAKE_THREAD_LIBS_INIT}
|
||||
)
|
||||
|
||||
if(NOT DEFINED LIB_INSTALL_PATH)
|
||||
set(LIB_INSTALL_PATH ${CMAKE_CURRENT_BINARY_DIR})
|
||||
endif()
|
||||
'''
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
top_code = self.gen_top()
|
||||
with open(self.output_dir + "CMakeLists.txt", "w") as f:
|
||||
f.write(top_code)
|
||||
@@ -0,0 +1,120 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import ast
|
||||
|
||||
fuse_gemm_info = [
|
||||
{
|
||||
'epilogue': {
|
||||
'tp': 'LeakyRelu', #'CustomizedLeaky_RELU'
|
||||
'bias': {'addbias': False, 'bias_tp': 'mat'},
|
||||
'args': [('float', 'leaky_alpha', 1.3), ],
|
||||
'func': '''
|
||||
y = max(leaky_alpha * x, x)
|
||||
y = y * x
|
||||
'''
|
||||
}
|
||||
},
|
||||
|
||||
]
|
||||
class AnalysisNodeVisitor(ast.NodeVisitor):
|
||||
def visit_Import(self,node):
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_ImportFrom(self,node):
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Assign(self,node):
|
||||
print('Node type: Assign and fields: ', node._fields)
|
||||
# print('Node type: Assign and targets value: ', node.targets, node.value)
|
||||
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_BinOp(self, node):
|
||||
print('Node type: BinOp and fields: ', node._fields)
|
||||
print('node op: ', type(node.op).__name__)
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Expr(self, node):
|
||||
print('Node type: Expr and fields: ', node._fields)
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Num(self,node):
|
||||
print('Node type: Num and fields: ', node._fields)
|
||||
print('Node type: Num: ', node.n)
|
||||
|
||||
def visit_Name(self,node):
|
||||
print('Node type: Name and fields: ', node._fields)
|
||||
print('Node type: Name and fields: ', type(node.ctx).__name__, node.id)
|
||||
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Str(self, node):
|
||||
print('Node type: Str and fields: ', node._fields)
|
||||
|
||||
class CodeVisitor(ast.NodeVisitor):
|
||||
def visit_BinOp(self, node):
|
||||
if isinstance(node.op, ast.Add):
|
||||
node.op = ast.Sub()
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Assign(self, node):
|
||||
print('Assign %s' % node.value)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Name(self, node):
|
||||
print("Name:", node.id)
|
||||
self.generic_visit(node)
|
||||
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
print('Function Name:%s'% node.name.op)
|
||||
self.generic_visit(node)
|
||||
func_log_stmt = ast.Print(
|
||||
dest = None,
|
||||
values = [ast.Str(s = 'calling func: %s' % node.name, lineno = 0, col_offset = 0)],
|
||||
nl = True,
|
||||
lineno = 0,
|
||||
col_offset = 0,
|
||||
)
|
||||
node.body.insert(0, func_log_stmt)
|
||||
|
||||
visitor = AnalysisNodeVisitor()
|
||||
|
||||
code = \
|
||||
'''
|
||||
|
||||
a=max(leaky_alpha * x, x +1)
|
||||
|
||||
'''
|
||||
|
||||
visitor.visit(ast.parse(code))
|
||||
477
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_device.py
Normal file
477
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_device.py
Normal file
@@ -0,0 +1,477 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
from typing import *
|
||||
|
||||
import helper
|
||||
import gen_ir
|
||||
|
||||
import gen_kernel as gen_ker
|
||||
|
||||
|
||||
class gen_device:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, cutlass_deps_root, project_root, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.raw_gemm_info = fuse_gemm_info
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
self.user_header_file = user_header_file
|
||||
self.args = {}
|
||||
# device arg struct memebr
|
||||
self.arg_member = []
|
||||
self.gen_class_name = gen_class_name
|
||||
self.gen_kernel_name = gen_class_name + "Kernel"
|
||||
self.tempalte_args = []
|
||||
self.__tempalate_arg_list = {'Stages': int, 'SplitKSerial': bool, 'IsBetaZero': bool, 'AlignmentA': int, 'AlignmentB': int}
|
||||
|
||||
self.file_name = output_dir + "/device/" +gen_class_name +".h"
|
||||
self.sample_dir = output_dir
|
||||
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
self.this_file_root = output_dir + "/device/"
|
||||
|
||||
self.first_use_1stage = False
|
||||
|
||||
## gen kernel
|
||||
self.gen_kernel = gen_ker.gen_kernel(self.tempalte_args, self.gen_class_name, self.b2b_num, output_dir, cutlass_deps_root, project_root)
|
||||
|
||||
|
||||
def __check_arg_type(self, temp_arg):
|
||||
if temp_arg in self.__tempalate_arg_list.keys():
|
||||
return self.__tempalate_arg_list[temp_arg]
|
||||
|
||||
find_sub = False
|
||||
for candidate_arg in self.__tempalate_arg_list.keys():
|
||||
if (temp_arg.find(candidate_arg) != -1):
|
||||
return self.__tempalate_arg_list[candidate_arg]
|
||||
|
||||
return 'typename'
|
||||
|
||||
# def gen_B2b2bGemm_class():
|
||||
def set_arch(self, sm_cap, mma_tp):
|
||||
if sm_cap == 75 or sm_cap == 80 or sm_cap == 86:
|
||||
self.arch = "cutlass::arch::Sm" + str(sm_cap)
|
||||
|
||||
if mma_tp is 'hmma1688':
|
||||
self.mma_shape = [16, 8, 8]
|
||||
self.mma_tp = 'hmma'
|
||||
elif mma_tp is 'imma8816':
|
||||
self.mma_tp = 'imma'
|
||||
self.mma_shape = [8, 8, 16]
|
||||
else:
|
||||
return 0
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''\
|
||||
/* Auto Generated code - Do not edit.*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include \"{cutlass_root}cutlass/cutlass.h\"
|
||||
#include \"{cutlass_root}cutlass/numeric_types.h\"
|
||||
#include \"{cutlass_root}cutlass/arch/arch.h\"
|
||||
#include \"{cutlass_root}cutlass/device_kernel.h\"
|
||||
|
||||
#include \"{cutlass_root}cutlass/gemm/threadblock/threadblock_swizzle.h\"
|
||||
|
||||
#include \"{cutlass_root}cutlass/gemm/device/default_gemm_configuration.h\"
|
||||
#include \"{cutlass_root}cutlass/epilogue/thread/linear_combination_relu.h\"
|
||||
#include \"{cutlass_root}cutlass/epilogue/thread/linear_combination.h\"
|
||||
|
||||
#include \"{project_root}../kernel/b2b_gemm.h\"
|
||||
#include \"{project_root}../kernel/default_b2b_gemm.h\"
|
||||
'''.format(cutlass_root=self.cutlass_deps_root, project_root=self.project_root, this_file_root=self.this_file_root)
|
||||
include_user_header = ""
|
||||
for header in self.user_header_file:
|
||||
include_user_header += "#include \"" + header + "\"\n"
|
||||
return code + include_user_header
|
||||
|
||||
def gen_code(self, sm_cap, mma_tp, ifprint = True):
|
||||
self.set_arch(sm_cap, mma_tp)
|
||||
|
||||
self.update_b2b_args()
|
||||
print(self.fuse_gemm_info)
|
||||
self.update_b2b_class_template_args()
|
||||
|
||||
func_code = self.gen_all_func()
|
||||
member_var_code = "private:\n typename B2bGemmKernel::Params params_;\n"
|
||||
|
||||
gen_code = gen_ir.gen_template_class(self.gen_class_name, self.tempalte_args, func_code + member_var_code)
|
||||
code = self.gen_include_header() + gen_ir.gen_namespace("cutlass", gen_ir.gen_namespace("gemm", gen_ir.gen_namespace("device", gen_code)))
|
||||
|
||||
if ifprint:
|
||||
print(code)
|
||||
|
||||
print("[INFO]: Gen device code output Dir: is ", self.file_name)
|
||||
with open(self.file_name, 'w+') as f:
|
||||
f.write(code)
|
||||
|
||||
|
||||
gen_kernel = self.gen_kernel.gen_code(self.first_use_1stage)
|
||||
print(gen_kernel)
|
||||
|
||||
def update_b2b_class_template_args(self):
|
||||
for arg in self.args.keys():
|
||||
self.tempalte_args.append([self.__check_arg_type(arg), arg, self.args[arg]])
|
||||
|
||||
def update_b2b_args(self):
|
||||
|
||||
self.args['ElementA'] = helper.type_2_cutlass_type(self.fuse_gemm_info[0]['A_tp'])
|
||||
self.args['LayoutA'] = helper.type_2_cutlass_type(self.fuse_gemm_info[0]['A_format'])
|
||||
|
||||
cnt = 0
|
||||
|
||||
warp_M_tile = 32
|
||||
|
||||
# Determine maxmimum N_tile
|
||||
Max_Ntile = 0
|
||||
for layer in self.fuse_gemm_info:
|
||||
n_tile = layer['mnk'][1]
|
||||
if n_tile > Max_Ntile:
|
||||
Max_Ntile = n_tile
|
||||
if Max_Ntile >= 256:
|
||||
warp_M_tile = 16
|
||||
|
||||
stages_temp = []
|
||||
|
||||
for layer in self.fuse_gemm_info:
|
||||
cnt_str = str(cnt)
|
||||
B_tp_str= 'ElementB' + cnt_str
|
||||
B_format_str = 'LayoutB' + cnt_str
|
||||
C_tp_str= 'ElementC' + cnt_str
|
||||
C_format_str = 'LayoutC' + cnt_str
|
||||
Acc_str = 'ElementAccumulator' + cnt_str
|
||||
|
||||
self.args[B_tp_str] = helper.type_2_cutlass_type(layer['B_tp'])
|
||||
self.args[B_format_str] = helper.type_2_cutlass_type(layer['B_format'])
|
||||
self.args[C_tp_str] = helper.type_2_cutlass_type(layer['C_tp'])
|
||||
self.args[C_format_str] = helper.type_2_cutlass_type(layer['C_format'])
|
||||
self.args[Acc_str] = helper.type_2_cutlass_type(layer['Acc_tp'])
|
||||
|
||||
|
||||
mnk = layer['mnk'][:]
|
||||
|
||||
tile_mnk = mnk[:]
|
||||
|
||||
tile_mnk[2] = 32 # force the ktile is 32
|
||||
|
||||
#N tile gen
|
||||
if mnk[1] > 1024:
|
||||
assert(0)
|
||||
elif mnk[1] > 512:
|
||||
tile_mnk[1] = 1024
|
||||
elif mnk[1] > 256:
|
||||
tile_mnk[1] = 512
|
||||
elif mnk[1] > 128:
|
||||
tile_mnk[1] = 256
|
||||
elif mnk[1] > 64:
|
||||
tile_mnk[1] = 128
|
||||
elif mnk[1] > 32:
|
||||
tile_mnk[1] = 64
|
||||
else :
|
||||
tile_mnk[1] = 32
|
||||
|
||||
if tile_mnk[1] == 512:
|
||||
stages_temp.append(1)
|
||||
else:
|
||||
stages_temp.append(2)
|
||||
|
||||
tile_mnk[0] = 4 * warp_M_tile
|
||||
|
||||
|
||||
|
||||
epilogue_setted_type = helper.get_epilogue_tp(layer)
|
||||
cutlass_epilogue_name = "LinearCombinationRelu"
|
||||
if epilogue_setted_type.lower() == 'leakyrelu':
|
||||
cutlass_epilogue_name = "LinearCombinationLeakyRelu"
|
||||
elif epilogue_setted_type.lower() == 'identity':
|
||||
cutlass_epilogue_name = "LinearCombination"
|
||||
|
||||
epilogue_str = 'EpilogueOutputOp' + cnt_str
|
||||
if cnt != len(self.fuse_gemm_info) - 1:
|
||||
n = layer['mnk'][1]
|
||||
Fragments = tile_mnk[1] // 8 * 2
|
||||
self.args[epilogue_str] = "cutlass::epilogue::thread::" + cutlass_epilogue_name + "<ElementC0_, " + str(Fragments) +", ElementAccumulator0_, ElementAccumulator0_>"
|
||||
else:
|
||||
n = layer['mnk'][1]
|
||||
n_mod_8 = n % 4
|
||||
N_align_elements = 1
|
||||
if n_mod_8 == 0:
|
||||
N_align_elements = 8
|
||||
elif n_mod_8 == 4:
|
||||
N_align_elements = 4
|
||||
elif n_mod_8 == 2 or n_mod_8 == 6:
|
||||
N_align_elements = 2
|
||||
|
||||
self.args[epilogue_str] = "cutlass::epilogue::thread::" + cutlass_epilogue_name+ "<ElementC0_, " + str(N_align_elements) + ", ElementAccumulator0_, ElementAccumulator0_>"
|
||||
|
||||
|
||||
|
||||
ThreadBlockShape_str = 'ThreadblockShape' + cnt_str
|
||||
|
||||
self.args[ThreadBlockShape_str] = helper.cvt_2_cutlass_shape(tile_mnk)
|
||||
|
||||
WarpShape_str = 'WarpShape' + cnt_str
|
||||
tile_mnk[0] = warp_M_tile
|
||||
self.args[WarpShape_str] = helper.cvt_2_cutlass_shape(tile_mnk)
|
||||
cnt += 1
|
||||
|
||||
|
||||
self.args['ElementD'] = helper.type_2_cutlass_type(self.fuse_gemm_info[self.b2b_num - 1]['C_tp'])
|
||||
self.args['LayoutD'] = helper.type_2_cutlass_type(self.fuse_gemm_info[self.b2b_num - 1]['C_format'])
|
||||
|
||||
self.args['InstructionShape'] = helper.cvt_2_cutlass_shape(self.mma_shape)
|
||||
self.args['OperatorClass'] = 'arch::OpClassTensorOp'
|
||||
self.args['ArchTag'] = self.arch
|
||||
self.args['ThreadblockSwizzle'] = 'threadblock::GemmBatchedIdentityThreadblockSwizzle'
|
||||
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
self.args[helper.var_idx('Stages', i)] = "2"
|
||||
|
||||
self.args['AlignmentA'] = str(8)
|
||||
self.args['AlignmentB'] = str(8)
|
||||
self.args['SplitKSerial'] = 'false'
|
||||
self.args['Operator'] = 'typename DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB0_, ElementC0_, ElementAccumulator0_>::Operator'
|
||||
self.args['IsBetaZero'] = 'false'
|
||||
|
||||
|
||||
def gen_using_kernel(self):
|
||||
code = "using B2bGemmKernel = typename kernel::DefaultB2bGemm<\n"
|
||||
code += " " + "ElementA,\n"
|
||||
code += " " + "LayoutA,\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("ElementB", i) + ",\n"
|
||||
code += " " + helper.var_idx("LayoutB", i) + ",\n"
|
||||
code += " " + helper.var_idx("ElementC", i) + ",\n"
|
||||
code += " " + helper.var_idx("LayoutC", i) + ",\n"
|
||||
code += " " + helper.var_idx("ElementAccumulator", i) + ",\n"
|
||||
code += " " + helper.var_idx("EpilogueOutputOp", i) + ",\n"
|
||||
code += " " + helper.var_idx("ThreadblockShape", i) + ",\n"
|
||||
code += " " + helper.var_idx("WarpShape", i) + ",\n"
|
||||
|
||||
code += " " + "ElementD,\n"
|
||||
code += " " + "LayoutD,\n"
|
||||
code += " " + "InstructionShape,\n"
|
||||
code += " " + "OperatorClass,\n"
|
||||
code += " " + "ArchTag,\n"
|
||||
code += " " + "ThreadblockSwizzle,\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("Stages", i) + ",\n"
|
||||
|
||||
|
||||
code += " " + "AlignmentA,\n"
|
||||
code += " " + "AlignmentB,\n"
|
||||
code += " " + "SplitKSerial,\n"
|
||||
code += " " + "Operator,\n"
|
||||
code += " " + "IsBetaZero_\n"
|
||||
|
||||
code += ">::B2bGemmKernel;\n\n"
|
||||
|
||||
return code
|
||||
|
||||
def gen_args(self):
|
||||
|
||||
def gen_arg_member(b2b_num):
|
||||
data_members = []
|
||||
|
||||
for i in range(b2b_num):
|
||||
member_type = "GemmCoord"
|
||||
member_name = "problem_size_" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
member_type = "TensorRef<ElementA const, LayoutA>"
|
||||
member_name = "ref_A0"
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
for i in range(b2b_num):
|
||||
member_type = "TensorRef<ElementB" + str(i) + " const, LayoutB" + str(i) +">"
|
||||
member_name = "ref_B" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
member_type = "TensorRef<ElementC" + str(i) + " const, LayoutC" + str(i) +">"
|
||||
member_name = "ref_C" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
member_type = "TensorRef<ElementD, LayoutD>"
|
||||
member_name = helper.var_idx("ref_D", b2b_num - 1)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
for i in range(b2b_num):
|
||||
member_type = "typename EpilogueOutputOp" + str(i) + "::Params"
|
||||
member_name = "epilogue" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
data_members.append(('int', 'batch_count'))
|
||||
|
||||
return data_members
|
||||
|
||||
def gen_arg_struct_default_ctor(struct_name, data_members, inital_param_num, inital_value):
|
||||
constructs_code = gen_ir.indentation + "CUTLASS_HOST_DEVICE\n" + \
|
||||
gen_ir.indentation + struct_name + " (): "
|
||||
for i in range(inital_param_num):
|
||||
final_param = ','
|
||||
if i == inital_param_num - 1:
|
||||
final_param = '{ }'
|
||||
constructs_code += data_members[i][1] + inital_value + final_param
|
||||
|
||||
constructs_code += "\n"
|
||||
return constructs_code
|
||||
|
||||
def gen_arg_struct_ctor(struct_name, data_members):
|
||||
constructs_code = gen_ir.indentation + "CUTLASS_HOST_DEVICE\n" + \
|
||||
gen_ir.indentation + struct_name + " (\n"
|
||||
cnt = 0
|
||||
param_num = len(data_members)
|
||||
for param in data_members:
|
||||
final = ',\n'
|
||||
if cnt == param_num - 1:
|
||||
final = '\n):\n'
|
||||
constructs_code += gen_ir.indentation + param[0] + " " + param[1] + "_" + final
|
||||
cnt += 1
|
||||
|
||||
cnt = 0
|
||||
for param in data_members:
|
||||
final = '),\n'
|
||||
if cnt == param_num - 1:
|
||||
final = ") { }\n"
|
||||
constructs_code += gen_ir.indentation + param[1] + "(" + param[1] + "_" + final
|
||||
cnt += 1
|
||||
|
||||
constructs_code += "\n"
|
||||
return constructs_code
|
||||
|
||||
# (variable type, variable name)
|
||||
struct_member = gen_arg_member(self.b2b_num)
|
||||
self.arg_member = struct_member
|
||||
|
||||
codeBody = ""
|
||||
for each_member in struct_member:
|
||||
codeBody += gen_ir.indentation + each_member[0] + " " + each_member[1] + ";\n"
|
||||
|
||||
codeBody += gen_arg_struct_default_ctor("Arguments", struct_member, self.b2b_num, "(0,0,0)") + "\n"
|
||||
codeBody += gen_arg_struct_ctor("Arguments", struct_member) + "\n"
|
||||
struct_code = gen_ir.gen_struct("Arguments", codeBody)
|
||||
return struct_code
|
||||
|
||||
def gen_func_constructs(self):
|
||||
code = self.gen_class_name +"() {}"
|
||||
return code
|
||||
|
||||
def gen_func_initialize(self):
|
||||
code = "Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {\n" + \
|
||||
"// Determine grid shape\n" + \
|
||||
"ThreadblockSwizzle threadblock_swizzle;\n" + \
|
||||
"cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(\n" + \
|
||||
" args.problem_size_0, \n" + \
|
||||
" { ThreadblockShape0::kM, ThreadblockShape0::kN, ThreadblockShape0::kK },\n" + \
|
||||
" args.batch_count);\n" + \
|
||||
"// Initialize the Params structure\n" + \
|
||||
"params_ = typename B2bGemmKernel::Params{\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx(" args.problem_size_", i) + ",\n"
|
||||
code += " grid_shape,\n" + \
|
||||
" args.ref_A0.non_const_ref(),\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx(" args.ref_B", i) + ".non_const_ref(),\n"
|
||||
code += helper.var_idx(" args.ref_C", i) + ".non_const_ref(),\n"
|
||||
|
||||
code += helper.var_idx(" args.ref_D", self.b2b_num - 1) + ",\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx(" args.epilogue", i) + ",\n"
|
||||
|
||||
code += " args.batch_count\n"
|
||||
code += "};\n" + \
|
||||
"return Status::kSuccess;\n" + \
|
||||
"}\n"
|
||||
return code
|
||||
|
||||
def gen_func_run(self):
|
||||
code = "Status run(cudaStream_t stream = nullptr) {\n" + \
|
||||
"\n" + \
|
||||
" ThreadblockSwizzle threadblock_swizzle;\n" + \
|
||||
"\n" + \
|
||||
" dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);\n" + \
|
||||
" dim3 block(B2bGemmKernel::kThreadCount, 1, 1);\n" + \
|
||||
"\n" + \
|
||||
" cudaError_t result;\n" + \
|
||||
"\n" + \
|
||||
" int smem_size = int(sizeof(typename B2bGemmKernel::SharedStorage));\n" + \
|
||||
" if (smem_size >= (48 << 10)) {\n" + \
|
||||
" result = cudaFuncSetAttribute(Kernel<B2bGemmKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);\n" + \
|
||||
"\n" + \
|
||||
" if (result != cudaSuccess) {\n" + \
|
||||
" return Status::kErrorInternal;\n" + \
|
||||
" }\n" + \
|
||||
"\n" + \
|
||||
" result = cudaFuncSetAttribute(\n" + \
|
||||
" Kernel<B2bGemmKernel>,\n" + \
|
||||
" cudaFuncAttributePreferredSharedMemoryCarveout, 100);\n" + \
|
||||
"\n" + \
|
||||
" if (result != cudaSuccess) {\n" + \
|
||||
" return Status::kErrorInternal;\n" + \
|
||||
" }\n" + \
|
||||
" }\n" + \
|
||||
" cutlass::Kernel<B2bGemmKernel><<<grid, block, smem_size, stream>>>(params_);\n" + \
|
||||
" result = cudaGetLastError();\n" + \
|
||||
" return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;\n" + \
|
||||
" }\n"
|
||||
|
||||
return code
|
||||
def gen_func_operator(self):
|
||||
opeartor_with_arg_code = "Status operator()(\n" + \
|
||||
" Arguments const &args,\n" + \
|
||||
" void *workspace = nullptr,\n" + \
|
||||
" cudaStream_t stream = nullptr) {\n" + \
|
||||
" Status status = initialize(args, workspace);\n" + \
|
||||
" \n" + \
|
||||
" if (status == Status::kSuccess) {\n" + \
|
||||
" status = run(stream);\n" + \
|
||||
" }\n" + \
|
||||
" return status;\n" + \
|
||||
"}\n"
|
||||
operator_code = "Status operator()(\n" + \
|
||||
" cudaStream_t stream = nullptr) {\n" + \
|
||||
" Status status = run(stream);\n" + \
|
||||
" return status;\n" + \
|
||||
"}\n"
|
||||
return opeartor_with_arg_code + "\n" + operator_code
|
||||
|
||||
def gen_all_func(self):
|
||||
return self.gen_using_kernel() + "\n" + \
|
||||
self.gen_args() + "\n" + \
|
||||
self.gen_func_constructs() + "\n" + \
|
||||
self.gen_func_initialize() + "\n" + \
|
||||
self.gen_func_run() + "\n" + \
|
||||
self.gen_func_operator()
|
||||
249
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_ir.py
Normal file
249
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_ir.py
Normal file
@@ -0,0 +1,249 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
|
||||
|
||||
indentation = " "
|
||||
|
||||
|
||||
def append_word(word):
|
||||
code = ""
|
||||
code += word
|
||||
code += " "
|
||||
return code
|
||||
|
||||
|
||||
def gen_namespace(namespace, codeBody):
|
||||
code_gen = "namespace " + namespace + " {\n"
|
||||
code_gen += codeBody
|
||||
code_gen += "} // namespace " + namespace + "\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_expression(type, lval, rval = None):
|
||||
code_gen = ""
|
||||
code_gen += append_word(type)
|
||||
code_gen += append_word(lval)
|
||||
if rval is not None:
|
||||
code_gen += append_word("=")
|
||||
code_gen += append_word(rval)
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_class(name, codeBody, inheritance_code = None):
|
||||
code_gen = ""
|
||||
if inheritance_code is None:
|
||||
code_gen = "class " + name + "{\n"
|
||||
else:
|
||||
code_gen = "class " + name + " : "+ inheritance_code + "{\n"
|
||||
code_gen += codeBody
|
||||
code_gen += "}; // class " + name + "\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_struct(name, codeBody, specialized = None):
|
||||
specialized_code = ""
|
||||
if specialized is not None:
|
||||
specialized_code = "<" + specialized + ">"
|
||||
code_gen = "struct " + name + specialized_code + "{\n"
|
||||
code_gen += codeBody
|
||||
code_gen += "}; // struct " + name + "\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_arg(arg_type, arg_name, default_val = None):
|
||||
rval = None
|
||||
if default_val is not None:
|
||||
rval = str(default_val)
|
||||
|
||||
arg_typename = ""
|
||||
if arg_type is int:
|
||||
arg_typename = "int"
|
||||
elif arg_type is bool:
|
||||
arg_typename = "bool"
|
||||
else:
|
||||
arg_typename = "typename"
|
||||
|
||||
internal_arg_name = arg_name + "_"
|
||||
|
||||
code_gen = indentation
|
||||
code_gen += gen_expression(arg_typename, internal_arg_name, rval)
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_args(args, set_default = True):
|
||||
arg_len = len(args)
|
||||
cnt = 1
|
||||
code_gen = ""
|
||||
for arg_tuple in args:
|
||||
arg_type = arg_tuple[0]
|
||||
arg_name = arg_tuple[1]
|
||||
arg_default_val = None
|
||||
if len(arg_tuple) == 3 and set_default:
|
||||
arg_default_val = arg_tuple[2]
|
||||
|
||||
code_gen += gen_template_arg(arg_type, arg_name, arg_default_val)
|
||||
if cnt != arg_len:
|
||||
code_gen += ",\n"
|
||||
cnt += 1
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_head(args, set_default = True):
|
||||
code_gen = "template <\n"
|
||||
code_gen += gen_template_args(args, set_default)
|
||||
code_gen += ">\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def export_template_args(args):
|
||||
code_gen = "public:\n"
|
||||
for arg_tuple in args:
|
||||
code_gen += indentation
|
||||
arg_type = arg_tuple[0]
|
||||
arg_name = arg_tuple[1]
|
||||
internal_arg_name = arg_name + "_"
|
||||
|
||||
typename = ""
|
||||
if arg_type is int:
|
||||
typename = "static int const"
|
||||
elif arg_type is bool:
|
||||
typename = "static bool const"
|
||||
else:
|
||||
typename = "using"
|
||||
|
||||
code_gen += gen_expression(typename, arg_name, internal_arg_name)
|
||||
code_gen += ";\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_class(class_name, args, codeBody, set_default = True, inheritance_code = None):
|
||||
code_gen = ""
|
||||
|
||||
code_gen += gen_template_head(args, set_default)
|
||||
code_gen += gen_class(class_name, export_template_args(args) + codeBody, inheritance_code)
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_struct(struct_name, args, codeBody, speicalized = None, set_default = True, export_args = True):
|
||||
code_gen = ""
|
||||
code_gen += gen_template_head(args, set_default)
|
||||
code = export_template_args(args) + codeBody
|
||||
if export_args is False:
|
||||
code = codeBody
|
||||
code_gen += gen_struct(struct_name, code , speicalized)
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_declare_template_struct(name, *params):
|
||||
code = name + "<"
|
||||
cnt = 0
|
||||
param_num = len(params)
|
||||
for param in params:
|
||||
final = ", "
|
||||
if cnt == param_num - 1:
|
||||
final = ""
|
||||
code += param + final
|
||||
cnt += 1
|
||||
code += ">;\n"
|
||||
return code
|
||||
|
||||
|
||||
def filtered_param(params, name_and_value_pair, keep_ = False):
|
||||
rtn_template_args = []
|
||||
speicalized_template_args = []
|
||||
|
||||
for param in params:
|
||||
param_name = ""
|
||||
if len(param) >= 1:
|
||||
param_name = param[1]
|
||||
else:
|
||||
param_name = param[0]
|
||||
|
||||
hit_flag = False
|
||||
set_value = ""
|
||||
for n_v_pair in name_and_value_pair:
|
||||
|
||||
filter_name = n_v_pair[0]
|
||||
set_value = n_v_pair[1]
|
||||
|
||||
if param_name == (filter_name + "_") or param_name == filter_name :
|
||||
hit_flag = True
|
||||
break
|
||||
|
||||
|
||||
if hit_flag is False:
|
||||
rtn_template_args.append(param)
|
||||
|
||||
if hit_flag is True:
|
||||
speicalized_template_args.append(set_value)
|
||||
else:
|
||||
if keep_ is True:
|
||||
speicalized_template_args.append(param_name + "_")
|
||||
else:
|
||||
speicalized_template_args.append(param_name)
|
||||
|
||||
|
||||
specialized_template_arg_str = helper.list_2_string(speicalized_template_args)
|
||||
|
||||
return rtn_template_args, specialized_template_arg_str
|
||||
|
||||
|
||||
def gen_func(func_name, arg_lists, code_body, only_declare = False, with_cudaStream = True):
|
||||
code = "void " + func_name + "(\n"
|
||||
for arg in arg_lists:
|
||||
arg_tp = arg[0]
|
||||
arg_nm = arg[1]
|
||||
code += " " + arg_tp + " " + arg_nm + ",\n"
|
||||
code += "cudaStream_t stream)"
|
||||
if only_declare :
|
||||
return code
|
||||
code += "{\n"
|
||||
|
||||
code += code_body + "\n"
|
||||
code += "}\n"
|
||||
return code
|
||||
|
||||
|
||||
def indent_level(code, level = 0):
|
||||
rtn_code = ""
|
||||
for i in range(level):
|
||||
rtn_code += " "
|
||||
|
||||
rtn_code += code
|
||||
|
||||
return rtn_code
|
||||
476
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_kernel.py
Normal file
476
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_kernel.py
Normal file
@@ -0,0 +1,476 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import gen_ir
|
||||
import helper
|
||||
import gen_threadblock as gen_tb
|
||||
|
||||
|
||||
class gen_default_Gemm:
|
||||
def __init__(self, template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root):
|
||||
self.gen_class_name = "B2bGemm"
|
||||
self.template_param = template_param
|
||||
self.b2b_num = b2b_num
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
|
||||
def gen_B2bMma(self, specialized_template_args):
|
||||
code = "using B2bMma = typename cutlass::gemm::threadblock::DefaultB2bMma<\n"
|
||||
code += specialized_template_args
|
||||
code += ">::ThreadblockB2bMma;\n"
|
||||
|
||||
# print(code)
|
||||
return code
|
||||
|
||||
def gen_epilogue(self):
|
||||
epilogue_code = ""
|
||||
epilogue_code += helper.var_idx("static const int kPartitionsK", self.b2b_num - 1) + helper.var_idx(" = ThreadblockShape", self.b2b_num - 1) + helper.var_idx("::kK / WarpShape", self.b2b_num - 1) + "::kK;\n"
|
||||
|
||||
epilogue_code += "using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<\n"
|
||||
epilogue_code += " " + helper.var_idx("ThreadblockShape", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("typename B2bMma::Operator", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("kPartitionsK", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("EpilogueOutputOp", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("EpilogueOutputOp", self.b2b_num - 1) + "::kCount\n"
|
||||
epilogue_code += ">::Epilogue;\n"
|
||||
|
||||
epilogue_code += "using B2bGemmKernel = kernel::B2bGemm<B2bMma, Epilogue, ThreadblockSwizzle, SplitKSerial>;\n\n"
|
||||
|
||||
return epilogue_code
|
||||
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''
|
||||
/* Auto Generated code - Do not edit.*/
|
||||
|
||||
#pragma once
|
||||
#include \"{cutlass_dir}cutlass/cutlass.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/layout/matrix.h\"
|
||||
#include \"{cutlass_dir}cutlass/numeric_types.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/epilogue.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/thread/linear_combination.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/gemm/gemm.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/kernel/gemm_pipelined.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_sm75.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_sm70.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_sm80.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_simt.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/threadblock_swizzle.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/default_epilogue_tensor_op.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/default_epilogue_simt.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/transform/threadblock/predicated_tile_iterator.h\"
|
||||
|
||||
#include \"../kernel/b2b_gemm.h\"
|
||||
#include \"../threadblock/default_b2b_mma.h\"
|
||||
'''.format(cutlass_dir=self.cutlass_deps_root)
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
gen_using = ''
|
||||
# Generate default template struct
|
||||
gen_code = gen_ir.gen_template_struct("Default" + self.gen_class_name, self.template_param,"", speicalized = None, set_default=False)
|
||||
|
||||
|
||||
filter_list = []
|
||||
filter_list.append(('Stages', 2))
|
||||
filter_list.append(("OperatorClass", "arch::OpClassTensorOp"))
|
||||
filter_list.append(("ArchTag", "arch::Sm75"))
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
filter_list.append((helper.var_idx("LayoutC", i), "layout::RowMajor"))
|
||||
|
||||
|
||||
rtn_template_args, speicalized_template_args = gen_ir.filtered_param(self.template_param, filter_list, keep_= True)
|
||||
|
||||
|
||||
B2bMma_code = self.gen_B2bMma(speicalized_template_args)
|
||||
epilogue_and_rest_code = self.gen_epilogue()
|
||||
|
||||
gen_special_code = gen_ir.gen_template_struct("Default" + self.gen_class_name, rtn_template_args, B2bMma_code + epilogue_and_rest_code, speicalized = speicalized_template_args, set_default=False)
|
||||
|
||||
code = gen_ir.gen_namespace("cutlass", gen_ir.gen_namespace("gemm", gen_ir.gen_namespace("kernel", gen_code + gen_special_code)))
|
||||
|
||||
return self.gen_include_header() + code
|
||||
|
||||
|
||||
class gen_Kernel:
|
||||
def __init__(self, template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root):
|
||||
self.gen_class_name = "B2bGemm"
|
||||
self.template_param = template_param
|
||||
self.b2bnum = b2b_num
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''
|
||||
#pragma once
|
||||
|
||||
#include \"{cutlass_dir}cutlass/cutlass.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/gemm.h\"
|
||||
#include \"{cutlass_dir}cutlass/matrix_coord.h\"\n'''.format(cutlass_dir=self.cutlass_deps_root)
|
||||
return code
|
||||
|
||||
def gen_Params(self):
|
||||
gen_param = ""
|
||||
for i in range(self.b2bnum):
|
||||
gen_param += " " + helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + ";\n"
|
||||
gen_param += " " + "cutlass::gemm::GemmCoord grid_tiled_shape;\n"
|
||||
gen_param += " " + "typename B2bMma::IteratorA0::Params params_A0;\n"
|
||||
gen_param += " " + "typename B2bMma::IteratorA0::TensorRef ref_A0;\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
gen_param += " " + helper.var_idx("typename B2bMma::IteratorB", i) + helper.var_idx("::Params params_B", i) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename B2bMma::IteratorB", i) + helper.var_idx("::TensorRef ref_B", i) + ";\n"
|
||||
if i == self.b2bnum - 1:
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::Params params_C", i) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_C", i) + ";\n"
|
||||
|
||||
else:
|
||||
gen_param += " " + helper.var_idx("typename FusedAddBiasEpilogue", i) + helper.var_idx("::OutputTileIterator::Params params_C", i) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename FusedAddBiasEpilogue", i) + helper.var_idx("::OutputTileIterator::TensorRef ref_C", i) + ";\n"
|
||||
|
||||
|
||||
|
||||
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::Params params_D", self.b2bnum - 1) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_D", self.b2bnum - 1) + ";\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
gen_param += " " + helper.var_idx("typename OutputOp", i) + helper.var_idx("::Params output_op_", i) + ";\n"
|
||||
|
||||
gen_param += " " + 'int batch_count' + ";\n"
|
||||
gen_param += " " + 'int gemm_k_iterations_0' + ";\n"
|
||||
|
||||
|
||||
return gen_param
|
||||
|
||||
def gen_Memberfunc(self):
|
||||
code_default = "\nCUTLASS_HOST_DEVICE\n"
|
||||
code_default += "Params()"
|
||||
|
||||
code_default += " { } \n\n"
|
||||
|
||||
code_construct = "\nCUTLASS_HOST_DEVICE\n"
|
||||
code_construct += "Params(\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("cutlass::gemm::GemmCoord const & problem_size_", i) + ",\n"
|
||||
|
||||
code_construct += " " + "cutlass::gemm::GemmCoord const & grid_tiled_shape,\n"
|
||||
|
||||
code_construct += " " + "typename B2bMma::IteratorA0::TensorRef ref_A0,\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("typename B2bMma::IteratorB", i) + helper.var_idx("::TensorRef ref_B", i) + ",\n"
|
||||
if i == self.b2bnum - 1:
|
||||
code_construct += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_C", i) + ",\n"
|
||||
else:
|
||||
code_construct += " " + helper.var_idx("typename FusedAddBiasEpilogue", i) + helper.var_idx("::OutputTileIterator::TensorRef ref_C", i) + ",\n"
|
||||
|
||||
code_construct += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_D", self.b2bnum - 1) + ",\n"
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("typename OutputOp", i) + helper.var_idx("::Params output_op_", i) + helper.var_idx(" = typename OutputOp", i) + "::Params(),\n"
|
||||
|
||||
code_construct += " " + "int batch_count = 1\n"
|
||||
|
||||
code_construct += "):\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("problem_size_", i) + helper.var_idx("(problem_size_", i) + "),\n"
|
||||
|
||||
code_construct += " " + "grid_tiled_shape(grid_tiled_shape),\n"
|
||||
code_construct += " " + "params_A0(ref_A0.layout()),\n"
|
||||
code_construct += " " + "ref_A0(ref_A0),\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("params_B", i) + helper.var_idx("(ref_B", i) + ".layout()),\n"
|
||||
code_construct += " " + helper.var_idx("ref_B", i) + helper.var_idx("(ref_B", i) + "),\n"
|
||||
code_construct += " " + helper.var_idx("params_C", i) + helper.var_idx("(ref_C", i) + ".layout()),\n"
|
||||
code_construct += " " + helper.var_idx("ref_C", i) + helper.var_idx("(ref_C", i) + "),\n"
|
||||
|
||||
code_construct += " " + helper.var_idx("params_D", self.b2bnum - 1) + helper.var_idx("(ref_D", self.b2bnum - 1) + ".layout()),\n"
|
||||
code_construct += " " + helper.var_idx("ref_D", self.b2bnum - 1) + helper.var_idx("(ref_D", self.b2bnum - 1) + "),\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("output_op_", i) + helper.var_idx("(output_op_", i) + "), \n"
|
||||
|
||||
code_construct += " " + "batch_count(batch_count) {\n"
|
||||
code_construct += " " + helper.var_idx("gemm_k_iterations_", 0) + helper.var_idx(" = (problem_size_", 0) + helper.var_idx(".k() + B2bMma::Shape", 0) + helper.var_idx("::kK - 1) / B2bMma::Shape", 0) + "::kK;\n"
|
||||
|
||||
code_construct += "}\n"
|
||||
|
||||
return code_default + code_construct
|
||||
|
||||
def gen_using(self):
|
||||
code_using = ""
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
code_using += " " + helper.var_idx("using OutputOp", i) + helper.var_idx(" = typename B2bMma::OutputOp", i) + ";\n"
|
||||
|
||||
code_using += " " + helper.var_idx("using OutputOp", self.b2bnum - 1) + " = typename Epilogue::OutputOp;\n"
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
code_using += " " + helper.var_idx("using FusedAddBiasEpilogue", i) + helper.var_idx(" = typename B2bMma::FusedAddBiasEpilogue", i) +";\n"
|
||||
|
||||
|
||||
code_using += " " + "using WarpCount0 = typename B2bMma::WarpCount0;\n"
|
||||
code_using += " " + "static int const kThreadCount = 32 * WarpCount0::kCount;\n"
|
||||
|
||||
code_using += gen_ir.gen_struct("Params", self.gen_Params() + self.gen_Memberfunc())
|
||||
|
||||
code_using += "union SharedStorage {\n"
|
||||
code_using += " " + "typename B2bMma::B2bMmaSharedStorage main_loop;\n"
|
||||
code_using += " " + "typename Epilogue::SharedStorage epilogue;\n"
|
||||
code_using += "};\n"
|
||||
|
||||
return code_using
|
||||
|
||||
def gen_can_implement(self):
|
||||
gen_code = ""
|
||||
return gen_code
|
||||
|
||||
def gen_operator_and_constr(self):
|
||||
ctr_code = "CUTLASS_HOST_DEVICE\n"
|
||||
ctr_code += self.gen_class_name + "() { } \n\n"
|
||||
operator_code = "CUTLASS_DEVICE\n"
|
||||
operator_code += "void operator()(Params const ¶ms, SharedStorage &shared_storage) {\n"
|
||||
operator_code += " " + "ThreadblockSwizzle threadblock_swizzle;\n"
|
||||
operator_code += " " + "cutlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.grid_tiled_shape);\n"
|
||||
operator_code += " " + "int batch_idx = threadblock_tile_offset.k();\n"
|
||||
operator_code += " " + "if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||\n"
|
||||
operator_code += " " + "params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {\n"
|
||||
operator_code += " " + " " + "return;\n"
|
||||
operator_code += " " + "}\n"
|
||||
|
||||
operator_code += " " + "cutlass::MatrixCoord tb_offset_A0{\n"
|
||||
operator_code += " " + " " + "threadblock_tile_offset.m() * B2bMma::Shape0::kM,\n"
|
||||
operator_code += " " + " " + "0\n"
|
||||
operator_code += " " + "};\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
operator_code += " " + helper.var_idx("cutlass::MatrixCoord tb_offset_B", i) + "{\n"
|
||||
operator_code += " " + " " + "0,\n"
|
||||
operator_code += " " + " " + helper.var_idx("threadblock_tile_offset.n() * B2bMma::Shape", i) + "::kN\n"
|
||||
operator_code += " " + "};\n"
|
||||
|
||||
operator_code += " " + "int thread_idx = threadIdx.x;\n\n"
|
||||
|
||||
operator_code += " " + "MatrixCoord threadblock_offset(\n"
|
||||
operator_code += " " + " " + helper.var_idx("threadblock_tile_offset.m() * B2bMma::Shape", self.b2bnum - 1) + "::kM,\n"
|
||||
operator_code += " " + " " + helper.var_idx("threadblock_tile_offset.n() * B2bMma::Shape", self.b2bnum - 1) + "::kN\n"
|
||||
operator_code += " " + ");\n"
|
||||
|
||||
operator_code += " " + "typename B2bMma::IteratorA0 iterator_A0(\n"
|
||||
operator_code += " " + " " + "params.params_A0,\n"
|
||||
operator_code += " " + " " + "params.ref_A0.data(),\n"
|
||||
operator_code += " " + " " + "params.problem_size_0.mk(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "tb_offset_A0);\n"
|
||||
|
||||
operator_code += " " + "iterator_A0.add_pointer_offset(batch_idx * params.problem_size_0.m() * params.problem_size_0.k());\n\n"
|
||||
|
||||
|
||||
for i in range (self.b2bnum):
|
||||
operator_code += " " + helper.var_idx("typename B2bMma::IteratorB", i ) + helper.var_idx(" iterator_B", i) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_B", i) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_B", i) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_", i) + ".kn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + helper.var_idx("tb_offset_B", i) + ");\n"
|
||||
operator_code += " " + helper.var_idx("iterator_B", i) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", i) + helper.var_idx(".n() * params.problem_size_", i) + ".k());\n\n"
|
||||
|
||||
|
||||
for i in range (self.b2bnum - 1):
|
||||
operator_code += " " + helper.var_idx("typename FusedAddBiasEpilogue", i ) + helper.var_idx("::OutputTileIterator iterator_C", i) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_C", i) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_C", i) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_" , i) + ".mn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "threadblock_offset" + ");\n"
|
||||
operator_code += " " + helper.var_idx("int ref_C", i) + helper.var_idx("_stride = params.ref_C", i) + ".stride()[0];\n"
|
||||
operator_code += " " + helper.var_idx("iterator_C", i) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", i) + helper.var_idx(".n() * (ref_C", i) + helper.var_idx("_stride == 0 ? 1 : params.problem_size_", i) + ".m()));\n\n"
|
||||
|
||||
|
||||
for i in range (self.b2bnum - 1):
|
||||
operator_code += " " + helper.var_idx("FusedAddBiasEpilogue", i ) + helper.var_idx(" epilogue_", i ) + ";\n"
|
||||
|
||||
|
||||
operator_code += " " + "int warp_idx = __shfl_sync(0x1f, threadIdx.x / 32, 0);\n"
|
||||
operator_code += " " + "int lane_idx = threadIdx.x % 32;\n"
|
||||
|
||||
for i in range (self.b2bnum - 1):
|
||||
operator_code += " " + helper.var_idx("OutputOp", i) + helper.var_idx(" output_op_", i) + helper.var_idx("(params.output_op_", i) + ");\n"
|
||||
|
||||
operator_code += " " + "B2bMma b2bMma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);\n"
|
||||
|
||||
operator_code += " " + "typename B2bMma::FragmentC0 src_accum;\n"
|
||||
operator_code += " " + helper.var_idx("typename B2bMma::FragmentC", self.b2bnum - 1)+ " accumulators;\n"
|
||||
|
||||
operator_code += " " + "src_accum.clear();\n"
|
||||
operator_code += " " + "accumulators.clear();\n"
|
||||
operator_code += " " + "b2bMma(params.gemm_k_iterations_0, accumulators, iterator_A0, "
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
operator_code += helper.var_idx("iterator_B", i) + ", "
|
||||
|
||||
operator_code += "src_accum"
|
||||
if self.b2bnum != 1:
|
||||
operator_code += ", "
|
||||
for i in range(self.b2bnum - 1):
|
||||
operator_code += helper.var_idx("output_op_", i) + ", "
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
operator_code += helper.var_idx("epilogue_", i) + ", "
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
final = ", "
|
||||
if i == self.b2bnum - 2:
|
||||
final =""
|
||||
operator_code += helper.var_idx("iterator_C", i) + final
|
||||
operator_code += ");\n"
|
||||
|
||||
operator_code += " " + helper.var_idx("OutputOp", self.b2bnum - 1) + helper.var_idx(" output_op_", self.b2bnum - 1) + helper.var_idx("(params.output_op_", self.b2bnum - 1) + ");\n"
|
||||
operator_code += " " + "threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.grid_tiled_shape);\n"
|
||||
|
||||
|
||||
|
||||
operator_code += " " + helper.var_idx("typename Epilogue::OutputTileIterator iterator_C", self.b2bnum - 1) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_C", self.b2bnum - 1) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_C", self.b2bnum - 1) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_", self.b2bnum - 1) + ".mn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "threadblock_offset\n"
|
||||
operator_code += " " + ");\n"
|
||||
operator_code += " " + helper.var_idx("int ref_C", self.b2bnum - 1) + helper.var_idx("_stride = params.ref_C", self.b2bnum - 1) + ".stride()[0];\n"
|
||||
|
||||
operator_code += " " + helper.var_idx("iterator_C", self.b2bnum - 1) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", self.b2bnum - 1) + helper.var_idx(".n() * (ref_C", self.b2bnum - 1) + helper.var_idx("_stride == 0 ? 1 : params.problem_size_", self.b2bnum - 1) + ".m()));\n\n"
|
||||
|
||||
operator_code += " " + helper.var_idx("typename Epilogue::OutputTileIterator iterator_D", self.b2bnum - 1) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_D", self.b2bnum - 1) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_D", self.b2bnum - 1) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_", self.b2bnum - 1) + ".mn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "threadblock_offset\n"
|
||||
operator_code += " " + ");\n"
|
||||
operator_code += " " + helper.var_idx("iterator_D", self.b2bnum - 1) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", self.b2bnum - 1) + helper.var_idx(".n() * params.problem_size_", self.b2bnum - 1) + ".m());\n\n"
|
||||
|
||||
|
||||
operator_code += " " + "Epilogue epilogue(\n"
|
||||
operator_code += " " + " " + "shared_storage.epilogue,\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "warp_idx,\n"
|
||||
operator_code += " " + " " + "lane_idx\n"
|
||||
operator_code += " " + ");\n"
|
||||
|
||||
operator_code += " " + "epilogue("
|
||||
operator_code += helper.var_idx("output_op_", self.b2bnum - 1) + ", "
|
||||
operator_code += helper.var_idx("iterator_D", self.b2bnum - 1) + ", "
|
||||
operator_code += "accumulators, "
|
||||
operator_code += helper.var_idx("iterator_C", self.b2bnum - 1) + ");\n"
|
||||
operator_code += "}\n"
|
||||
|
||||
return ctr_code + operator_code
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''
|
||||
#pragma once
|
||||
|
||||
#include \"{cutlass_dir}cutlass/cutlass.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/gemm/gemm.h\"
|
||||
#include \"{cutlass_dir}cutlass/matrix_coord.h\"
|
||||
#include \"{cutlass_dir}cutlass/semaphore.h\"
|
||||
'''.format(cutlass_dir=self.cutlass_deps_root)
|
||||
return code
|
||||
def gen_code(self):
|
||||
|
||||
template_param = []
|
||||
template_param.append(("typename", "B2bMma"))
|
||||
template_param.append(("typename", "Epilogue"))
|
||||
template_param.append(("typename", "ThreadblockSwizzle"))
|
||||
template_param.append((bool, "SplitKSerial"))
|
||||
|
||||
code_body = ""
|
||||
code_body += self.gen_using()
|
||||
code_body += self.gen_operator_and_constr()
|
||||
|
||||
struct_code = gen_ir.gen_template_struct(self.gen_class_name, template_param, code_body)
|
||||
code = self.gen_include_header()
|
||||
code += gen_ir.gen_namespace("cutlass", gen_ir.gen_namespace("gemm", gen_ir.gen_namespace("kernel", struct_code)))
|
||||
|
||||
return self.gen_include_header() + code
|
||||
|
||||
|
||||
|
||||
class gen_kernel:
|
||||
def __init__(self, template_param, gen_class_name, b2b_num, output_dir, cutlass_deps_root, project_root):
|
||||
self.template_param = template_param
|
||||
|
||||
self.gen_class_name = "B2bGemm"
|
||||
self.gen_kernel_name = gen_class_name + "Kernel"
|
||||
self.tempalte_args = []
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
|
||||
self.gen_default_b2b_gemm = gen_default_Gemm(template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root)
|
||||
self.gen_Kerenl = gen_Kernel(template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root)
|
||||
|
||||
# Include gen_threadBlock
|
||||
self.gen_threadBlock = gen_tb.gen_threadblock(template_param, gen_class_name, b2b_num, output_dir, cutlass_deps_root, project_root)
|
||||
|
||||
self.file_dir = output_dir + "/kernel/"
|
||||
|
||||
def gen_code(self, first_use_1stage):
|
||||
|
||||
default_b2b_gemm = self.gen_default_b2b_gemm.gen_code()
|
||||
|
||||
print("[INFO]: Gen kernel code [default_b2b_gemm.h]output Dir: is ", self.file_dir)
|
||||
|
||||
with open(self.file_dir + "default_b2b_gemm.h", "w+") as f:
|
||||
f.write(default_b2b_gemm)
|
||||
|
||||
kernel = self.gen_Kerenl.gen_code()
|
||||
print("[INFO]: Gen kernel code [b2b_gemm.h]output Dir: is ", self.file_dir)
|
||||
|
||||
with open(self.file_dir + "b2b_gemm.h", "w+") as f:
|
||||
f.write(kernel)
|
||||
|
||||
# Call code to gen threadblock
|
||||
self.gen_threadBlock.gen_code(first_use_1stage)
|
||||
232
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_sample.py
Normal file
232
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_sample.py
Normal file
@@ -0,0 +1,232 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
import gen_ir as ir
|
||||
|
||||
class gen_test:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.gen_class_name = gen_class_name
|
||||
self.user_header_file = user_header_file
|
||||
self.sample_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
def gen_cpp_sample(self):
|
||||
code = "/* Auto Generated code - Do not edit.*/\n"
|
||||
code += "#include <stdio.h> \n"
|
||||
|
||||
code += "#include \"cutlass/gemm/device/gemm_batched.h\" \n"
|
||||
code += "#include \"cutlass/cutlass.h\" \n"
|
||||
|
||||
code += "#include \"../cutlass_irrelevant.h\" \n"
|
||||
code += "#include \"../cutlass_verify.h\" \n"
|
||||
|
||||
code += "#include \"leaky_bias.h\" \n"
|
||||
|
||||
code += "#include \"utils.h\" \n"
|
||||
|
||||
|
||||
|
||||
code += "int main(int args, char * argv[]) {\n"
|
||||
code += " " + "int M = atoi(argv[1]);\n"
|
||||
code += " " + "int K0 = " + str(self.fuse_gemm_info[0]['mnk'][0]) + ";\n"
|
||||
code += " " + "if(args == 3);\n"
|
||||
code += " " + " " + "K0 = atoi(argv[2]);\n"
|
||||
code += " " + "int B = 1;\n"
|
||||
code += " " + "if(args == 4);\n"
|
||||
code += " " + " " + "B = atoi(argv[3]);\n"
|
||||
|
||||
code += " " + "srand(1234UL);\n"
|
||||
code += " " + "int device_id = 0;\n"
|
||||
code += " " + "cudaGetDevice(&device_id);\n"
|
||||
code += " " + "cudaDeviceProp prop;\n"
|
||||
code += " " + "cudaGetDeviceProperties(&prop, device_id);\n"
|
||||
code += " " + "int sm = prop.major *10 + prop.minor;\n"
|
||||
code += "using ElementCompute = cutlass::half_t;\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("ElementCompute alpha", i) + " = ElementCompute(1);\n"
|
||||
addbias = helper.get_epilogue_add_bias_or_not( self.fuse_gemm_info[i])
|
||||
if addbias:
|
||||
code += " " + helper.var_idx("ElementCompute beta", i) + " = ElementCompute(1);\n"
|
||||
else:
|
||||
code += " " + helper.var_idx("ElementCompute beta", i) + " = ElementCompute(0);\n"
|
||||
|
||||
code += " " + "size_t flops = 0;\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
m = self.fuse_gemm_info[i]['mnk'][0]
|
||||
n = self.fuse_gemm_info[i]['mnk'][1]
|
||||
k = self.fuse_gemm_info[i]['mnk'][2]
|
||||
|
||||
bias_shape = helper.get_epilogue_bias_shape(self.fuse_gemm_info[i])
|
||||
|
||||
this_k = "K0"
|
||||
if (i > 0):
|
||||
this_k = str(k)
|
||||
|
||||
code += " " + "flops += size_t(2) * size_t(M) * size_t(B) * " + "size_t(" + str(n) + ") * size_t(" + this_k + ");\n"
|
||||
|
||||
code += " " + helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + "(" + "M" + ", " + str(n) + ", " + this_k + ");\n"
|
||||
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_A", i) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".m() * problem_size_", i) + ".k());\n"
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_B", i) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".n() * problem_size_", i) + ".k());\n"
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_C", i) + "(B * " + str(bias_shape[0]) + " * " + str(bias_shape[1]) + ");\n"
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_D_cutlass_ref", i) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".m() * problem_size_", i) + ".n());\n"
|
||||
|
||||
code += " " + helper.var_idx("Mat_A", i) + ".init();\n"
|
||||
code += " " + helper.var_idx("Mat_B", i) + ".init();\n"
|
||||
code += " " + helper.var_idx("Mat_C", i) + ".init();\n"
|
||||
|
||||
|
||||
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_D", self.b2b_num - 1) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".m() * problem_size_",self.b2b_num - 1) + ".n());\n"
|
||||
|
||||
params = []
|
||||
params.append("M")
|
||||
params.append("B")
|
||||
|
||||
params.append("Mat_A0.device_ptr")
|
||||
for i in range(self.b2b_num):
|
||||
params.append(helper.var_idx("Mat_B", i) + ".device_ptr")
|
||||
params.append(helper.var_idx("Mat_C", i) + ".device_ptr")
|
||||
if i != self.b2b_num-1:
|
||||
params.append(helper.var_idx("Mat_D_cutlass_ref", i) + ".device_ptr")
|
||||
params.append(helper.var_idx("Mat_D", self.b2b_num - 1) + ".device_ptr")
|
||||
|
||||
code += " " + "Param arguments = {\n"
|
||||
code += " " + " " + "M,\n"
|
||||
code += " " + " " + "K0,\n"
|
||||
code += " " + " " + "B,\n"
|
||||
|
||||
code += " " + " " + "reinterpret_cast<const void*>(Mat_A0.device_ptr),\n"
|
||||
cnt = 1
|
||||
for i in range(self.b2b_num):
|
||||
bias_flag = helper.get_epilogue_add_bias_or_not( self.fuse_gemm_info[i])
|
||||
code += " " + " " + "reinterpret_cast<const void*>(" + helper.var_idx("Mat_B", i) + ".device_ptr" + "),\n"
|
||||
cnt += 1
|
||||
if bias_flag:
|
||||
code += " " + " " + "reinterpret_cast<const void*>(" + helper.var_idx("Mat_C", i) + ".device_ptr" + "),\n"
|
||||
cnt += 1
|
||||
else:
|
||||
code += " " + " " + "reinterpret_cast<const void*>(NULL),\n"
|
||||
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_value = str(arg[2])
|
||||
|
||||
code += " " + " " + helper.type_2_cutlass_type(acc_tp) + "(" + arg_value + "),\n"
|
||||
|
||||
if i != self.b2b_num - 1:
|
||||
code += " " + " " + "reinterpret_cast<void*>(" + helper.var_idx("Mat_D_cutlass_ref", i) + ".device_ptr" + "),\n"
|
||||
else:
|
||||
code += " " + " " + "reinterpret_cast<void*>(" + helper.var_idx("Mat_D", i) + ".device_ptr" + ")};\n"
|
||||
|
||||
|
||||
|
||||
|
||||
code += " " + "TI(FUSED_CUTLASS);\n"
|
||||
code += " " + "for(int i = 0; i < 100; i++){\n"
|
||||
code += " " + " " + "one_api(arguments, sm, NULL);\n"
|
||||
|
||||
code += " " + "}\n"
|
||||
code += " " + "TO(FUSED_CUTLASS, \"FUSED_CUTLASS\", 100);\n"
|
||||
|
||||
code += "\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
|
||||
N_str = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
code_this += " " + helper.var_idx("typename Gemm", i) + helper.var_idx("::Arguments arguments_", i) + "{\n"
|
||||
code_this += " " + " " + helper.var_idx("problem_size_", i) + ",\n"
|
||||
ldmA = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
ldmA = "K0"
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
ldmB = "K0"
|
||||
ldmC = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
ldmBias = str(helper.get_epilogue_bias_ldm(self.fuse_gemm_info[i]))
|
||||
|
||||
if self.fuse_gemm_info[i]['A_format'] is 'Col':
|
||||
ldmA = "M"
|
||||
if self.fuse_gemm_info[i]['B_format'] is 'Row':
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
if self.fuse_gemm_info[i]['C_format'] is 'Col':
|
||||
ldmC = "M"
|
||||
|
||||
if i == 0:
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("Mat_A", i) + ".device_ptr), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
else:
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("Mat_D_cutlass_ref", i - 1) + ".device_ptr), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + "*>(" + helper.var_idx("Mat_B", i) + ".device_ptr), " + ldmB + "}, " + N_str + " * " + ldmB + ",\n"
|
||||
|
||||
M_bias = str(helper.get_epilogue_bias_shape(self.fuse_gemm_info[i])[0])
|
||||
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("Mat_C", i) + ".device_ptr), " + ldmBias + "}, " + M_bias + " * " + N_str + ",\n"
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("Mat_D_cutlass_ref", i) + ".device_ptr), " + ldmC + "}, " + "M * " + ldmC + ",\n"
|
||||
code_this += " " + " " + "{ " + helper.var_idx("alpha", i) + ", " + helper.var_idx("beta", i)
|
||||
for epilogue_arg in helper.get_epilogue_args(self.fuse_gemm_info[i]):
|
||||
arg_value = str(epilogue_arg[2])
|
||||
code_this += ", " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(" + str(arg_value) + ")"
|
||||
code_this += " " + " },\n"
|
||||
code_this += " " + " " + "B};\n"
|
||||
|
||||
code += code_this
|
||||
|
||||
|
||||
|
||||
code += " " + "TI(UNFUSED_CUTLASS);\n"
|
||||
code += " " + "for(int i = 0; i < 100; i++){\n"
|
||||
code += " " + " " + self.gen_class_name + "_verify(\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + " " + " " + helper.var_idx("arguments_", i) + ",\n"
|
||||
code += " " + " " + " " + "NULL);\n"
|
||||
|
||||
code += " " + "}\n"
|
||||
code += " " + "TO(UNFUSED_CUTLASS, \"UNFUSED_CUTLASS\", 100);\n"
|
||||
|
||||
code += " " + helper.var_idx("Mat_D_cutlass_ref", self.b2b_num - 1) + ".d2h();\n"
|
||||
code += " " + helper.var_idx("Mat_D", self.b2b_num - 1) + ".d2h();\n"
|
||||
code += " " + helper.var_idx("check_result(Mat_D_cutlass_ref", self.b2b_num - 1) + helper.var_idx(".host_ptr, Mat_D", self.b2b_num - 1) \
|
||||
+ helper.var_idx(".host_ptr, Mat_D", self.b2b_num - 1) + ".elements);\n"
|
||||
|
||||
code += "\n\n}\n"
|
||||
|
||||
with open(self.sample_dir + "sample.cu", "w+") as f:
|
||||
f.write(code)
|
||||
1013
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_threadblock.py
Normal file
1013
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_threadblock.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,456 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
import gen_ir as ir
|
||||
|
||||
class gen_turing_impl:
|
||||
def __init__(self,fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.class_name = gen_class_name
|
||||
self.gen_class_name = gen_class_name + "_turing_impl"
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.output_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
self.gen_turing_unfused = gen_volta_turing_fuse_act_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
|
||||
def gen_using(self):
|
||||
code_using = "using b2b_gemm = typename cutlass::gemm::device::" + self.class_name + "<cutlass::half_t>;"
|
||||
|
||||
return code_using + "\n"
|
||||
|
||||
def gen_initialize(self):
|
||||
code = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " alpha", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(1);\n"
|
||||
beta = "(1)"
|
||||
|
||||
if helper.get_epilogue_add_bias_or_not(self.fuse_gemm_info[i]) is False:
|
||||
beta = "(0)"
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " beta", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + beta + ";\n"
|
||||
k_str = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
k_str = "K0"
|
||||
code_this += helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + "(M, " + str(self.fuse_gemm_info[i]['mnk'][1]) + ", " + k_str + ");\n"
|
||||
code += code_this
|
||||
code += "typename b2b_gemm::Arguments arguments{\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("problem_size_", i) + ",\n"
|
||||
|
||||
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("A", 0) + "), " + helper.var_idx("problem_size_", 0) + ".k()},\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
ldmB = "K0"
|
||||
|
||||
if self.fuse_gemm_info[i]['B_format'] is 'Row':
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
ldmC = str(helper.get_epilogue_bias_ldm(self.fuse_gemm_info[i]))
|
||||
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + "*>(" + helper.var_idx("B", i) + "), " + ldmB + "},\n"
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("C", i) + "), " + ldmC + "},\n"
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("D", self.b2b_num -1) + "), " + helper.var_idx("problem_size_", self.b2b_num - 1) + ".n()},\n"
|
||||
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + "{ " + helper.var_idx("alpha", i) + ", " + helper.var_idx("beta", i)
|
||||
for epilogue_arg in helper.get_epilogue_args(self.fuse_gemm_info[i]):
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + epilogue_arg[1]
|
||||
code += ", " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(" + str(arg_name) + ")"
|
||||
code += "},\n"
|
||||
code += " " + "Batch};\n\n"
|
||||
|
||||
code += " " "b2b_gemm gemm_op;\n"
|
||||
code += " " + "gemm_op.initialize(arguments);\n"
|
||||
return code + "\n"
|
||||
|
||||
|
||||
|
||||
def gen_run(self):
|
||||
code = " " + "gemm_op(stream);\n"
|
||||
|
||||
return code
|
||||
|
||||
def gen_wrapper(self):
|
||||
code_body = ""
|
||||
|
||||
arg_lists = []
|
||||
arg_lists.append(["int", "M"])
|
||||
arg_lists.append(["int", "K0"])
|
||||
arg_lists.append(["int", "Batch"])
|
||||
arg_lists.append(["void*", helper.var_idx("A", 0)])
|
||||
for i in range(self.b2b_num):
|
||||
arg_lists.append(["void*", helper.var_idx("B", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("C", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("D", i)])
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_tp = arg[0]
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
arg_lists.append([arg_tp, arg_name])
|
||||
|
||||
if self.b2b_num == 1:
|
||||
code_body += self.gen_turing_unfused.gen_using(False) #False -> Turing, True -> Volta
|
||||
code_body += self.gen_turing_unfused.gen_initialize()
|
||||
code_body += self.gen_turing_unfused.gen_run()
|
||||
else:
|
||||
code_body += self.gen_using()
|
||||
code_body += self.gen_initialize()
|
||||
code_body += self.gen_run()
|
||||
|
||||
code = ir.gen_func(self.gen_class_name, arg_lists, code_body)
|
||||
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
|
||||
code = self.gen_wrapper()
|
||||
helper.write_2_headfile("turing_impl.h", self.output_dir, self.user_header_file + "\n" + code)
|
||||
|
||||
class gen_volta_turing_fuse_act_impl:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.gen_class_name = gen_class_name + "_volta_impl"
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.output_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
def perf_tiling(self, layer_mnk):
|
||||
mnk = layer_mnk[:]
|
||||
block_tile = mnk[:]
|
||||
block_tile[2] = 32 # force the K tile to be 32
|
||||
|
||||
# M tile gen
|
||||
block_tile[0] = 32
|
||||
|
||||
# N tile gen
|
||||
if mnk[1] > 128:
|
||||
block_tile[1] = 256
|
||||
elif mnk[1] > 64:
|
||||
block_tile[1] = 128
|
||||
elif mnk[1] > 32:
|
||||
block_tile[1] = 64
|
||||
else :
|
||||
block_tile[1] = 32
|
||||
|
||||
warp_tile = block_tile[:]
|
||||
if block_tile[1] == 256:
|
||||
warp_tile[1] = 64
|
||||
elif block_tile[1] == 128:
|
||||
warp_tile[1] = 32
|
||||
elif block_tile[1] == 64:
|
||||
warp_tile[1] = 32
|
||||
else :
|
||||
warp_tile[1] = 32
|
||||
|
||||
warp_tile[0] = 32
|
||||
|
||||
return block_tile, warp_tile
|
||||
|
||||
|
||||
def process_epilogue(self, epilogue_tp, n, C_tp, Acc_tp):
|
||||
epilogue_setted_type = epilogue_tp
|
||||
cutlass_epilogue_name = "LinearCombinationRelu"
|
||||
if epilogue_setted_type.lower() == 'leakyrelu':
|
||||
cutlass_epilogue_name = "LinearCombinationLeakyRelu"
|
||||
elif epilogue_setted_type.lower() == 'identity':
|
||||
cutlass_epilogue_name = "LinearCombination"
|
||||
|
||||
|
||||
n_mod_8 = n % 4
|
||||
N_align_elements = 1
|
||||
if n_mod_8 == 0:
|
||||
N_align_elements = 8
|
||||
elif n_mod_8 == 4:
|
||||
N_align_elements = 4
|
||||
elif n_mod_8 == 2 or n_mod_8 == 6:
|
||||
N_align_elements = 2
|
||||
|
||||
epilogue_str = "cutlass::epilogue::thread::" + cutlass_epilogue_name+ "<" + C_tp + ", " + str(N_align_elements) + ", " + Acc_tp + ", " + Acc_tp + ">"
|
||||
|
||||
return epilogue_str
|
||||
|
||||
def gen_using(self, volta = True):
|
||||
code_using = ""
|
||||
volta_arch = "cutlass::arch::Sm70"
|
||||
volta_tc = "cutlass::gemm::GemmShape<8, 8, 4>"
|
||||
|
||||
turing_arch = "cutlass::arch::Sm75"
|
||||
turing_tc = "cutlass::gemm::GemmShape<16, 8, 8>"
|
||||
|
||||
arch = ""
|
||||
tc = ""
|
||||
if volta:
|
||||
arch = volta_arch
|
||||
tc = volta_tc
|
||||
else:
|
||||
arch = turing_arch
|
||||
tc = turing_tc
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
|
||||
k = self.fuse_gemm_info[i]['mnk'][2]
|
||||
|
||||
k_mod_8 = k % 4
|
||||
ab_ldm = 1
|
||||
if k_mod_8 == 0:
|
||||
ab_ldm = 8
|
||||
elif k_mod_8 == 4:
|
||||
ab_ldm = 4
|
||||
elif k_mod_8 == 2 or k_mod_8 == 6:
|
||||
ab_ldm = 2
|
||||
|
||||
block_tile, warp_tile = self.perf_tiling(self.fuse_gemm_info[i]['mnk'])
|
||||
|
||||
this_gemm_config = helper.var_idx("using Gemm", i) + " = cutlass::gemm::device::GemmBatched<\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_format']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_format']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_format']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + ",\n"
|
||||
this_gemm_config += " " + "cutlass::arch::OpClassTensorOp,\n"
|
||||
this_gemm_config += " " + arch + ",\n"
|
||||
this_gemm_config += " " + "cutlass::gemm::GemmShape<" + str(block_tile[0]) + ", " + str(block_tile[1]) + ", " + str(block_tile[2]) + ">,\n"
|
||||
this_gemm_config += " " + "cutlass::gemm::GemmShape<" + str(warp_tile[0]) + ", " + str(warp_tile[1]) + ", " + str(warp_tile[2]) + ">,\n"
|
||||
this_gemm_config += " " + tc + ",\n"
|
||||
this_gemm_config += " " + self.process_epilogue(helper.get_epilogue_tp(self.fuse_gemm_info[i]), self.fuse_gemm_info[i]['mnk'][1], helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']), helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp'])) + ",\n"
|
||||
this_gemm_config += " " + "cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,\n"
|
||||
this_gemm_config += " " + "2,\n"
|
||||
this_gemm_config += " " + str(ab_ldm) + ",\n"
|
||||
this_gemm_config += " " + str(ab_ldm) + ">;\n"
|
||||
|
||||
code_using += this_gemm_config + "\n"
|
||||
|
||||
return code_using + "\n"
|
||||
|
||||
def gen_initialize(self):
|
||||
code = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
|
||||
N_str = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " alpha", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(1);\n"
|
||||
beta = "(1)"
|
||||
if helper.get_epilogue_add_bias_or_not( self.fuse_gemm_info[i]) is False:
|
||||
beta = "(0)"
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " beta", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + beta + ";\n"
|
||||
|
||||
k_str = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
k_str = "K0"
|
||||
code_this += helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + "(M, " + str(self.fuse_gemm_info[i]['mnk'][1]) + ", " + k_str + ");\n"
|
||||
code_this += helper.var_idx("typename Gemm", i) + helper.var_idx("::Arguments arguments_", i) + "{\n"
|
||||
code_this += " " + helper.var_idx("problem_size_", i) + ",\n"
|
||||
ldmA = k_str
|
||||
ldmB = k_str
|
||||
ldmC = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
ldmBias = str(helper.get_epilogue_bias_ldm(self.fuse_gemm_info[i]))
|
||||
|
||||
if self.fuse_gemm_info[i]['A_format'] is 'Col':
|
||||
ldmA = "M"
|
||||
if self.fuse_gemm_info[i]['B_format'] is 'Row':
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
if self.fuse_gemm_info[i]['C_format'] is 'Col':
|
||||
ldmC = "M"
|
||||
|
||||
if i == 0:
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("A", i) + "), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
else:
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("D", i - 1) + "), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + "*>(" + helper.var_idx("B", i) + "), " + ldmB + "}, " + N_str + " * " + ldmB + ",\n"
|
||||
|
||||
M_bias = str(helper.get_epilogue_bias_shape(self.fuse_gemm_info[i])[0])
|
||||
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("C", i) + "), " + ldmBias + "}, " + M_bias + " * " + N_str + ",\n"
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("D", i) + "), " + ldmC + "}, " + "M * " + ldmC + ",\n"
|
||||
code_this += " " + "{ " + helper.var_idx("alpha", i) + ", " + helper.var_idx("beta", i)
|
||||
for epilogue_arg in helper.get_epilogue_args(self.fuse_gemm_info[i]):
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + epilogue_arg[1]
|
||||
code_this += ", " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(" + str(arg_name) + ")"
|
||||
code_this += " },\n"
|
||||
code_this += " " + "Batch};\n"
|
||||
|
||||
code_this += " " + helper.var_idx("Gemm", i) + helper.var_idx(" gemm_op_", i) + ";\n"
|
||||
code_this += " " + helper.var_idx("gemm_op_", i) + helper.var_idx(".initialize(arguments_", i) + ", nullptr);\n"
|
||||
|
||||
code += code_this + "\n"
|
||||
return code + "\n"
|
||||
|
||||
|
||||
def gen_run(self):
|
||||
code = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
code_this += " " + helper.var_idx("gemm_op_", i) + "(stream);\n"
|
||||
|
||||
code += code_this
|
||||
return code
|
||||
|
||||
def gen_wrapper(self):
|
||||
code_body = ""
|
||||
|
||||
arg_lists = []
|
||||
arg_lists.append(["int", "M"])
|
||||
arg_lists.append(["int", "K0"])
|
||||
arg_lists.append(["int", "Batch"])
|
||||
arg_lists.append(["void*", helper.var_idx("A", 0)])
|
||||
for i in range(self.b2b_num):
|
||||
arg_lists.append(["void*", helper.var_idx("B", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("C", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("D", i)])
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_tp = arg[0]
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
arg_lists.append([arg_tp, arg_name])
|
||||
code_body += self.gen_using()
|
||||
code_body += self.gen_initialize()
|
||||
code_body += self.gen_run()
|
||||
|
||||
code = ir.gen_func(self.gen_class_name, arg_lists, code_body)
|
||||
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
code = self.gen_wrapper()
|
||||
helper.write_2_headfile("volta_impl.h", self.output_dir, self.user_header_file + "\n" + code)
|
||||
|
||||
class gen_one_API:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.gen_class_name = gen_class_name
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.output_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
self.gen_volta = gen_volta_turing_fuse_act_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
|
||||
self.gen_turing = gen_turing_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
|
||||
def gen_CUTLASS_irrelevant_API(self):
|
||||
code = ""
|
||||
code += "#include <cuda_runtime.h>\n"
|
||||
code += "#include <assert.h>\n"
|
||||
|
||||
param_name = "Fused" + str(self.b2b_num) + "xGemm_"
|
||||
for i in range(self.b2b_num):
|
||||
param_name += str(self.fuse_gemm_info[i]['mnk'][1]) + "_"
|
||||
param_name += "Params"
|
||||
params = ""
|
||||
params += " " + "int M;\n"
|
||||
params += " " + "int K0;\n"
|
||||
params += " " + "int Batch;\n"
|
||||
params += " " + "const void* A0;\n"
|
||||
for i in range(self.b2b_num):
|
||||
params += " " + "const void* " + helper.var_idx("B", i) + ";\n"
|
||||
params += " " + "const void* " + helper.var_idx("C", i) + ";\n"
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_tp = arg[0]
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
params += " " + arg_tp + " " + arg_name + ";\n"
|
||||
params += " " + "void* " + helper.var_idx("D", i) + ";\n"
|
||||
code += ir.gen_struct(param_name, params)
|
||||
code += "using Param = " + param_name + ";\n"
|
||||
code += "void one_api( const Param & param, int sm, cudaStream_t stream);\n"
|
||||
|
||||
|
||||
return code
|
||||
|
||||
def gen_one_api(self):
|
||||
code = ""
|
||||
code += "/* Auto Generated code - Do not edit.*/\n"
|
||||
code += "#include \"cutlass_irrelevant.h\"\n"
|
||||
code += "#include \"api.h\"\n"
|
||||
code += "void one_api( const Param & param, int sm, cudaStream_t stream) {\n"
|
||||
|
||||
code += " " + "if (sm == 70) \n"
|
||||
code += " " + " " + self.gen_class_name + "_volta_impl(param.M, param.K0, param.Batch, const_cast<void*>(param.A0), "
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx("const_cast<void*>(param.B", i) + "), "
|
||||
code += helper.var_idx("const_cast<void*>(param.C", i) + "), "
|
||||
code += helper.var_idx("param.D", i) + ", "
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
code += "param." + arg_name + ", "
|
||||
code += "stream);\n"
|
||||
code += " " + "else if(sm >= 75) \n"
|
||||
code += " " + " " + self.gen_class_name + "_turing_impl(param.M, param.K0, param.Batch, const_cast<void*>(param.A0), "
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx("const_cast<void*>(param.B", i) + "), "
|
||||
code += helper.var_idx("const_cast<void*>(param.C", i) + "), "
|
||||
code += helper.var_idx("param.D", i) + ", "
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
code += "param." + arg_name + ", "
|
||||
code += "stream);\n"
|
||||
code += " " + "else assert(0);\n"
|
||||
code += "}\n"
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
|
||||
turing_code = self.gen_turing.gen_wrapper()
|
||||
volta_code = self.gen_volta.gen_wrapper()
|
||||
cutlass_irrelevant_code = self.gen_CUTLASS_irrelevant_API()
|
||||
|
||||
one_api_code = self.gen_one_api()
|
||||
with open(self.output_dir + "one_api.cu", "w+") as f:
|
||||
f.write(one_api_code)
|
||||
|
||||
helper.write_2_headfile("cutlass_irrelevant.h", self.output_dir, cutlass_irrelevant_code)
|
||||
|
||||
helper.write_2_headfile("api.h", self.output_dir, self.user_header_file + "\n" + turing_code + volta_code)
|
||||
92
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_verify.py
Normal file
92
examples/44_multi_gemm_ir_and_codegen/ir_gen/gen_verify.py
Normal file
@@ -0,0 +1,92 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
import gen_ir as ir
|
||||
|
||||
import gen_turing_and_volta as gen_basic
|
||||
|
||||
|
||||
class gen_verify:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.name = gen_class_name + "_verify"
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
self.params = []
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.seperate_cutlass = gen_basic.gen_volta_turing_fuse_act_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
self.gen_params()
|
||||
self.output_dir = output_dir
|
||||
|
||||
|
||||
def gen_code(self):
|
||||
code = ""
|
||||
code += self.user_header_file
|
||||
code += self.seperate_cutlass.gen_using(False) #False -> Turing, True -> Volta
|
||||
|
||||
code_body = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_body += " " + helper.var_idx("Gemm", i) + helper.var_idx(" gemm_op_", i) + ";\n"
|
||||
code_body += " " + helper.var_idx("gemm_op_", i) + helper.var_idx(".initialize(Arguments_", i) + ", nullptr);\n"
|
||||
|
||||
code_body += self.seperate_cutlass.gen_run()
|
||||
|
||||
code += ir.gen_func(self.name, self.params, code_body)
|
||||
helper.write_2_headfile("cutlass_verify.h", self.output_dir, code)
|
||||
|
||||
|
||||
def gen_params(self):
|
||||
for i in range(self.b2b_num):
|
||||
self.params.append(
|
||||
(
|
||||
helper.var_idx("typename Gemm", i)+ "::Arguments",
|
||||
helper.var_idx("Arguments_", i)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_params(self, declartion = True):
|
||||
code = ""
|
||||
if declartion:
|
||||
for param in self.params:
|
||||
code += param[0] + " " + param[1] + ";\n"
|
||||
|
||||
return code
|
||||
|
||||
|
||||
def gen_initialize():
|
||||
code = ""
|
||||
initialize_code = self.seperate_cutlass.gen_initialize()
|
||||
|
||||
code = ir.gen_func("initialize", [[]])
|
||||
52
examples/44_multi_gemm_ir_and_codegen/ir_gen/generate.sh
Executable file
52
examples/44_multi_gemm_ir_and_codegen/ir_gen/generate.sh
Executable file
@@ -0,0 +1,52 @@
|
||||
#!/bin/bash
|
||||
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
NUM_ARGS=3
|
||||
if [ $# -ne $NUM_ARGS ]; then
|
||||
echo "Usage: $0 <config_file> <output_directory> <cutlass_directory>"
|
||||
echo " config_file: JSON file containing configuration to run"
|
||||
echo " output_directory: directory to store results"
|
||||
echo " cutlass_directory: directory containing cutlass source"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
config_file=$1
|
||||
output_dir=$2
|
||||
cutlass_dir=$3
|
||||
|
||||
python3 gen_all_code.py \
|
||||
--config-file $config_file \
|
||||
--gen-name FusedMultiGemmForward \
|
||||
--output-dir $output_dir \
|
||||
--cutlass-dir $cutlass_dir
|
||||
135
examples/44_multi_gemm_ir_and_codegen/ir_gen/helper.py
Normal file
135
examples/44_multi_gemm_ir_and_codegen/ir_gen/helper.py
Normal file
@@ -0,0 +1,135 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
def type_2_cutlass_type(input_type = "fp16"):
|
||||
# float point type
|
||||
if input_type == "fp32":
|
||||
return "float"
|
||||
if input_type == "bf16":
|
||||
return "cutlass::bfloat16_t"
|
||||
if input_type == "fp16":
|
||||
return "cutlass::half_t"
|
||||
|
||||
# integer type
|
||||
if(input_type == "int32"):
|
||||
return "int32_t"
|
||||
if(input_type == "int8"):
|
||||
return "int8_t"
|
||||
|
||||
if input_type == 'Row':
|
||||
return 'cutlass::layout::RowMajor'
|
||||
if input_type == 'Col':
|
||||
return 'cutlass::layout::ColumnMajor'
|
||||
|
||||
def cvt_2_cutlass_shape(gemm_shape):
|
||||
# gemm shape
|
||||
if len(gemm_shape) == 3:
|
||||
val = "cutlass::gemm::GemmShape<" \
|
||||
+ str(gemm_shape[0]) + ", " \
|
||||
+ str(gemm_shape[1]) + ", " \
|
||||
+ str(gemm_shape[2]) + ">"
|
||||
return val
|
||||
|
||||
|
||||
def write_2_headfile(filename, file_dir, string):
|
||||
with open(file_dir + filename, 'w') as f:
|
||||
f.write("/* Auto Generated code - Do not edit.*/\n\n\n#pragma once\n" + string)
|
||||
|
||||
def var_idx(varaiable, index):
|
||||
return varaiable + str(index)
|
||||
|
||||
|
||||
def list_2_string(input_list, ):
|
||||
rtn_string = ""
|
||||
|
||||
cnt = 0
|
||||
|
||||
for element in input_list:
|
||||
final = ", \n"
|
||||
if cnt == len(input_list) - 1:
|
||||
final = "\n"
|
||||
cnt += 1
|
||||
rtn_string += str(element) + final
|
||||
|
||||
return rtn_string
|
||||
|
||||
|
||||
def get_epilouge_info(layer_info):
|
||||
return layer_info['epilogue']
|
||||
|
||||
def get_epilogue_tp(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['tp']
|
||||
|
||||
def get_epilogue_add_bias_or_not(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['bias']['addbias']
|
||||
|
||||
def get_epilogue_add_bias_tp(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['bias']['bias_tp']
|
||||
|
||||
def get_epilogue_args(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['args']
|
||||
|
||||
def get_epilogue_bias_shape(layer_info):
|
||||
bias_tp = get_epilogue_add_bias_tp(layer_info).lower()
|
||||
mn_shape = layer_info['mnk'][:-1]
|
||||
|
||||
if bias_tp == 'mat':
|
||||
mn_shape[0] = 'M'
|
||||
return mn_shape
|
||||
elif bias_tp == 'vec':
|
||||
mn_shape[0] = 1
|
||||
return mn_shape
|
||||
else:
|
||||
assert(0)
|
||||
|
||||
def get_epilogue_bias_ldm(layer_info):
|
||||
bias_tp = get_epilogue_add_bias_tp(layer_info).lower()
|
||||
mn_shape = layer_info['mnk'][:-1]
|
||||
|
||||
c_layout = layer_info['C_format'].lower()
|
||||
|
||||
if c_layout != 'row':
|
||||
assert(0)
|
||||
|
||||
if bias_tp == 'mat':
|
||||
return mn_shape[1]
|
||||
elif bias_tp == 'vec':
|
||||
return 0
|
||||
else:
|
||||
assert(0)
|
||||
|
||||
def get_epilogue_compute_tp(layer_info):
|
||||
return layer_info['Acc_tp']
|
||||
@@ -0,0 +1,67 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import os
|
||||
|
||||
class replace_fix_impl:
|
||||
def __init__(self, src_dir, dst_dir, cutlass_deps_root):
|
||||
self.src_dir = src_dir
|
||||
self.dst_dir = dst_dir
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
|
||||
|
||||
|
||||
def gen_code(self):
|
||||
for sub_dir in os.walk(self.src_dir):
|
||||
files_in_sub_dir = sub_dir[2]
|
||||
|
||||
src_dirs = sub_dir[0]
|
||||
output_dirs = self.dst_dir + sub_dir[0][len(self.src_dir):]
|
||||
|
||||
if not os.path.exists(output_dirs):
|
||||
os.mkdir(output_dirs)
|
||||
|
||||
for f in files_in_sub_dir:
|
||||
with open(src_dirs +"/" + f, 'r') as current_file:
|
||||
output_lines = []
|
||||
lines = current_file.readlines()
|
||||
|
||||
for line in lines:
|
||||
if(len(line) >= len("#include \"cutlass") and line[:len("#include \"cutlass")] == "#include \"cutlass"):
|
||||
new_line = "#include \"" + self.cutlass_deps_root + line[len("#include \""):]
|
||||
# print(new_line)
|
||||
output_lines.append(new_line)
|
||||
else:
|
||||
output_lines.append(line)
|
||||
|
||||
with open(output_dirs + "/" + f, "w+") as dest_file:
|
||||
dest_file.writelines(output_lines)
|
||||
292
examples/44_multi_gemm_ir_and_codegen/leaky_bias.h
Normal file
292
examples/44_multi_gemm_ir_and_codegen/leaky_bias.h
Normal file
@@ -0,0 +1,292 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
#pragma once
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
template <typename T>
|
||||
__device__
|
||||
T add(T const & a, T const &b){
|
||||
return (a + b);
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__
|
||||
half2 add(half2 const & a, half2 const &b){
|
||||
return (__hadd2(a,b));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct RELU{
|
||||
__device__
|
||||
T operator()(T const & a){
|
||||
return a > T(0) ? a : T(0);
|
||||
}
|
||||
__device__
|
||||
half2 operator()(half2 const & a){
|
||||
float2 a_fp32x2 = __half22float2(a);
|
||||
a_fp32x2.x = a_fp32x2.x > 0.f ? a_fp32x2.x : 0.f;
|
||||
a_fp32x2.y = a_fp32x2.y > 0.f ? a_fp32x2.y : 0.f;
|
||||
if(a_fp32x2.x < 0.f || a_fp32x2.y < 0.f)
|
||||
printf(" %f %f\n", a_fp32x2.x ,a_fp32x2.y);
|
||||
return __float22half2_rn(a_fp32x2);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct LEAKY_RELU{
|
||||
__device__
|
||||
T operator()(T const & a, T const & scale = half(1)){
|
||||
return a > T(0) ? a : scale * a;
|
||||
}
|
||||
__device__
|
||||
half2 operator()(half2 const & a, half const & scale = half(1)){
|
||||
half2 zero = __half2half2(half(0));
|
||||
half2 gt_zero = __hge2(a, zero);
|
||||
half2 le_zero = __hle2(a, zero);
|
||||
|
||||
|
||||
half2 scale_f16x2 = __half2half2(scale);
|
||||
half2 mask_scale_f16x2 = __hfma2(le_zero, scale_f16x2, gt_zero);
|
||||
return __hmul2(a, mask_scale_f16x2);
|
||||
}
|
||||
};
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
__global__ void leaky_and_activation(half* inout, half* bias, half scale, bool mat_bias){
|
||||
|
||||
constexpr bool N_MOD_2 = N & 1 ? false : true;
|
||||
|
||||
using Access_tp = typename std::conditional<N_MOD_2, half2, half>::type;
|
||||
|
||||
constexpr int Access_elements = sizeof(Access_tp) / sizeof(half);
|
||||
|
||||
constexpr int iter = (N + (BLOCKDIM * Access_elements) - 1 ) / (BLOCKDIM * Access_elements);
|
||||
|
||||
LEAKY_RELU<half> Act;
|
||||
Access_tp src_v[iter];
|
||||
Access_tp bias_v[iter];
|
||||
|
||||
int batch_id = blockIdx.y;
|
||||
int batch_offset = batch_id * gridDim.x * N;
|
||||
|
||||
for(int i = 0; i < iter; i++){
|
||||
int idx = (i * BLOCKDIM + threadIdx.x) * Access_elements;
|
||||
if (idx < N){
|
||||
src_v[i] = *reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset);
|
||||
if (mat_bias)
|
||||
bias_v[i] = *reinterpret_cast<Access_tp*>(bias + blockIdx.x * N + idx + batch_offset);
|
||||
else
|
||||
bias_v[i] = *reinterpret_cast<Access_tp*>(bias + idx + batch_id * N);
|
||||
*reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset) = Act(add(src_v[i],bias_v[i]),scale);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
__global__ void leaky_and_activation(half* inout, half scale){
|
||||
|
||||
constexpr bool N_MOD_2 = N & 1 ? false : true;
|
||||
|
||||
using Access_tp = typename std::conditional<N_MOD_2, half2, half>::type;
|
||||
|
||||
constexpr int Access_elements = sizeof(Access_tp) / sizeof(half);
|
||||
|
||||
constexpr int iter = (N + (BLOCKDIM * Access_elements) - 1 ) / (BLOCKDIM * Access_elements);
|
||||
|
||||
int batch_id = blockIdx.y;
|
||||
int batch_offset = batch_id * gridDim.x * N;
|
||||
|
||||
LEAKY_RELU<half> Act;
|
||||
Access_tp src_v[iter];
|
||||
|
||||
for(int i = 0; i < iter; i++){
|
||||
int idx = (i * BLOCKDIM + threadIdx.x) * Access_elements;
|
||||
if (idx < N){
|
||||
src_v[i] = *reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset);
|
||||
*reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset) = Act(src_v[i], scale);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
void leaky_and_activation(half* inout, half* bias, int m, int b, half scale, bool mat_bias){
|
||||
|
||||
dim3 grid(m, b);
|
||||
if (bias == nullptr)
|
||||
leaky_and_activation<N, BLOCKDIM><<<grid , BLOCKDIM>>>(inout, scale);
|
||||
else
|
||||
leaky_and_activation<N, BLOCKDIM><<<grid , BLOCKDIM>>>(inout, bias, scale, mat_bias);
|
||||
}
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
__global__ void relu_and_activation(half* inout, half* bias, bool mat_bias){
|
||||
|
||||
constexpr bool N_MOD_2 = N & 1 ? false : true;
|
||||
|
||||
using Access_tp = typename std::conditional<N_MOD_2, half2, half>::type;
|
||||
|
||||
constexpr int Access_elements = sizeof(Access_tp) / sizeof(half);
|
||||
|
||||
constexpr int iter = (N + (BLOCKDIM * Access_elements) - 1 ) / (BLOCKDIM * Access_elements);
|
||||
|
||||
RELU<half> Act;
|
||||
Access_tp src_v[iter];
|
||||
Access_tp bias_v[iter];
|
||||
|
||||
int batch_id = blockIdx.y;
|
||||
int batch_offset = batch_id * gridDim.x * N;
|
||||
|
||||
for(int i = 0; i < iter; i++){
|
||||
int idx = (i * BLOCKDIM + threadIdx.x) * Access_elements;
|
||||
if (idx < N){
|
||||
src_v[i] = *reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset);
|
||||
if (mat_bias)
|
||||
bias_v[i] = *reinterpret_cast<Access_tp*>(bias + blockIdx.x * N + idx + batch_offset);
|
||||
else
|
||||
bias_v[i] = *reinterpret_cast<Access_tp*>(bias + idx + batch_id * N);
|
||||
*reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset) = Act(add(src_v[i],bias_v[i]));
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
__global__ void relu_and_activation(half* inout){
|
||||
|
||||
constexpr bool N_MOD_2 = N & 1 ? false : true;
|
||||
|
||||
using Access_tp = typename std::conditional<N_MOD_2, half2, half>::type;
|
||||
|
||||
constexpr int Access_elements = sizeof(Access_tp) / sizeof(half);
|
||||
|
||||
constexpr int iter = (N + (BLOCKDIM * Access_elements) - 1 ) / (BLOCKDIM * Access_elements);
|
||||
|
||||
int batch_id = blockIdx.y;
|
||||
int batch_offset = batch_id * gridDim.x * N;
|
||||
|
||||
RELU<half> Act;
|
||||
Access_tp src_v[iter];
|
||||
|
||||
for(int i = 0; i < iter; i++){
|
||||
int idx = (i * BLOCKDIM + threadIdx.x) * Access_elements;
|
||||
if (idx < N){
|
||||
src_v[i] = *reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset);
|
||||
*reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset) = Act(src_v[i]);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
void relu_and_activation(half* inout, half* bias, int m, int b, bool mat_bias){
|
||||
dim3 grid(m, b);
|
||||
if (bias == nullptr)
|
||||
relu_and_activation<N, BLOCKDIM><<<grid , BLOCKDIM>>>(inout);
|
||||
else
|
||||
relu_and_activation<N, BLOCKDIM><<<grid , BLOCKDIM>>>(inout, bias, mat_bias);
|
||||
}
|
||||
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
__global__ void identity_and_activation(half* inout, half* bias, bool mat_bias){
|
||||
|
||||
constexpr bool N_MOD_2 = N & 1 ? false : true;
|
||||
|
||||
using Access_tp = typename std::conditional<N_MOD_2, half2, half>::type;
|
||||
|
||||
constexpr int Access_elements = sizeof(Access_tp) / sizeof(half);
|
||||
|
||||
constexpr int iter = (N + (BLOCKDIM * Access_elements) - 1 ) / (BLOCKDIM * Access_elements);
|
||||
|
||||
int batch_id = blockIdx.y;
|
||||
int batch_offset = batch_id * gridDim.x * N;
|
||||
|
||||
Access_tp src_v[iter];
|
||||
Access_tp bias_v[iter];
|
||||
|
||||
for(int i = 0; i < iter; i++){
|
||||
int idx = (i * BLOCKDIM + threadIdx.x) * Access_elements;
|
||||
if (idx < N){
|
||||
src_v[i] = *reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset);
|
||||
if (mat_bias)
|
||||
bias_v[i] = *reinterpret_cast<Access_tp*>(bias + blockIdx.x * N + idx + batch_offset);
|
||||
else
|
||||
bias_v[i] = *reinterpret_cast<Access_tp*>(bias + idx + batch_id * N);
|
||||
*reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset) = (add(src_v[i],bias_v[i]));
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
__global__ void identity_and_activation(half* inout){
|
||||
|
||||
constexpr bool N_MOD_2 = N & 1 ? false : true;
|
||||
|
||||
using Access_tp = typename std::conditional<N_MOD_2, half2, half>::type;
|
||||
|
||||
constexpr int Access_elements = sizeof(Access_tp) / sizeof(half);
|
||||
|
||||
constexpr int iter = (N + (BLOCKDIM * Access_elements) - 1 ) / (BLOCKDIM * Access_elements);
|
||||
|
||||
int batch_id = blockIdx.y;
|
||||
int batch_offset = batch_id * gridDim.x * N;
|
||||
Access_tp src_v[iter];
|
||||
|
||||
for(int i = 0; i < iter; i++){
|
||||
int idx = (i * BLOCKDIM + threadIdx.x) * Access_elements;
|
||||
if (idx < N){
|
||||
src_v[i] = *reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset);
|
||||
*reinterpret_cast<Access_tp*>(inout + blockIdx.x * N + idx + batch_offset) = (src_v[i]);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
template <int N, int BLOCKDIM>
|
||||
void identity_and_activation(half* inout, half* bias, int m, int b, bool mat_bias){
|
||||
dim3 grid(m, b);
|
||||
if (bias == nullptr)
|
||||
identity_and_activation<N, BLOCKDIM><<<grid , BLOCKDIM>>>(inout);
|
||||
else
|
||||
identity_and_activation<N, BLOCKDIM><<<grid , BLOCKDIM>>>(inout, bias, mat_bias);
|
||||
}
|
||||
94
examples/44_multi_gemm_ir_and_codegen/utils.h
Normal file
94
examples/44_multi_gemm_ir_and_codegen/utils.h
Normal file
@@ -0,0 +1,94 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
#pragma once
|
||||
#define TI(tag) \
|
||||
cudaEvent_t _event_start_ ##tag; \
|
||||
cudaEvent_t _event_end_ ##tag; \
|
||||
float _event_time_ ##tag; \
|
||||
cudaEventCreate(& _event_start_ ##tag); \
|
||||
cudaEventCreate(& _event_end_ ##tag); \
|
||||
cudaEventRecord(_event_start_ ##tag);
|
||||
|
||||
#define TO(tag, str, times) \
|
||||
cudaEventRecord(_event_end_ ##tag); \
|
||||
cudaEventSynchronize(_event_end_ ##tag); \
|
||||
cudaEventElapsedTime(&_event_time_ ##tag, _event_start_ ##tag, _event_end_ ##tag); \
|
||||
float _event_time_once_ ##tag = _event_time_ ##tag / times; \
|
||||
printf("%20s:\t %10.3fus\t", str, _event_time_once_ ##tag * 1000); \
|
||||
cudaDeviceSynchronize(); \
|
||||
printf("%20s string: %s\n",str, cudaGetErrorString(cudaGetLastError()));
|
||||
|
||||
template<typename T>
|
||||
struct memory_unit{
|
||||
T* host_ptr;
|
||||
T* device_ptr;
|
||||
int size_bytes;
|
||||
int elements;
|
||||
void h2d(){
|
||||
cudaMemcpy(device_ptr, host_ptr, size_bytes, cudaMemcpyHostToDevice);
|
||||
}
|
||||
void d2h(){
|
||||
cudaMemcpy(host_ptr, device_ptr, size_bytes, cudaMemcpyDeviceToHost);
|
||||
}
|
||||
void free_all(){
|
||||
free(host_ptr);
|
||||
cudaFree(device_ptr);
|
||||
}
|
||||
memory_unit(int elements_): size_bytes(elements_ * sizeof(T)), elements(elements_){
|
||||
host_ptr = (T*) malloc(elements_ * sizeof(T));
|
||||
cudaMalloc((void**)&device_ptr, elements_ * sizeof(T));
|
||||
}
|
||||
void init(int abs_range = 1){
|
||||
for(int i = 0; i < elements; i++){
|
||||
host_ptr[i] = T(rand() % 100 / float(100) * 2 * abs_range - abs_range);
|
||||
}
|
||||
h2d();
|
||||
}
|
||||
};
|
||||
|
||||
template<typename T>
|
||||
int check_result(T * a, T * b, int N){
|
||||
int cnt = 0;
|
||||
for(int i = 0; i < N; i ++){
|
||||
float std = float(a[i]);
|
||||
float my = float(b[i]);
|
||||
|
||||
if(abs(std - my) / abs(std) > 1e-2)
|
||||
{
|
||||
// printf("my: %f , std: %f\n", my, std);
|
||||
cnt++;
|
||||
}
|
||||
|
||||
}
|
||||
printf("total err: %d / %d\n", cnt, N);
|
||||
return cnt;
|
||||
}
|
||||
Reference in New Issue
Block a user