使用Chisel搭建Systolic Array

luoganttcc

已于 2023-01-31 16:13:01 修改

阅读量808

点赞数 1

分类专栏：数字电路文章标签：人工智能

于 2023-01-31 16:12:30 首次发布

本文链接：https://blog.csdn.net/luoganttcc/article/details/128819226

版权

数字电路专栏收录该内容

100 篇文章 11 订阅

订阅专栏

最近听到非常多人吹Chisel，为了方便快速做算法实现，就去尝试学了下，发现确实很香，有种相见恨晚的感觉。

今天是使用Chisel搭建个脉动阵列（Systolic Array, SA）^[1]，脉动阵列是神经网络中最基础也是最常用的组件之一。它主要是实现高效的矩阵乘法，而神经网络中到处都是矩阵乘法，从全连接层到卷积层（GEMM），都可以用矩阵乘法表示。UCB的Chipyard中专门集成了个脉动阵列的生成器Gemmini。

GitHub - ucb-bar/gemmini: Berkeley's Spatial Array Generator github.com/ucb-bar/gemmini

Gemmini考虑了很多的场景，比较复杂。本实现在参考Gemmini的基础上实现了个基础功能的用于全连接层的Weight-stationary的脉动阵列，即权重存储在PE内，只有输入激励在脉动阵列里传输。这样可以简化权重的加载。

SA是由许多个PE（Processing Element）组成的，每个PE执行的是乘加运算，然后将部分和的结果打拍后放入下一个PE，下一个PE将乘法结果与部分和相加得到新的部分和后传递到再下一个PE。最终在每列的最后一个PE得到总和结果，也就是一个神经元的输出。

本实现中的PE采用的是Weight-Stationary架构，如下图所示。Weight存储在Register中，Weight与输入的Input Activation相乘的结果再与上方PE的部分和结果相加，再传入下一个PE，其中的MUX的目的是让部分和与权重共用一组总线，这样可以节省布线资源。在加载权重时，就不让部分和传递，而是让权重一层层传递下去。（注：这里没有将下方和右方的Register画出来）

与Weight-Stationary相对的有Ouput-Stationary架构，就是每个PE中存储的是Output Activation，由于没有做实现，因此不展开叙述了。

首先实现PE，Chisel代码如下。这里默认读者有点Chisel基础，可以参考Chisel-book^[2]，并跑完Chisel Bootcamp^[3]做入门学习。下面代码定义了5个输入输出，包括输入的激励inAct，传递到右边PE的输入激励的打拍结果outAct，权重和部分和Partial Sum共用的输入和输出线inWtPS和outWtPS，还有个inwtValid信号是告诉PE当前应该传递权重而不是部分和。这里将激励位宽与权重位宽作为了可配置的参数actWidth和weightWidth。

import chisel3._
import chisel3.util._ 
class PE(actWidth: Int = 16, weightWidth: Int = 16) extends Module{
    val io = IO(new Bundle{
        val inAct       = Input(SInt(actWidth.W))
        val inWtPS      = Input(SInt((actWidth+weightWidth).W))
        val outWtPS     = Output(SInt((actWidth+weightWidth).W))
        val outAct      = Output(SInt(actWidth.W)) // 向右传递
        val inwtValid   = Input((UInt(1.W))) 
    })
    val parSumWidth = actWidth + weightWidth;
    // 不要初始化，避免复位
    val weightReg   = Reg(SInt(weightWidth.W))
    val actReg      = Reg(SInt(actWidth.W))
    val parSumReg   = Reg(SInt(parSumWidth.W))
    actReg    := io.inAct;
    parSumReg := io.inWtPS + weightReg * actReg;
    io.outAct := actReg;
    // 会自动截低位，不用担心位宽，甚至还会在高位补符号位（有符号数）和0（无符号数）
    when(io.inwtValid === 1.U){
        weightReg   := io.inWtPS
        io.outWtPS  := weightReg
    } .otherwise {
        weightReg   := weightReg
        io.outWtPS  := parSumReg
    }
}

然后是实现Systolic Array，也就是例化多个PE，这里例化的个数由meshRows和meshColumns决定，分别表示Systolic Array的高度和宽度。例化时采用了fill函数，即List中有meshRows个元素，每个元素也是个List，含有meshColumns个PE，因此总PE个数是meshRows * meshColumns。

然后就是inwtValid线连接到每个PE上，这里使用了Scala的语法糖foreach（Scala推荐函数式编程，能不用for循环就不用for循环），在PEs.foreach(_.foreach(_.io.inwtValid := io.inwtValid)) 中，第一个foreach表示对每个PEs的元素（也是个List）做操作，操作仍然是_.foreach，因为List内每个元素仍然是个List，第二个foreach才是对单个PE做操作，指定inwtValid的连接。

import chisel3._
import chisel3.util._ 
class SystolicArray(actWidth: Int = 16, weightWidth: Int = 16,
                    meshRows: Int = 4, meshColumns: Int = 4) extends Module{
  val io = IO(new Bundle{
    val inAct     = Input(Vec(meshRows, SInt(actWidth.W)))
    val inWeight  = Input(Vec(meshColumns, SInt(weightWidth.W)))
    val outSum    = Output(Vec(meshColumns, SInt((actWidth+weightWidth).W)))
    val inwtValid = Input((UInt(1.W)))
  })
  val parSumWidth = actWidth + weightWidth; 
  // PEs(r)(c)表示第r行，第c列的PE，从0开始
  val PEs: Seq[Seq[PE]] = Seq.fill(meshRows, meshColumns)(Module(new PE(actWidth, weightWidth)))
  PEs.foreach(_.foreach(_.io.inwtValid := io.inwtValid));
  for(row <- 0 until meshRows){
    for(col <- 0 until meshColumns){
      // 第一行的输入连接到SA端口
      if(row == 0)  PEs(row)(col).io.inWtPS := io.inWeight(col)
      else          PEs(row)(col).io.inWtPS := PEs(row-1)(col).io.outWtPS
      // 最后一行的输出连接到模块外
      if(row == meshRows-1)   io.outSum(col) := PEs(row)(col).io.outWtPS  
      // 第一列
      if(col == 0)  PEs(row)(col).io.inAct := io.inAct(row)
      else          PEs(row)(col).io.inAct := PEs(row)(col-1).io.outAct
    }
  }
}