普通视图

发现新文章,点击刷新页面。
昨天以前SatanWoo

TrampolineHook - 解决栈污染问题支持变参 Hook

作者 SatanWoo
2020年5月18日 01:27

在之前的文章《基于桥的全量方法 Hook 方案(3)- TrampolineHook》 的文末,我说如果对汇编熟悉的同学可能会发现我之前实现的一个错误 - 关于上下文污染

一提到到上下文污染,可能我们绝大多数人想到的都是寄存器污染,但是实际上还有一个不容我们忽视的上下文资源:,过去可能大家常见的 Hook 代码关注比较少,正好这次在借助 TrampolineHook 修复这个方面的问题,我们来一起探讨下。

先看一个例子

假设有这样一个类 TestObject和 不定参函数 method,定义如下:

@interface TestObject : NSObject
@end

@implementation TestObject

- (void)method:(int *)value,...
{
     va_list list;
    va_start(list, value);

    while (value) {
        NSLog(@"orig value is %d", *value);
        value = va_arg(list, int *);
    }
    va_end(list);
}

@end

如果要使用 TrampolineHook 来拦截 method 的调用,也非常简单。如下所示:

THInterceptor *sharedInterceptor = [THInterceptor sharedInterceptorWithFunction:(IMP)wzq_check_variadic];

Method m = class_getInstanceMethod([TestObject class], @selector(method:));
IMP imp = method_getImplementation(m);

THInterceptorResult *result = [sharedInterceptor interceptFunction:(IMP)imp];
if (result.state == THInterceptStateSuccess) {
    method_setImplementation(m, (IMP)result.replacedAddress);
}

// 拦截函数
void wzq_check_variadic(id a, char * methodName, int *v, ...)
{
        NSLog(@"haha checked %@ %s", a, methodName);
}

当我们使用如下方式调用 -[TestObject method:] 的时候,你会发现一切正常,毫无问题。

TestObject *obj = [[TestObject alloc] init];
int a = 0;
int b = 1;
int c = 2;
int d = 3;
int e = 4;
int f = 5;
int g = 6;
int h = 7;
int i = 8;
[obj method:&a, &b, &c, &d, &e, &f, &g, &h, &i, nil];

但是如果你将拦截函数中添加打印参数的语句后,如下所示:

void wzq_check_variadic(id a, char * methodName, int *v, ...)
{
    NSLog(@"haha checked %@ %s", a, methodName);
    va_list args;
    va_start(args, v);
    while (v != NULL) { 
        NSLog(@"v is %d", *v);
        v = va_arg(args, int *); // crash 
    }
    va_end(args);
}

你会发现出现了必现的崩溃情形,而且是必定崩溃在第二次读取变参列表中的参数的时候。
为什么添加了读取参数的代码就导致运行崩溃了?有点意思。

了解变参的传递过程。

为了避免优化的干扰,如下汇编生成的优化选项为 -O0

为了看运行时的栈结构是如何生成的,我们通过汇编结合图的形式来一探究竟没 Hook 的时候的调用情况。

首先先看 Caller 函数,即 [obj method:&a, &b, &c, &d, &e, &f, &g, &h, &i, nil]; 这段代码所处的函数,汇编如下:

 // prologue
