Visão Geral do Operaodr aclnnMatmulAllReduceV3

O operador aclnnMatmulAllReduceV3 é projetado para fundir a operação de multiplicação de matrizes (MatMul) com a comunicação coletiva AllReduce. Esta fusão é compatível com as funcionalidades da versão anterior (aclnnMatmulAllReduceV2), mas introduz o parâmetro commMode, permitindo a seleção dinâmica do motor de comunicação.

A fórmula matemática executada por este operador é:

output = AllReduce(x1 @ x2 + bias + x3)

Compatibilidade de Hardware

Nota: É obrigatório que os pacotes de firmware do driver e do CANN sejam da versão 8.0.RC2 ou superior e compatíveis entre si. Versões incompatíveis podem resultar em falhas críticas, como Bus Error.

Arquitetura da API

A interface segue um modelo de duas etapas. A primeira etapa calcula o tamanho da memória de workspace necessária e gera o executor do operador. A segunda etapa executa a computação real.

Etapa 1: Obtenção do Workspace e Executor

aclnnStatus aclnnMatmulAllReduceV3GetWorkspaceSize(
    const aclTensor *input_left,
    const aclTensor *input_right,
    const aclTensor *bias_tensor,
    const aclTensor *residual_tensor,
    const char      *comm_group,
    const char      *reduction_type,
    const char      *engine_mode,
    int64_t          comm_turns,
    int64_t          stream_config,
    const aclTensor *result_tensor,
    uint64_t        *required_workspace,
    aclOpExecutor  **op_executor)

Etapa 2: Execução do Operador

aclnnStatus aclnnMatmulAllReduceV3(
    void              *workspace_ptr,
    uint64_t           workspace_bytes,
    aclOpExecutor     *op_executor,
    const aclrtStream  execution_stream)

Detalhamento dos Parâmetros (Etapa 1)

• input_left (x1): Matriz esquerda. Suporta 2D ou 3D. Formatos ND. Tipos: FLOAT16, BFLOAT16.
• input_right (x2): Matriz direita. Apenas 2D. Suporta transposição na ND. Tipos: FLOAT16, BFLOAT16.
• residual_tensor (x3): Tansor residual adicionado após a multiplicação. Shape deve corresponder ao resultado do MatMul. Suporta tensores não contíguos.
• bias_tensor: Vetor de bias (1D). Suporta tensores não contíguos.
• comm_group: Nome do domínio de comunicação (obtido via HcclGetCommName).
• reduction_type: Tipo de redução. Atualmente, apenas "sum" é suportado.
• engine_mode: Motor de comunicação ("ai_cpu", "ccu" ou vazio para adaptativo).
• comm_turns: Fator de divisão de dados. Atualmente, apenas 0 é suportado.
• stream_config: Modo de fluxo AscendCL. Atualmente, apenas 1 é suportado.
• result_tensor: Tensor de saída. Dimensões iguais a input_left.

Restrições e Diretrizes de Uso

• Motores de Comunicação:
  • Atlas A2: Não suporta seleção manual. engine_mode deve ser vazio (usa AICPU por padrão).
  • Ascend 950PR/DT: Suporta "ai_cpu", "ccu" ou vazio (adaptativo, padrão CCU). O mesmo motor deve ser usado em toda a cadeia de comunicação.
• Computação Determinística:
  • Atlas A2: Não determinístico por padrão. Ative via variável de ambiente HCCL_DETERMINISTIC=true.
  • Ascend 950PR/DT: Determinístico por padrão.
• Topologia e Shape:
  • Suporta apenas topologia all mesh via link HCCS.
  • Limites de dimensão: b*s, m, k, n não podem exceder 2147483647.
  • Se input_left for (b, s, k), a saída será (b, s, n). Se for (m, k), a saída será (m, n).
  • Tensores vazios são suportados (ex: k=0 resulta em bias + x3; outras dimensões zeradas retornam imediatamente sem executar o kernel).

Exemplo de Implementação em C++

O exemplo a seguir demonstra a inicialização do HCCL, criação de tensores e a execução do operador fundido em um ambiente multi-thread.

#include <iostream>
#include <vector>
#include <thread>
#include <memory>
#include "hccl/hccl.h"
#include "acl/acl.h"

#define NUM_DEVICES 2
#define ENSURE_SUCCESS(call, msg) \
    do { \
        auto status = (call); \
        if (status != ACL_SUCCESS) { \
            std::cerr << msg << " Error code: " << status << std::endl; \
            return status; \
        } \
    } while(0)

// Função auxiliar para calcular o tamanho total do shape
size_t CalculateTotalElements(const std::vector<int64_t>& dims) {
    size_t total = 1;
    for (auto d : dims) total *= d;
    return total;
}

// Estrutura para encapsular os argumentos de cada thread
struct ThreadContext {
    int32_t device_id;
    HcclComm comm_handle;
    aclrtStream compute_stream;
    aclrtContext dev_context;
};

