wehub-resource-sync
140 lines
4.4 KiB
Plaintext
140 lines
4.4 KiB
Plaintext
//
|
|
// MatMulExecution.cu
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2026/02/25.
|
|
// Copyright © 2026, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "core/MusaBackend.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "MNN_generated.h"
|
|
#include <musa_runtime.h>
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
// MUSA kernel for matrix multiplication
|
|
__global__ void MatMulKernel(const float* A, const float* B, float* C,
|
|
int M, int N, int K) {
|
|
int row = blockIdx.y * blockDim.y + threadIdx.y;
|
|
int col = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (row >= M || col >= N) return;
|
|
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < K; ++i) {
|
|
sum += A[row * K + i] * B[i * N + col];
|
|
}
|
|
|
|
C[row * N + col] = sum;
|
|
}
|
|
|
|
// MUSA kernel for batched matrix multiplication
|
|
__global__ void BatchMatMulKernel(const float* A, const float* B, float* C,
|
|
int batch, int M, int N, int K) {
|
|
int b = blockIdx.z;
|
|
int row = blockIdx.y * blockDim.y + threadIdx.y;
|
|
int col = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (row >= M || col >= N) return;
|
|
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < K; ++i) {
|
|
sum += A[b * M * K + row * K + i] * B[b * K * N + i * N + col];
|
|
}
|
|
|
|
C[b * M * N + row * N + col] = sum;
|
|
}
|
|
|
|
class MatMulExecution : public Execution {
|
|
public:
|
|
MatMulExecution(Backend* backend) : Execution(backend) {}
|
|
|
|
virtual ErrorCode onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) override {
|
|
mShapeChanged = true;
|
|
return NO_ERROR;
|
|
}
|
|
|
|
virtual ErrorCode onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) override {
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("start MatMulExecution onExecute...\n");
|
|
#endif
|
|
|
|
auto input0 = inputs[0];
|
|
auto input1 = inputs[1];
|
|
auto output = outputs[0];
|
|
|
|
auto input0Shape = input0->shape();
|
|
auto input1Shape = input1->shape();
|
|
auto outputShape = output->shape();
|
|
|
|
void* input0Ptr = (void*)input0->deviceId();
|
|
void* input1Ptr = (void*)input1->deviceId();
|
|
void* outputPtr = (void*)output->deviceId();
|
|
|
|
if (input0Shape.size() == 2 && input1Shape.size() == 2) {
|
|
// 2D matrix multiplication
|
|
int M = input0Shape[0];
|
|
int K = input0Shape[1];
|
|
int N = input1Shape[1];
|
|
|
|
dim3 threadsPerBlock(16, 16);
|
|
dim3 blocksPerGrid((N + 15) / 16, (M + 15) / 16);
|
|
|
|
MatMulKernel<<<blocksPerGrid, threadsPerBlock>>>(
|
|
(const float*)input0Ptr, (const float*)input1Ptr, (float*)outputPtr, M, N, K);
|
|
} else {
|
|
// Batched matrix multiplication
|
|
int batch = 1;
|
|
int M = input0Shape[input0Shape.size() - 2];
|
|
int K = input0Shape[input0Shape.size() - 1];
|
|
int N = input1Shape[input1Shape.size() - 1];
|
|
|
|
for (size_t i = 0; i < input0Shape.size() - 2; ++i) {
|
|
batch *= input0Shape[i];
|
|
}
|
|
|
|
dim3 threadsPerBlock(16, 16);
|
|
dim3 blocksPerGrid((N + 15) / 16, (M + 15) / 16, batch);
|
|
|
|
BatchMatMulKernel<<<blocksPerGrid, threadsPerBlock>>>(
|
|
(const float*)input0Ptr, (const float*)input1Ptr, (float*)outputPtr, batch, M, N, K);
|
|
}
|
|
|
|
// Check for kernel launch errors
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
MNN_ERROR("MUSA MatMul kernel launch failed: %s\n", musaGetErrorString(err));
|
|
}
|
|
|
|
// Synchronize to ensure completion
|
|
auto musaBackend = static_cast<MusaBackend*>(backend());
|
|
musaBackend->getMusaRuntime()->device_sync();
|
|
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("end MatMulExecution onExecute...\n");
|
|
#endif
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
private:
|
|
bool mShapeChanged{false};
|
|
};
|
|
|
|
// Creator for MatMul operations
|
|
class MatMulCreator : public MusaBackend::Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const MNN::Op* op, Backend* backend) const override {
|
|
return new MatMulExecution(backend);
|
|
}
|
|
};
|
|
|
|
MusaCreatorRegister<MatMulCreator> __MatMulExecution(OpType_MatMul);
|
|
MusaCreatorRegister<MatMulCreator> __MatMulInt8Execution(OpType_MatMulInt8);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|