tvm在CPU上优化GEMM结果

本文展示如何只添加18行code，在矩阵乘法上获得200+倍的加速。

通常，CPU上的计算密集型任务有2个优化点：

  • 提高内存访问的缓存命中率
  • SIMD指令加速

对于gemm的优化手段已有现成的总结，基本都可以在这篇文档how to optimize gemm找到。

tvm已经实现了其中的一些优化方法，但由于tvm本身的限制，还有一些方法没有实现。

本文逐步优化，不断提升程序性能。首先用没有优化的code和numpy运行结果做对比，如下：

这是我自己的容器里运行的结果

Numpy running time: 0.004862
Baseline: 2.646903

原始ir如下：

produce C {
  for (x, 0, 1024) {
    for (y, 0, 1024) {
      C[((x*1024) + y)] = 0.000000f
      for (k, 0, 1024) {
        C[((x*1024) + y)] = (C[((x*1024) + y)] + (A[((x*1024) + k)]*B[(y + (k*1024))]))
      }
    }
  }
}

1 blocking

使用blocking的技术可以显著提升缓存命中率，因为数据被分块进行计算，块内的数据在缓存中的访问都是相邻的。

结果如下：

Numpy running time: 0.004732
Baseline: 2.892019
Opt1: 0.701635

优化后ir

produce C {
  for (x.outer, 0, 32) {
    for (y.outer, 0, 32) {
      for (x.inner.init, 0, 32) {
        for (y.inner.init, 0, 32) {
          C[(((((x.outer*1024) + y.outer) + (x.inner.init*32))*32) + y.inner.init)] = 0.000000f
        }
      }
      for (k.outer, 0, 256) {
        for (k.inner, 0, 4) {
          for (x.inner, 0, 32) {
            for (y.inner, 0, 32) {
              C[(((((x.outer*1024) + y.outer) + (x.inner*32))*32) + y.inner)] = (C[(((((x.outer*1024) + y.outer) + (x.inner*32))*32) + y.inner)] + (A[(((((x.outer*8192) + k.outer)*4) + k.inner) + (x.inner*1024))]*B[((((y.outer + (k.outer*128)) + (k.inner*32))*32) + y.inner)]))
            }
          }
        }
      }
    }
  }
}

2 Vectorization

向量化。
结果：

Numpy running time: 0.004964
Baseline: 2.884543
Opt1: 0.713341
Opt2: 0.331218

优化后ir

produce C {
  for (x.outer, 0, 32) {
    for (y.outer, 0, 32) {
      for (x.inner.init, 0, 32) {
        C[ramp(((((x.outer*1024) + y.outer) + (x.inner.init*32))*32), 1, 32)] = x32(0.000000f)
      }
      for (k.outer, 0, 256) {
        for (k.inner, 0, 4) {
          for (x.inner, 0, 32) {
            C[ramp(((((x.outer*1024) + y.outer) + (x.inner*32))*32), 1, 32)] = (C[ramp(((((x.outer*1024) + y.outer) + (x.inner*32))*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer)*4) + k.inner) + (x.inner*1024))])*B[ramp((((y.outer + (k.outer*128)) + (k.inner*32))*32), 1, 32)]))
          }
        }
      }
    }
  }
}

3 Loop Permutation

结果：

Numpy running time: 0.005203
Baseline: 2.646298
Opt1: 0.691242
Opt2: 0.330293
Opt3: 0.147917

优化后ir：

produce C {
  for (x.outer, 0, 32) {
    for (y.outer, 0, 32) {
      for (x.inner.init, 0, 32) {
        C[ramp(((((x.outer*1024) + y.outer) + (x.inner.init*32))*32), 1, 32)] = x32(0.000000f)
      }
      for (k.outer, 0, 256) {
        for (x.inner, 0, 32) {
          for (k.inner, 0, 4) {
            C[ramp(((((x.outer*1024) + y.outer) + (x.inner*32))*32), 1, 32)] = (C[ramp(((((x.outer*1024) + y.outer) + (x.inner*32))*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.inner*256))*4) + k.inner)])*B[ramp((((y.outer + (k.outer*128)) + (k.inner*32))*32), 1, 32)]))
          }
        }
      }
    }
  }
}

4 Array Packing

结果：

