深入解析华为CANN算子开发：从入图到动态Shape推导

随着AI计算的不断发展，华为昇腾AI处理器及其CANN算子开发框架在高性能算子实现和图优化中扮演着核心角色。本文将深入解析CANN算子开发的全流程，重点讲解入图阶段的Shape与DataType推导、数据依赖处理及动态输出Shape推导，为开发者提供实践参考。

一、算子开发与入图的概述

在传统算子开发中，开发者通常关注算子核心计算逻辑，即通过输入Tensor计算输出Tensor。然而，在图模式（Graph Mode）下，算子开发需要考虑更复杂的场景：

1. Tensor Shape与DataType推导
   在图生成阶段，框架会提前推导每个Tensor的形状与数据类型。这样可以：

   • 在执行前验证输入输出合法性；
   • 为输出Tensor静态分配内存，避免动态分配的性能开销。

2. 算子入图代码文件
   除了算子核函数，开发者还需提供入图逻辑，包括：

   • DataType推导函数
   • Shape推导函数
   • ShapeRange推导函数（动态输出Shape场景）
   • 数据依赖声明

二、数据类型推导（DataType Inference）

DataType推导是算子入图的第一步。以一个自定义AddCustom算子为例，输出Tensor的DataType可以直接继承输入Tensor的类型：

namespace ge {
static graphStatus InferDataType(gert::InferDataTypeContext* context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
} // namespace ge

对于更复杂的场景，比如输入为DT_INT4时输出为DT_INT32，推导逻辑可灵活处理：

if (context->GetInputDataType(0) == DT_INT4) {
    context->SetOutputDataType(0, DT_INT32);
}

这种设计确保了算子在图模式下可以自动处理多种数据类型输入，保持计算正确性。

三、Shape推导（Shape Inference）

Shape推导是图模式下的核心环节。开发者可以通过两种方式实现：

3.1 Follow模式

若输出Shape与某输入Shape完全一致，可使用Follow接口快速表达：

this->Output("y1")
    .ParamType(REQUIRED)
    .Follow("x1", FollowType::SHAPE);

3.2 自定义InferShape函数

对于输出Shape与输入Shape存在复杂关系的算子，如Reshape，需编写自定义InferShape函数：

ge::graphStatus InferShapeForReshape(InferShapeContext *context) {
    const gert::Shape *x_shape = context->GetInputShape(0);
    const gert::Tensor *shape_tensor = context->GetInputTensor(1);
    gert::Shape *output_shape = context->GetOutputShape(0);

    if (!x_shape || !shape_tensor || !output_shape) return ge::GRAPH_FAILED;

    auto reshape_size = static_cast<int32_t>(shape_tensor->GetShapeSize());
    if (shape_tensor->GetDataType() == ge::DT_INT32) {
        int32_t *reshape_data = shape_tensor->GetData<int32_t>();
        return ReshapeInferShapeImpl<int32_t>(reshape_data, *x_shape, *output_shape, reshape_size);
    } else {
        int64_t *reshape_data = shape_tensor->GetData<int64_t>();
        return ReshapeInferShapeImpl<int64_t>(reshape_data, *x_shape, *output_shape, reshape_size);
    }
}

3.3 数据依赖算子

部分算子在Shape推导时，需要依赖输入的真实值，如Reshape依赖shape输入。此类输入需通过ValueDepend(REQUIRED)声明：

this->Input("shape")
    .ParamType(REQUIRED)
    .ValueDepend(REQUIRED);

四、动态Shape与ShapeRange推导

有些算子（如Unique）的输出Shape在编译阶段无法确定，必须在执行时才能得出。这时需要ShapeRange推导，用于预估最大输出内存：

ge::graphStatus UniqueInferShapeRangeFunc(gert::InferShapeRangeContext *context) {
    auto x_shape_range = context->GetInputShapeRange(0U);
    auto y_shape_range = context->GetOutputShapeRange(0U);

    y_shape_range->GetMax()->SetDim(0, x_shape_range->GetMax()->GetDim(0));
    y_shape_range->GetMin()->SetDim(0, x_shape_range->GetMin()->GetDim(0));
    return ge::GRAPH_SUCCESS;
}

通过ShapeRange推导，框架可以安全地为动态输出分配内存，保证算子执行的正确性。

五、获取算子属性与特殊输入

5.1 获取IR属性

算子在注册时定义的属性（如src_format、dst_format）可通过TilingContext访问：

const RuntimeAttrs *attrs = context->GetAttrs();
const char *src_format = attrs->GetAttrPointer<char>(0);
const char *dst_format = attrs->GetAttrPointer<char>(1);

5.2 Optional与Dynamic输入

某些算子包含多个可选或动态输入（如DynamicRNNV3），其实例化后的输入Index可能不固定。此时可使用：

auto shape = context->GetOptionalInputShape(original_index);
auto dyn_shape = context->GetDynamicInputShape(ir_index, relative_index);

保证在InferShape和Tiling函数中正确获取输入Shape信息。

六、总结

华为CANN算子开发不仅仅是核函数实现，更包含入图阶段的静态分析、Shape/DataType推导及数据依赖处理。核心要点包括：

1. 静态推导优先：尽量在编译阶段推导输出Shape和DataType，减少运行时开销。
2. 数据依赖算子：通过ValueDepend声明，保证在InferShape阶段可访问输入Tensor数据。
3. 动态Shape处理：ShapeRange推导提供安全内存分配策略，支持如Unique等动态输出算子。
4. 属性和特殊输入处理：Optional和Dynamic输入需通过对应接口获取，以保证推导正确性。

掌握这些开发技巧，可以帮助开发者在CANN框架下高效构建高性能算子，同时保证算子在图模式执行中的正确性与稳定性。