0x1009dc428 <+0>:   sub    sp, sp, #0xb0             ; =0xb0 
0x1009dc42c <+4>:   stp    x29, x30, [sp, #0xa0]
0x1009dc430 <+8>:   add    x29, sp, #0xa0            ; =0xa0 

// 构造 int 变量 0 - 8
0x1009dc498 <+112>: stur   wzr, [x29, #-0x2c]
0x1009dc49c <+116>: mov    w11, #0x1
0x1009dc4a0 <+120>: stur   w11, [x29, #-0x30]
0x1009dc4a4 <+124>: mov    w11, #0x2
0x1009dc4a8 <+128>: stur   w11, [x29, #-0x34]
0x1009dc4ac <+132>: mov    w11, #0x3
0x1009dc4b0 <+136>: stur   w11, [x29, #-0x38]
0x1009dc4b4 <+140>: mov    w11, #0x4
0x1009dc4b8 <+144>: stur   w11, [x29, #-0x3c]
0x1009dc4bc <+148>: mov    w11, #0x5
0x1009dc4c0 <+152>: stur   w11, [x29, #-0x40]
0x1009dc4c4 <+156>: mov    w11, #0x6
0x1009dc4c8 <+160>: stur   w11, [x29, #-0x44]
0x1009dc4cc <+164>: mov    w11, #0x7
0x1009dc4d0 <+168>: stur   w11, [x29, #-0x48]
0x1009dc4d4 <+172>: mov    w11, #0x8
0x1009dc4d8 <+176>: stur   w11, [x29, #-0x4c]
0x1009dc4dc <+180>: ldur   x9, [x29, #-0x28]
0x1009dc4e0 <+184>: ldr    x1, [x8]
0x1009dc4e4 <+188>: mov    x8, sp
0x1009dc4e8 <+192>: mov    x10, #0x0

// 把对应 int 变量的地址存入栈中 
0x1009dc4ec <+196>: str    x10, [x8, #0x40]
0x1009dc4f0 <+200>: sub    x10, x29, #0x4c           ; =0x4c 
0x1009dc4f4 <+204>: str    x10, [x8, #0x38]
0x1009dc4f8 <+208>: sub    x10, x29, #0x48           ; =0x48 
0x1009dc4fc <+212>: str    x10, [x8, #0x30]
0x1009dc500 <+216>: sub    x10, x29, #0x44           ; =0x44 
0x1009dc504 <+220>: str    x10, [x8, #0x28]
0x1009dc508 <+224>: sub    x10, x29, #0x40           ; =0x40 
0x1009dc50c <+228>: str    x10, [x8, #0x20]
0x1009dc510 <+232>: sub    x10, x29, #0x3c           ; =0x3c 
0x1009dc514 <+236>: str    x10, [x8, #0x18]
0x1009dc518 <+240>: sub    x10, x29, #0x38           ; =0x38 
0x1009dc51c <+244>: str    x10, [x8, #0x10]
0x1009dc520 <+248>: sub    x10, x29, #0x34           ; =0x34 
0x1009dc524 <+252>: str    x10, [x8, #0x8]
0x1009dc528 <+256>: sub    x10, x29, #0x30           ; =0x30 
0x1009dc52c <+260>: str    x10, [x8]

 // 其余参数 x0, x1, x2
0x1009dc530 <+264>: sub    x2, x29, #0x2c            ; =0x2c 
0x1009dc534 <+268>: mov    x0, x9

// 调用 method 函数
0x1009dc538 <+272>: bl     0x1009e8d9c               ; symbol stub for: objc_msgSend

上述这段函数,简要而言,就是干了四件事:

  • 分配 176 byte 的栈内存

  • 在栈上分配 a = 0, b = 1 等等 9 个变量

  • 把 &b, &c 等 8个 int 变量的地址压栈。

  • x0 (obj), x1 (method), x2(&a)

特别注意,变参列表的第一个参数也是通过寄存器来传递。

  • 调用 method 函数

如果不理解, 可以参考这张图:

而当进入 method: 函数时,汇编如下:

重点看两行蓝色汇编断点的地方,其实是在暗示一种循环,也从底层实现上对应上了我们不断循环获取变参列表的逻辑。

简要来说,就是从变参列表的第一参数(寄存器中的值代表地址),开始读取,循环遍历。这里的循环利用了栈空间在函数调用间的连续性,不断将偏移地址从原来 caller 函数的 sp 回溯,读取处于高地址的 caller 栈空间中的 int 变量地址。

看到这,我想大家也知道了为什么是必定崩溃在第二次读取变参的时候。

  • x0, x1 不用说,是寄存器参数。和变参不变参函数无关,这也能解释为什么只读取 id obj 和 SEL selector 不会崩溃。

  • x2,即变参函数列表的第一个参数,我这里把他称为变参的锚点参数,它也是通过寄存器传递,所以读取的时候没问题。

  • 变参列表的后续参数都是分配在调用函数(caller)中,而 TrampolineHook 在调用 interceptor 之前利用了栈(操作 SP)来保存上下文,如下所示,因此破坏了栈资源上下文,导致循环从栈地址获取参数的时候崩溃

    stp q0,  q1,   [sp, #-32]!
    stp q2,  q3,   [sp, #-32]!
    stp q4,  q5,   [sp, #-32]!
    stp q6,  q7,   [sp, #-32]!
    stp lr,  x10,  [sp, #-16]!
    stp x0,  x1,   [sp, #-16]!
    
  • 而调用原函数的时候,由于栈已经复原了,所以就不会出现崩溃了。

解决方案

了解了问题出现的原因,解决办法就很简单了,我们要让调用 inteceptor 时候的上下文和调用原函数一样。

  • 还是构造一堆的动态 trampoline ,让原函数替换到 trampoline,同时保存原函数的 IMP。

  • 依然保存原先需要的上下文,比如通用寄存器、浮点寄存器,但是不能使用栈了。

  • 调用 interceptor。

  • 恢复上下文,调用到原函数。

其实整个步骤和原先基本一样,唯一需要考虑的就是如何在一点也不用栈的资源的前提下保存寄存器上下文?

堆上。
堆上。
堆上。

简单而言,我们把上下文一股脑都保存到堆上就行。需要保存的上下文大致类似于一个结构体:

typedef struct _THPageVariadicContext {
    int64_t gR[10];              // general registers x0-x8 + x13
    int64_t vR[16];              // float   registers q0-q7
    int64_t linkRegister;        // lr
    int64_t originIMPRegister;   // origin
} THPageVariadicContext;

当然,这里的结构体只是形象化表示内存中的数据顺序和含义,真正使用汇编操作内存的时候,没有结构体。

保存上下文解决了我们不污染栈的诉求,但是同时也引出了一个新的问题,堆分配的地址我们保存在哪?跨函数调用后恢复上下文必须要让我们分配出的堆地址得到“持久化”存储啊。

  • 保存到栈上?这肯定不可能,自己打自己脸嘛。

  • 保存到寄存器上?如果是 caller-saved 寄存器,那不能保证跨函数调用完后,寄存器里面的内容还是我们原先设定的那样;而如果是 callee-saved 寄存器,确实可以解决跨函数调用后数据还原成我们保存的那样。但是同样的,我们自身也是其他 caller 函数的 callee,我们侵占了一个寄存器,怎么在返回到 caller 函数之前复原这个 callee-saved 寄存器呢?

上面这段话有点绕。

所以,我们在分配堆内存的时候,要多分配一个 8 byte 的空间,把侵占的 callee-saved register 的值保存到堆内存中,然后再继续存我们原先要保留的上下文。

关键代码简要概括如下:

  • 第一步,在拦截到函数调用后,先进入我们的 pre 操作,这里是在堆上对应上下文空间大小的地方。需要注意的是,调用分配内存的函数是使用 malloc,我们并不知道 malloc 究竟会破坏哪些寄存器,因为也需要作一次额外的寄存器上下文保存,不过这个保存时短暂的,分配结束后就恢复。然后将这些上下文都保存到堆上。

    attribute((naked))
    void THPageVariadicContextPre(void)
    {

    // 先保存,避免调用 malloc 破坏寄存器
    saveRegs();
    
    // 分配堆上内存 extra 16 byte + sizeof(THPageVariadicContext)
    __asm volatile ("mov x0, #0xF0");
    __asm volatile ("bl _malloc");
    
    // 返回的分配内存地址保存起来 callee-saved
    __asm volatile ("str x19, [x0]");
    __asm volatile ("mov x19, x0");
    
    // 恢复堆栈,避免影响变参所处在的堆栈
    restoreRegs();
    
    // 用堆上空间保存数据
    __asm volatile ("stp x0, x1,  [x19, #(16 + 0 * 16)]");
    __asm volatile ("stp x2, x3,  [x19, #(16 + 1 * 16)]");
    __asm volatile ("stp x4, x5,  [x19, #(16 + 2 * 16)]");
    __asm volatile ("stp x6, x7,  [x19, #(16 + 3 * 16)]");
    __asm volatile ("stp x8, x13, [x19, #(16 + 4 * 16)]");
    
    __asm volatile ("stp q0, q1,  [x19, #(16 + 5 * 16 + 0 * 32)]");
    __asm volatile ("stp q2, q3,  [x19, #(16 + 5 * 16 + 1 * 32)]");
    __asm volatile ("stp q4, q5,  [x19, #(16 + 5 * 16 + 2 * 32)]");
    __asm volatile ("stp q6, q7,  [x19, #(16 + 5 * 16 + 3 * 32)]");
    
    __asm volatile ("stp lr, x10, [x19, #(16 + 5 * 16 + 4 * 32)]");
    
    __asm volatile ("ret");
    

    }

  • 调用完拦截函数,我们需要销毁堆空间,由于我们之前使用的是 callee-saved 的寄存器,我们能确保寄存器的值还是调用之前的。所以我们放心的将其中的值取出来,然后销毁对应的占空间,然后恢复寄存器即可。

    __attribute__((__naked__))
    void THPageVariadicContextPost(void)
    {
        // x19 肯定是正确的地址,使用x19恢复对应的数据
        __asm volatile ("ldp lr, x10, [x19, #(16 + 5 * 16 + 4 * 32)]");
        __asm volatile ("ldp q6, q7,  [x19, #(16 + 5 * 16 + 3 * 32)]");
        __asm volatile ("ldp q4, q5,  [x19, #(16 + 5 * 16 + 2 * 32)]");
        __asm volatile ("ldp q2, q3,  [x19, #(16 + 5 * 16 + 1 * 32)]");
        __asm volatile ("ldp q0, q1,  [x19, #(16 + 5 * 16 + 0 * 32)]");
    
        __asm volatile ("ldp x8, x13, [x19, #(16 + 4 * 16)]");
        __asm volatile ("ldp x6, x7,  [x19, #(16 + 3 * 16)]");
        __asm volatile ("ldp x4, x5,  [x19, #(16 + 2 * 16)]");
        __asm volatile ("ldp x2, x3,  [x19, #(16 + 1 * 16)]");
        __asm volatile ("ldp x0, x1,  [x19, #(16 + 0 * 16)]");
    
        // 保存一下,避免 free 的影响。
        saveRegs();
    
        // 恢复原先的 x19, 调用free
        __asm volatile ("mov x0, x19");
        __asm volatile ("ldr x19, [x19]");
        __asm volatile ("bl _free");
    
        // 恢复堆栈
        restoreRegs();
    
        __asm volatile ("mov lr, x13");
        __asm volatile ("br x10");
    }
    
  • 需要注意的是,我们这里用了__attribute__((__naked__)),这个作用是为了让我们的函数不会额外的生成函数 prologue/epilogue 中的压栈消栈操作。

至此,变参 Hook 就完成了,大家可以前往 Github 查看最新的 THVaradicInterceptor 来使用。

后记

有的朋友会问,为什么很多网上常见的 Hook 方案,都不要这么复杂的上下文保存流程?

其实道理很简单,保存什么上下文取决你的拦截或者 Hook 函数的目的以及使用方式。

举个非常常见的统计函数调用耗时的例子,在这个情形中,一般只用关注 x0x1 两个参数 来记录是什么类什么函数的调用。这种情况下,你的上下文保存可以极简,甚至只要保存 x0, x1 即可。

TrampolineHook 想要提供的拦截器,是一个通用的拦截器,我不能保证其内部的实现,因为我需要保留的上下文就必须很完整。

后续 TrampolineHook 除了完善对 x86_64 的支持外,还有两个比较大的技术目标,也会慢慢完善。如果有什么使用中遇到的问题或者 Bug 也欢迎提交代码。

基于桥的全量方法 Hook 方案(3)- TrampolineHook

作者 SatanWoo
2020年4月26日 00:28

本来以为是双休日,结果五一调休本周末只休一天,懵逼。不过还算完成了承诺,赶了出来。

开源地址:https://github.com/SatanWoo/TrampolineHook

TrampolineHook 是什么

之前杨萧玉在看到我《基于桥的全量方法 Hook 方案(2) - 全新升级》 后就问我这个和直接用 method_exchangeImplementation 之类的 runtime 方法交换 IMP 性能对比咋样?

所以这篇文章开头先占用大家宝贵的两分钟,简要说明下。

TrampolineHook 本质上不是用来 Swizzling 的框架,取 Hook 这个名字只是为了读起来顺口。它实际上是一个中心重定向框架。 换句话说,你可以认为它是为了通过一个函数替换/拦截所有你想要函数的框架。

其实这个中心重定向的思想并不新潮,很多人(包括我自己)在内就曾经利用重载 objc_msgForward 干过这样的事。

但是这个方式我在之前的文章里也提到过对应的缺点,比如:

  • 性能慢
  • 不能替换/拦截同一个继承链上的多个类。

所以可以认为 TrampolineHook 是一个让你不用关注底层架构Calling Convention(因为涉及到汇编),不用关心上下文信息保存、恢复,不用担心引入传统 Swizzle 方案在大型项目中有奇奇怪怪 Crash 问题的中心重定向框架。

TrampolineHook 技术原理

整个技术原理其实可以分为三部分:

  • vm_remap 技术。

  • 流程设计。

  • 汇编实现。

vm_remap 的价值

通俗意义上,我们访问的内存都是按照页来组织。而在程序加载后分配的页之中,会对应有不同的权限,比如代码占用的页,就是可读且可执行,但是一般不具备可写的权限;而存放数据的页呢,就对应是可读且可写,但不能拥有可执行权限。

在绝大多数情况下,当我们编写完一个程序运行的时候,动态分配的页都是用来做数据保存、访问的,不太会有涉及执行权限。

而要做到可以将动态分配出来的内存页具备可执行权限,就需要利用 vm_remap。 它的定义是这样的:

On Darwin, vm_remap() provides support for mapping an existing code page at new address, while retaining the existing page protections; using vm_remap(), we can create multiple copies of existing, executable code, placed at arbitrary addresses.

从定义中我们可以知道两点信息:

  • vm_remap 可以让内存页具备被 map 的页的特性,如果是可执行页被 map,那新创建的页自然而然页具备了这个权限。

  • vm_remap 也不是肆无忌惮的创建任何可执行的页,通俗理解,它只是一个 copy 映射。

上述图片引用自Implementing imp_implementationWithBlock()

因此,我们可以通过在编写代码的过程中,精心构造、预留在程序二进制的代码页,在运行时不断“复制映射”,来完成特殊的使命。

在我们的定义中,我们是构造了连续的两个页

流程设计

要构造特殊的程序二进制代码,首先还是要梳理我们的目的,我们的诉求是所有的函数都能先进入我们的一个中心重定向函数,执行自定义的操作,然后返回原函数,同时这个调用栈不能乱。

  • 把一个我们要替换的原方法 IMP A 取出来,保存起来。
  • 给这个原方法塞一个动态分配的可执行地址 B。
  • 当执行这个原方法的时候,会跳转到 可执行地址 B。
  • 这个 B 经过一段简短的运算操作,可以获取到原先保存的 IMP A。
  • 在跳转回 IMP A 之前,统一拦截函数先做些事情,比如检查是不是主线程调用之类的。

【注意】:在整个过程中,我们要保证参数寄存器、返回地址等不能错乱。

汇编实现

既然 vm_remap 是按页的维度来映射,我们要构造的代码自然而然要页对齐在 arm64 中,一页是 0x4000,也就是 16KB,所以首先就是 .align 14 来确保。

然后上一下最关键部分的代码,感兴趣的还是去 Github 上阅读完整的代码吧。

_th_entry:

// 不要小看这五行汇编
nop
nop
nop
nop
nop

sub x12, lr,   #0x8
sub x12, x12,  #0x4000
mov lr,  x13

ldr x10, [x12]

stp q0,  q1,   [sp, #-32]!
stp q2,  q3,   [sp, #-32]!
stp q4,  q5,   [sp, #-32]!
stp q6,  q7,   [sp, #-32]!

stp lr,  x10,  [sp, #-16]!
stp x0,  x1,   [sp, #-16]!
stp x2,  x3,   [sp, #-16]!
stp x4,  x5,   [sp, #-16]!
stp x6,  x7,   [sp, #-16]!
str x8,        [sp, #-16]!

// 加载自定义的拦截器,并跳转过去。
ldr x8,  interceptor
blr x8

ldr x8,        [sp], #16
ldp x6,  x7,   [sp], #16
ldp x4,  x5,   [sp], #16
ldp x2,  x3,   [sp], #16
ldp x0,  x1,   [sp], #16
ldp lr,  x10,  [sp], #16

ldp q6,  q7,   [sp], #32
ldp q4,  q5,   [sp], #32
ldp q2,  q3,   [sp], #32
ldp q0,  q1,   [sp], #32

br  x10

.rept 2032
mov x13, lr
bl _th_entry;
.endr

整段汇编可以分为几个部分:

  • 设计一大堆的动态可执行地址,即:

    .rept 2032
    mov x13, lr
    bl _th_entry;
    .endr
    

    这里最早我的实现是复制粘贴一大堆重复性代码,在 HookZZ 作者的指导下,我优化成了上述这样。

  • 执行统一的运行过程,通过偏移计算等方式获取保留的原始 IMP。

  • 要注意特定的寄存器用处,x8-x18是临时寄存器,里面的值在函数调用后可能被修改,这些寄存器为caller-saved。所以在我们自身函数可以用,但是要在调用别的函数之前保存好。

  • 要特别注意对 LR 寄存器的处理,没处理好,调用栈就回不去了。

  • 保存对应的参数、浮点参数等寄存器,避免上下文被我们自己的处理函数破坏。

  • b / bl 的跳转范围非常有限,由于我们是动态地址分配,不能保证拦截函数的范围偏移,所以要采用 blr 的方式。

TrampolineHook 用处

和传统的 Swizzle 需要提供对应的替换后的函数实现不同,中心化重定向思想可以帮助你实现很多有意思的事情:

  • 比如网上很常见的 hook objc_msgSend,可以帮你查看任意被 Hook 二进制中的函数耗时和调用链路。

  • 比如 Bang / AnyMethodLog 这样的重定向 Log 日志框架等等。

苹果著名的 MainThreadChecker 也用了类似的技术。由于我才疏学浅,只是大致完成了对其实现的逆向,通过 TrampolineHook 进行了重写。 因为效果还不错,所以也开源了出来,地址是:https://github.com/SatanWoo/TrampolineHook/tree/master/Example/MainThreadChecker

这次在重写 MainThreadChecker 的过程中,我也对比了下和 2017 年苹果实现的差异。在整体流程上没有比较大的差异,但是还是有一些细节可以分享分享:

  • iOS 10 的时候对应的二进制是 UIKit,到了 iOS 12/13 成了 UIKitCore,所以原先获取二进制的逻辑失效了,为了避免后续版本的变更干扰,我采用了苹果自身的守候,通过 class_getImageName([UIResponder class]) 来保证获取的就是我们理解上的 UIKit 动态库。

当然 TrampolineHook 的作用不止于此,争取过段时间把我的一些想法做完善再和大家交流。

后续思考

本质上 Trampolinevm_remap 技术不是新的技术,很早就有人应用了,构造 Trampoline 实际上在苹果自身关于 Block 的实现中就有。业界也有 SwiftTrace 也是用了对应的技术。

真正的关键在于你用 Trampoline 做什么?用途的不同也决定了效果的不同,这也是我把之前的代码重写 TrampolineHook 中所收获的,而且随着 TrampolineHook 相对我自身之前实现的优化,我发现眼前豁然开朗,能玩的事情还有很多,哈哈。

对了,如果有朋友对 arm64 的汇编比较熟悉,同时对函数调用也比较了解的话,会很快的发现我上述提供的汇编代码存在一个漏洞(虽然这个漏洞绝大多数人用不到),感兴趣的朋友可以微信交流下。

开源地址:https://github.com/SatanWoo/TrampolineHook 如果大家有什么想法或者遇到了自身项目中的 Bug,欢迎 issue。

了解 SIMD 指令

作者 SatanWoo
2019年12月1日 23:39

了解 SIMD 指令

SIMD 是一种常见的利用单指令完成多数据量处理的计算方式。本文作为 SIMD 文章的引子,先来了解简单的 SIMD 使用和概念。

SIMD 的含义

SIMD 的全称是 Single Instruction Multiple Data。简要来说,就是通过一条指令完成多条数据处理的行为。我们知道,虽然程序是由一条条机器指令组成,但是实际上执行一条机器码包含了多个过程,包含取指令、分析指令到执行等,如下图所示(暂时先忽略流水线并行)

而在这其中,每一个阶段,都会消耗一个或多个机器周期。如果我们认为,取指令和分析指令(译码)可以近似的认为是一个机器周期内完成,那么不同的指令,在执行阶段耗费的机器周期则大不相同。

举个例子,可能加法指令的执行阶段需要两个机器周期;而乘法可能需要5-6个机器周期。那么,当我们无法缩短指令的执行周期缩短的时候,利用 SIMD 技术,则可以在相同的执行周期内完成更多的数据处理,这样也同等的提升了单位时间内的数据吞吐,提高了计算性能。

在 Intel 的手册上,提供了包含 MMX, SSE, AVX 等系列的并行指令,面向不同长度的数据并行,比如:

  • MMX 并行计算 64bit 的数据。
  • SSE 并行计算 128bit 的数据。
  • AVX 并行计算 256bit 的数据。
  • AVX512 并行计算 512bit 的数据。

更多详细的使用可以参考:

Intel 手册

SIMD 的使用方式

由于绝大多数的人对 SIMD 还不甚了解,因此本文基于大家比较熟悉的环境 Xcode + x86/64 架构来完成。

主要是我懒,不想再翻 ARM 的手册了。

这里我们以一个简单的 256bit (32 byte) 加法改写成 SIMD 的形式来验证:

原始版本:

double input1[k] = {1, 2, 3, 4};
double input2[k] = {5, 6, 7, 8};
double result[k] = {0};

for (int i = 0; i < k; i++) {
    result[i] = input1[i] + input2[i];
}

SIMD 版本:

const int k = 4;
double input1[k] = {1, 2, 3, 4};
double input2[k] = {5, 6, 7, 8};
double result[k] = {0};

__m256d a = _mm256_load_pd(input1);
__m256d b = _mm256_load_pd(input2);

__m256d c = _mm256_add_pd(a, b);
_mm256_store_pd(result, c);

原始版本比较好懂,我们主要来深入看下 SIMD 中代码的意思:

  • _mm256_load_pd 就是从内存中读取一个地址,这个地址返回为 __m256d 的向量(256bit)。其中, __mm256d的定义为下:

    typedef double __m256d __attribute__((__vector_size__(32)));
    

    这个含义的意思就是 __m256d 的长度是 32 byte(256bit),而这个 32 byte 是按照 4 个 double 元素构成的。

  • _mm256_add_pd 就是对两个 256bit 的向量元素进行直接相加。

  • _mm256_store_pd 就是 _mm256_load_pd的逆运算,不再赘述。

注意:如果提示需要 AVX 支持的话,请在 Xcode 对应的代码文件处添加 Compiler Flag: -mavx

用 SIMD 实现求和加法

既然说了 SIMD 的本质还是为了提升单位时间内的计算吞吐量,我们还是用一个简单的例子,加法求和来实践一下:

常规的代码如下:

double CommonAdd(double *data, int count)
{
    double result = 0;

    for (int i = 0 ; i < count; i++) {
        result += data[i];
    }

    return result;
}

SIMD 的代码如下:

double AVXAdd(double *data, int count)
{
    int offset = 0;

    __m256d v1;
    __m256d sum = _mm256_setzero_pd();

    double ret = 0;

    for (int i = 0; i < count/4; i++) {
        v1 = _mm256_load_pd(data + offset);
        sum = _mm256_add_pd(sum, v1);
        offset += 4;
    }

    sum = _mm256_hadd_pd(sum, sum); // 水平求和

    ret += sum[0];
    ret += sum[2];

    return ret;
}

测试代码如下:

int main() {

    struct  timeval   start;
    struct  timeval   end;


    const int k = 512 * 512;

    const int loop = 1;

    double input1[k];

    for (int i = 0; i < k; i++) {
        input1[i] = i;
    }

    gettimeofday(&start, nullptr);

    for (int j = 0; j < loop; j++) {
        CommonAdd(input1, k);
    }

    gettimeofday(&end, nullptr);

    printf("tv_sec:%ld\n",end.tv_sec - start.tv_sec);
    printf("tv_usec:%d\n", end.tv_usec - start.tv_usec);

    std::cout << " ======================= " << std::endl;

    gettimeofday(&start, nullptr);

    for (int j = 0; j < loop; j++) {
        AVXAdd(input1, k);
    }

    gettimeofday(&end, nullptr);

    printf("tv_sec:%ld\n",end.tv_sec - start.tv_sec);
    printf("tv_usec:%d\n", end.tv_usec - start.tv_usec);

    return 0;
}

这里,我们选择了图像处理里面比较常见的 512 * 512 大小来做验证,在我的 2015款 MacBookPro 上可以得到大致如下两个性能耗时:

  • 常规方法 【774 us】
  • SIMD 【560 us】

别小看这一点的性能差距,对于大运算量的端侧深度学习可就有很显著的差距了。

后记

本文只是仅仅介绍了最常规的 SIMD 使用方式。但是在实际设计的过程中,不可能像我们这么简单的去应用。随之而来的,你会发现伴随着许多不同的坑,包含不规范的应用导致性能的下降崩溃问题。这些都会留在后面我们去解决。

❌
❌