Numpy running time: 0.005159
Baseline: 2.884619
Opt1: 0.693074
Opt2: 0.332173
Opt3: 0.149278
Opt4: 0.233195

这一步优化，连续跑了两次，性能都反而变差了。

优化后ir：


// attr [packedB] storage_scope = "global"
allocate packedB[float32x32 * 32 * 1024 * 1]
produce packedB {
  parallel (x, 0, 32) {
    for (y, 0, 1024) {
      packedB[ramp((((x*1024) + y)*32), 1, 32)] = B[ramp(((x + (y*32))*32), 1, 32)]
    }
  }
}
produce C {
  for (x.outer, 0, 32) {
    for (y.outer, 0, 32) {
      for (x.inner.init, 0, 32) {
        C[ramp(((((x.outer*1024) + y.outer) + (x.inner.init*32))*32), 1, 32)] = x32(0.000000f)
      }
      for (k.outer, 0, 256) {
        for (x.inner, 0, 32) {
          for (k.inner, 0, 4) {
            C[ramp(((((x.outer*1024) + y.outer) + (x.inner*32))*32), 1, 32)] = (C[ramp(((((x.outer*1024) + y.outer) + (x.inner*32))*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.inner*256))*4) + k.inner)])*packedB[ramp((((((y.outer*256) + k.outer)*4) + k.inner)*32), 1, 32)]))
          }
        }
      }
    }
  }
}

5 Write cache for blocks

结果：

Numpy running time: 0.005358
Baseline: 2.654734
Opt1: 0.689408
Opt2: 0.329072
Opt3: 0.148742
Opt4: 0.231431
Opt5: 0.211086

优化后ir：

// attr [packedB] storage_scope = "global"
allocate packedB[float32x32 * 32 * 1024 * 1]
// attr [C.global] storage_scope = "global"
allocate C.global[float32 * 32 * 32]
produce packedB {
  parallel (x, 0, 32) {
    for (y, 0, 1024) {
      packedB[ramp((((x*1024) + y)*32), 1, 32)] = B[ramp(((x + (y*32))*32), 1, 32)]
    }
  }
}
produce C {
  for (x.outer, 0, 32) {
    for (y.outer, 0, 32) {
      produce C.global {
        for (x.c.init, 0, 32) {
          C.global[ramp((x.c.init*32), 1, 32)] = x32(0.000000f)
        }
        for (k.outer, 0, 256) {
          for (x.c, 0, 32) {
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[((((x.outer*8192) + k.outer) + (x.c*256))*4)])*packedB[ramp((((y.outer*256) + k.outer)*128), 1, 32)]))
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.c*256))*4) + 1)])*packedB[ramp(((((y.outer*256) + k.outer)*128) + 32), 1, 32)]))
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.c*256))*4) + 2)])*packedB[ramp(((((y.outer*256) + k.outer)*128) + 64), 1, 32)]))
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.c*256))*4) + 3)])*packedB[ramp(((((y.outer*256) + k.outer)*128) + 96), 1, 32)]))
          }
        }
      }
      for (x.inner, 0, 32) {
        for (y.inner, 0, 32) {
          C[(((((x.outer*1024) + y.outer) + (x.inner*32))*32) + y.inner)] = C.global[((x.inner*32) + y.inner)]
        }
      }
    }
  }
}

6 Parallel

结果：

Numpy running time: 0.005989
Baseline: 2.635383
Opt1: 0.691006
Opt2: 0.328837
Opt3: 0.149464
Opt4: 0.233010
Opt5: 0.213697
Opt6: 0.018374

优化后ir：