// Função para alocar e criar um tensor ACL
aclTensor* BuildAclTensor(const std::vector<int64_t>& shape, aclDataType dtype, void* dev_ptr) {
    std::vector<int64_t> strides(shape.size(), 1);
    for (int i = shape.size() - 2; i >= 0; --i) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }
    return aclCreateTensor(shape.data(), shape.size(), dtype, strides.data(), 0, 
                           ACL_FORMAT_ND, shape.data(), shape.size(), dev_ptr);
}

void ExecuteFusedOperation(ThreadContext& ctx) {
    aclrtSetCurrentContext(ctx.dev_context);
    
    char comm_name[128];
    ENSURE_SUCCESS(HcclGetCommName(ctx.comm_handle, comm_name), "Failed to get HCCL comm name");

    // Definição dos shapes (exemplo: 16x32 @ 32x64)
    std::vector<int64_t> dim_x1 = {16, 32};
    std::vector<int64_t> dim_x2 = {32, 64};
    std::vector<int64_t> dim_bias = {64};
    std::vector<int64_t> dim_x3 = {16, 64};
    
    size_t size_x1 = CalculateTotalElements(dim_x1) * sizeof(aclFloat16);
    size_t size_x2 = CalculateTotalElements(dim_x2) * sizeof(aclFloat16);
    size_t size_bias = CalculateTotalElements(dim_bias) * sizeof(aclFloat16);
    size_t size_x3 = CalculateTotalElements(dim_x3) * sizeof(aclFloat16);

    void *ptr_x1, *ptr_x2, *ptr_bias, *ptr_x3, *ptr_out;
    
    // Alocação de memória no device
    aclrtMalloc(&ptr_x1, size_x1, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&ptr_x2, size_x2, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&ptr_bias, size_bias, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&ptr_x3, size_x3, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&ptr_out, size_x3, ACL_MEM_MALLOC_HUGE_FIRST);

    // Criação dos tensores
    aclTensor *t_x1 = BuildAclTensor(dim_x1, ACL_FLOAT16, ptr_x1);
    aclTensor *t_x2 = BuildAclTensor(dim_x2, ACL_FLOAT16, ptr_x2);
    aclTensor *t_bias = BuildAclTensor(dim_bias, ACL_FLOAT16, ptr_bias);
    aclTensor *t_x3 = BuildAclTensor(dim_x3, ACL_FLOAT16, ptr_x3);
    aclTensor *t_out = BuildAclTensor(dim_x3, ACL_FLOAT16, ptr_out);

    uint64_t workspace_size = 0;
    aclOpExecutor* executor = nullptr;
    void* workspace_ptr = nullptr;

    // Etapa 1: Obter workspace e executor
    aclnnStatus ret = aclnnMatmulAllReduceV3GetWorkspaceSize(
        t_x1, t_x2, t_bias, t_x3, comm_name, "sum", "", 0, 1, t_out, &workspace_size, &executor
    );
    
    if (ret == ACL_SUCCESS && workspace_size > 0) {
        aclrtMalloc(&workspace_ptr, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST);
    }

    // Etapa 2: Executar o operador
    if (ret == ACL_SUCCESS) {
        ret = aclnnMatmulAllReduceV3(workspace_ptr, workspace_size, executor, ctx.compute_stream);
    }

    // Sincronização
    aclrtSynchronizeStream(ctx.compute_stream);

    // Limpeza de recursos
    aclDestroyTensor(t_x1); aclDestroyTensor(t_x2); aclDestroyTensor(t_bias);
    aclDestroyTensor(t_x3); aclDestroyTensor(t_out);
    
    aclrtFree(ptr_x1); aclrtFree(ptr_x2); aclrtFree(ptr_bias);
    aclrtFree(ptr_x3); aclrtFree(ptr_out);
    if (workspace_ptr) aclrtFree(workspace_ptr);
    
    aclrtDestroyStream(ctx.compute_stream);
    HcclCommDestroy(ctx.comm_handle);
    aclrtDestroyContext(ctx.dev_context);
    aclrtResetDevice(ctx.device_id);
}

int main() {
    aclInit(nullptr);
    
    int32_t device_list[NUM_DEVICES] = {0, 1};
    HcclComm comms[NUM_DEVICES];
    
    for(int i=0; i<NUM_DEVICES; ++i) aclrtSetDevice(device_list[i]);
    HcclCommInitAll(NUM_DEVICES, device_list, comms);
    
    std::vector<std::thread> workers;
    ThreadContext contexts[NUM_DEVICES];
    
    for (int i = 0; i < NUM_DEVICES; ++i) {
        aclrtSetDevice(device_list[i]);
        aclrtCreateContext(&contexts[i].dev_context, device_list[i]);
        aclrtCreateStream(&contexts[i].compute_stream);
        contexts[i].device_id = device_list[i];
        contexts[i].comm_handle = comms[i];
        
        workers.emplace_back(ExecuteFusedOperation, std::ref(contexts[i]));
    }
    
    for (auto& t : workers) t.join();
    
    aclFinalize();
    return 0;
}