51 lines
2.4 KiB
Plaintext
51 lines
2.4 KiB
Plaintext
/*************************************************************************
|
|
* Copyright (C) [2023-2024] by Cambricon, Inc.
|
|
*
|
|
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
|
|
* OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
|
* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
|
|
* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
|
|
* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
|
|
* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
|
|
* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
*************************************************************************/
|
|
#ifndef CSRC_KERNELS_CAST_GATING_MLUH_
|
|
#define CSRC_KERNELS_CAST_GATING_MLUH_
|
|
|
|
#include "../kernel_utils.h"
|
|
#include "cnnl.h"
|
|
namespace tmo {
|
|
/**
|
|
* @brief Convert input to float32 and do gating operation.
|
|
* @param queue: The queue for mlu.
|
|
* @param input: Input. Pointer to the MLU memory that stores the input,
|
|
* the shape must be [input_row, hidden_size].
|
|
* @param filter: Input. Pointer to the MLU memory that stores the weight,
|
|
* the shape must be [expert_num, hidden_size].
|
|
* @param output: Output. Pointer to the MLU memory that stores the output,
|
|
* the shape must be [input_row, expert_num].
|
|
* @param input_row: Input.
|
|
* @param expert_num: Input.
|
|
* @param hidden_size: Input.
|
|
* @param a_dtype: Input. The data-type of input.
|
|
* @param workspace: Input. Pointer to the MLU workspace.
|
|
* @param workspace_size_bytes: Input. The size of workspace in bytes.
|
|
* @note: a_dtype must be CNNL_DTYPE_BFLOAT16 or CNNL_DTYPE_HALF.
|
|
* expert_num must be in range [1, 128].
|
|
* If workspace is NOT NULL, workspace_size_bytes must NOT be smaller than 16 * 1024 * 1024.
|
|
* The data-type of filter and output must be float.
|
|
* cast_gating only supports MLU500 device or higher.
|
|
*/
|
|
KernelStatus invokeCastGating(cnrtQueue_t queue,
|
|
void *input,
|
|
void *filter,
|
|
void *output,
|
|
int input_row,
|
|
int expert_num,
|
|
int hidden_size,
|
|
cnnlDataType_t a_dtype,
|
|
void *workspace,
|
|
size_t workspace_size_bytes);
|
|
} // namespace tmo
|
|
#endif // CSRC_KERNELS_CAST_GATING_MLUH_
|