// attr [packedB] storage_scope = "global"
allocate packedB[float32x32 * 32 * 1024 * 1]
produce packedB {
  parallel (x, 0, 32) {
    for (y, 0, 1024) {
      packedB[ramp((((x*1024) + y)*32), 1, 32)] = B[ramp(((x + (y*32))*32), 1, 32)]
    }
  }
}
produce C {
  parallel (x.outer, 0, 32) {
    // attr [C.global] storage_scope = "global"
    allocate C.global[float32 * 32 * 32]
    for (y.outer, 0, 32) {
      produce C.global {
        for (x.c.init, 0, 32) {
          C.global[ramp((x.c.init*32), 1, 32)] = x32(0.000000f)
        }
        for (k.outer, 0, 256) {
          for (x.c, 0, 32) {
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[((((x.outer*8192) + k.outer) + (x.c*256))*4)])*packedB[ramp((((y.outer*256) + k.outer)*128), 1, 32)]))
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.c*256))*4) + 1)])*packedB[ramp(((((y.outer*256) + k.outer)*128) + 32), 1, 32)]))
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.c*256))*4) + 2)])*packedB[ramp(((((y.outer*256) + k.outer)*128) + 64), 1, 32)]))
            C.global[ramp((x.c*32), 1, 32)] = (C.global[ramp((x.c*32), 1, 32)] + (x32(A[(((((x.outer*8192) + k.outer) + (x.c*256))*4) + 3)])*packedB[ramp(((((y.outer*256) + k.outer)*128) + 96), 1, 32)]))
          }
        }
      }
      for (x.inner, 0, 32) {
        for (y.inner, 0, 32) {
          C[(((((x.outer*1024) + y.outer) + (x.inner*32))*32) + y.inner)] = C.global[((x.inner*32) + y.inner)]
        }
      }
    }
  }
}

参考： https://docs.tvm.ai/tutorials/optimize/opt_gemm.html

©著作权归作者所有,转载或内容合作请联系作者
  • 人面猴
    序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
    沈念sama阅读 216,287评论 6赞 498
  • 死咒
    序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
    沈念sama阅读 92,346评论 3赞 392
  • 救了他两次的神仙让他今天三更去死
    文/潘晓璐 我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
    开封第一讲书人阅读 162,277评论 0赞 353
  • 道士缉凶录：失踪的卖姜人
    文/不坏的土叔 我叫张陵，是天一观的道长。 经常有香客问我，道长，这世上最难降的妖魔是什么？ 我笑而不...
    开封第一讲书人阅读 58,132评论 1赞 292
  • 港岛之恋（遗憾婚礼）
    正文 为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
    茶点故事阅读 67,147评论 6赞 388
  • 恶毒庶女顶嫁案：这布局不是一般人想出来的
    文/花漫 我一把揭开白布。 她就那样静静地躺着，像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上，一...
    开封第一讲书人阅读 51,106评论 1赞 295
  • 城市分裂传说
    那天，我揣着相机与录音，去河边找鬼。 笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。 我是一名探鬼主播，决...
    沈念sama阅读 40,019评论 3赞 417
  • 双鸳鸯连环套：你想象不到人心有多黑
    文/苍兰香墨 我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
    开封第一讲书人阅读 38,862评论 0赞 274
  • 万荣杀人案实录
    序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
    沈念sama阅读 45,301评论 1赞 310
  • 护林员之死
    正文 独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,521评论 2赞 332
  • 白月光启示录
    正文 我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,682评论 1赞 348
  • 活死人
    序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
    沈念sama阅读 35,405评论 5赞 343
  • 日本核电站爆炸内幕
    正文 年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
    茶点故事阅读 40,996评论 3赞 325
  • 男人毒药：我在死后第九天来索命
    文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,651评论 0赞 22
  • 一桩弑父案，背后竟有这般阴谋
    文/苍兰香墨 我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,803评论 1赞 268
  • 情欲美人皮
    我被黑心中介骗来泰国打工， 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。 一个月前我还...
    沈念sama阅读 47,674评论 2赞 368
  • 代替公主和亲
    正文 我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
    茶点故事阅读 44,563评论 2赞 352

推荐阅读更多精彩内容

  • Android - 收藏集
    Android 自定义View的各种姿势1 Activity的显示之ViewRootImpl详解 Activity...
    passiontim阅读 172,067评论 25赞 707
  • Android - 收藏集
    用两张图告诉你，为什么你的 App 会卡顿? - Android - 掘金 Cover 有什么料？ 从这篇文章中你...
    hw1212阅读 12,712评论 2赞 59
  • 小燕子出生第14天
    今天小燕子的爸爸跑步4.9公里去大华超市（99 ranch market）买了木瓜、鲫鱼、排骨、猪脚等下奶食材。跑...
    李林燕她爸阅读 235评论 0赞 0
  • 渐行渐远的我们
    我不知道从何时起，我们的距离越来越远。直到现在我还能记起我们第一次见面的情景，你跟在璐璐旁边，默默的看着我笑。那一...
    橘枳qie阅读 250评论 0赞 1
  • 采撷春色
    耿俨阅读 214评论 4赞 2