183 lines
6.1 KiB
C++
183 lines
6.1 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2017 - 2023 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 binding CUTLASS C++ APIs to Python
|
|
*/
|
|
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl_bind.h>
|
|
|
|
#include "builtin_types.h"
|
|
#include "device_launch_parameters.h"
|
|
#include "stddef.h"
|
|
#include "cutlass/cutlass.h"
|
|
|
|
#include "include/conv/convolution.h"
|
|
#include "include/gemm/gemm.h"
|
|
#include "include/types.h"
|
|
#include "include/layout/layout.h"
|
|
#include "include/tensor_coord.h"
|
|
#include "include/arch.h"
|
|
#include "include/tensor_ref_view.h"
|
|
#include "include/swizzling.h"
|
|
#include "test/conv/convolution.h"
|
|
#include "test/gemm/gemm.h"
|
|
|
|
|
|
// Data Types
|
|
#include "library.h"
|
|
|
|
// compiler
|
|
#include "compiler.h"
|
|
|
|
|
|
namespace py = pybind11;
|
|
|
|
|
|
PYBIND11_MODULE(cutlass, m) {
|
|
|
|
// module doc
|
|
m.doc() = "cutlass C++ binding";
|
|
|
|
//
|
|
// Bind data type
|
|
//
|
|
bind_cutlass_types(m);
|
|
|
|
//
|
|
// Bind layout
|
|
//
|
|
bind_layout(m);
|
|
|
|
//
|
|
// Bind tensor coord
|
|
//
|
|
bind_tensor_coord(m);
|
|
|
|
//
|
|
// Bind tensor ref
|
|
//
|
|
bind_tensor_refs_and_views(m);
|
|
|
|
//
|
|
// Bind opcode
|
|
//
|
|
bind_opcode(m);
|
|
|
|
//
|
|
// Bind convolution
|
|
//
|
|
py::module_ conv_submodule = m.def_submodule("conv");
|
|
bind_convolution(conv_submodule);
|
|
|
|
//
|
|
// Bind gemm
|
|
//
|
|
py::module_ gemm_submodule = m.def_submodule("gemm");
|
|
bind_gemm(gemm_submodule);
|
|
|
|
//
|
|
// Bind swizzling
|
|
//
|
|
bind_threadblock_swizzle(m);
|
|
|
|
|
|
//
|
|
// Bind test units
|
|
//
|
|
py::module_ test = m.def_submodule("test");
|
|
py::module_ test_conv = test.def_submodule("conv");
|
|
bind_convolution_test(test_conv);
|
|
|
|
py::module_ test_gemm = test.def_submodule("gemm");
|
|
bind_gemm_test(test_gemm);
|
|
|
|
// data types
|
|
py::enum_<cutlass::DataType>(m, "dtype")
|
|
.value("b1", cutlass::DataType::kB1)
|
|
.value("u2", cutlass::DataType::kU2)
|
|
.value("u4", cutlass::DataType::kU4)
|
|
.value("u8", cutlass::DataType::kU8)
|
|
.value("u16", cutlass::DataType::kU16)
|
|
.value("u32", cutlass::DataType::kU32)
|
|
.value("u64", cutlass::DataType::kU64)
|
|
.value("s2", cutlass::DataType::kS2)
|
|
.value("s4", cutlass::DataType::kS4)
|
|
.value("s16", cutlass::DataType::kS16)
|
|
.value("s64", cutlass::DataType::kS64)
|
|
.value("cf16", cutlass::DataType::kCF16)
|
|
.value("cbf16", cutlass::DataType::kCBF16)
|
|
.value("cf32", cutlass::DataType::kCF32)
|
|
.value("ctf32", cutlass::DataType::kCTF32)
|
|
.value("cf64", cutlass::DataType::kCF64)
|
|
.value("cs2", cutlass::DataType::kCS2)
|
|
.value("cs4", cutlass::DataType::kCS4)
|
|
.value("cs8", cutlass::DataType::kCS8)
|
|
.value("cs16", cutlass::DataType::kCS16)
|
|
.value("cs32", cutlass::DataType::kCS32)
|
|
.value("cs64", cutlass::DataType::kCS64)
|
|
.value("cu2", cutlass::DataType::kCU2)
|
|
.value("cu4", cutlass::DataType::kCU4)
|
|
.value("cu8", cutlass::DataType::kCU8)
|
|
.value("cu16", cutlass::DataType::kCU16)
|
|
.value("cu32", cutlass::DataType::kCU32)
|
|
.value("cu64", cutlass::DataType::kCU64)
|
|
.value("invalid", cutlass::DataType::kInvalid);
|
|
|
|
// layout types
|
|
py::enum_<cutlass::LayoutType>(m, "layout")
|
|
.value("ColumnMajorInterleaved2", cutlass::LayoutType::kColumnMajorInterleaved2)
|
|
.value("RowMajorInterleaved2", cutlass::LayoutType::kRowMajorInterleaved2)
|
|
.value("ColumnMajorInterleaved64", cutlass::LayoutType::kColumnMajorInterleaved64)
|
|
.value("RowMajorInterleaved64", cutlass::LayoutType::kRowMajorInterleaved64)
|
|
.value("TensorNDHWC", cutlass::LayoutType::kTensorNDHWC)
|
|
.value("TensorNCHW", cutlass::LayoutType::kTensorNCHW)
|
|
.value("TensorNGHWC", cutlass::LayoutType::kTensorNGHWC)
|
|
.value("TensorNC64HW64", cutlass::LayoutType::kTensorNC64HW64)
|
|
.value("TensorC64RSK64", cutlass::LayoutType::kTensorC64RSK64);
|
|
|
|
// transform types
|
|
py::enum_<cutlass::ComplexTransform>(m, "complex_transform")
|
|
.value("none", cutlass::ComplexTransform::kNone)
|
|
.value("conj", cutlass::ComplexTransform::kConjugate);
|
|
|
|
//
|
|
// Compiler
|
|
//
|
|
py::class_<cutlass::CompileCache>(m, "CompileCache")
|
|
.def(py::init<>())
|
|
.def("at", &cutlass::CompileCache::at)
|
|
.def("insert", &cutlass::CompileCache::insert)
|
|
.def("size", &cutlass::CompileCache::size)
|
|
.def("clear", &cutlass::CompileCache::clear);
|
|
|
|
}
|