本文共 3746 字，大约阅读时间需要 12 分钟。

本文主体内容转载自，因为我的C语言不熟悉，所以本篇大体还是借鉴于罗神的文章。不过相比罗神，我比较侥幸成为了PyTorch官方的contributor，一共也没改多少代码…

多重派发

首先我们需要了解什么是多重派发，直观来说就是根据一些条件（函数签名）在运行时分发相关的函数。PyTorch中使用了这个方法，就使得CPU上可以根据环境变量，硬件指令实现等条件选择SIMD加速的backend。首先我们用一个C的demo来展示如何在C语言中模拟多重派发。

目标

我们这个demo的目标是根据C语言的内置类型对不同函数进行分发，具体来说就是现在有（以下是伪代码）

func_1(int32, int64)func_2(int32, float32)...

他们都是根据不同类型实现的类似功能的函数，我们希望在调用一个入口函数func的时候自动根据输入变量的类型进行方法的派发。

首先我们需要创建一个枚举类型来标记一下

typedef enum {    tagInt32 = 0,    tagInt64,    tagFloat32,    tagFloat64,    ntypes, // 这个在后面的函数调用列表中有意义} Tag;

这里，tagInt64 = 1; tagFloat32 = 2; tagFloat64 = 3; ntypes = 4;

然后定义一个结构体作为类型

typedef struct{    void *data;    size_t size;    Tag tag; // tag的类型用int也ok     //int tag; 与Tag tag效果一样。} Type;

之后我们需要对不同的类型定义一些工厂函数来产生相对应的实例，所以这里定义一个宏来批量产生函数定义

#define MAKE_TYPE(NAME, CTYPE)                      \    Type *make_##NAME(CTYPE data)                   \    {                                               \        Type *type;                                 \        type = (Type *)malloc(sizeof(Type));        \        type->size = sizeof(CTYPE);                 \        type->data = (void *)malloc(sizeof(CTYPE)); \        type->tag = tag##NAME;                      \        CTYPE *temp_ptr = (CTYPE *)type->data;      \        *temp_ptr = data;                           \        return type;                                \    }MAKE_TYPE(Int32, int32_t)MAKE_TYPE(Int64, int64_t)MAKE_TYPE(Float32, float)MAKE_TYPE(Float64, double)

这一步是生成make_Int32, make_Int64, make_Float32, make_Float64这4个宏函数，方式类似于Python中的装饰器(decorator)。##在宏定义中的连接的意思：比如define torch##7就是torch7。

然后定义函数类型，我们这里的函数都是输入参数为2（可变参数的话，看情况换成指针的指针之类的方案）

typedef int (*FuncType)(Type *a, Type *b);

关于C/C++函数指针的内容，请参考

然后我们有一些函数在不同类型上的实现

// int32 float32int Func_0x001(Type *a, Type *b){    int32_t *a_data = a->data;    float *b_data = b->data;    printf("input: int32 %d, float32 %f", *a_data, *b_data);    return 0;}// int32 float64int Func_0x002(Type *a, Type *b){    int32_t *a_data = a->data;    double *b_data = b->data;    printf("input: int32 %d, float64 %lf", *a_data, *b_data);    return 0;}

然后别忘了让没有相关实现的类型fall back到错误处理或者默认方法上去（PyTorch中是回退到一般的没有SIMD指令的实现上去）

int fallback(){    printf("Error: MethodError: No method match input type\n");    exit(-1);    return 0;}

接下来定义一个函数调用表，表的各个维度就是表示各个类型

FuncType FUNC_CALL_LIST[ntypes][ntypes] = {    {        NULL,        // int32 int32        NULL,        // int32 int64        &Func_0x001, // int32 float32        &Func_0x002, // int32 float64    },    {        NULL,        // int64 int32        NULL,        // int64 int64        NULL,        // int64 float32        NULL,        // int64 float64    },    {        NULL,        // float32 int32        NULL,        // float32 int64        NULL,        // float64 int32        NULL,        // float64 int64    },};

然后定义函数入口

int FuncEntry(Type *a, Type *b){    FuncType func_ptr = FUNC_CALL_LIST[a->tag][b->tag];    if (func_ptr != NULL)    {        (*func_ptr)(a, b);    }    else     {        return fallback();    }}

最后提供用来释放（析构）的函数

void type_free(Type *type){    free(type->data);    free(type);}

int main(){    Type *a = make_Int32(2); // a -> tag = 0    Type *b = make_Float32(1.5); // b -> tag = 2    // 调用的是函数Func_0x001    FuncEntry(a, b);    type_free(a);    type_free(b);    printf("\n");    return 0;}

编译一下gcc main.c -o main，就可以看到效果，程序会在运行时分发对应的方法，然后输出结果，如果没有对应的方法就回退到默认实现上去。

然后需要头文件

#include 
   
    #include 
    
     #include 
     
    
   

PyTorch中的TH库的实现机制是类似的，但是没有类型判断而是根据支持的SIMD指令和环境变量来分发方法。

类似在demo里使用的，每个SIMD实现都是

• name：函数名
• EXT：对应的SIMD类型，有：DEFAULT，AVX，AVX2，SSE等

相关的信息会记录在 FunctionDescription这个结构体中，然后最后根据对应的条件进行派发。如果有新的实现只需要插入到generic/THVectorDispatch.cpp中即可，不需要管其它部分，当硬件和环境变量的相关条件满足的时候会自动分配过去，这是运行时分配的，因为是直接访问地址，所以复杂度也是O(1)，不需要进行重新编译。