PyTorch学习笔记(7)——多重派发

本文主体内容转载自罗秀哲——PyTorch源码浅析（五），因为我的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++函数指针的内容，请参考typedef定义函数类型的用法

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

// int32 float32
int 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 float64
int 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 <stdio.h>
#include <stdlib.h>
#include <stdint.h>

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

类似在demo里使用的，每个SIMD实现都是
THVector_(name_EXT)的形式，其中

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

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