what is a low-latency hardware algorithm to find the maximum value in an array

Viewed 136

I need to build a low-latency, single-cycle hardware module to find the index and value of the maximum element in an array. Currently, I am using a comparator tree, but the latency is unsatisfactory. So are there any other algorithms that might have lower latency?

I expect the input array to be large (256 to 4096 elements) and the values to be small (3 to 5 bits). Also, I expect the array to be sparse i.e. many small values and few large values.

I am primarily concerned with latency; area is not quite as important.

My current implementation that uses a comparator tree looks something like this:

implicit class reduceTreeOp[A](seq: Seq[A]) {
  def reduceTree[B >: A](op: (B, B) => B): B = {
    if(seq.length == 0)
      throw new NoSuchElementException("cannot reduce empty Seq")
    var rseq: Seq[B] = seq
    while(rseq.length != 1)
      rseq = rseq.grouped(2).toSeq
        .map(s => if(s.length == 1) s(0) else op(s(0), s(1)))
    rseq(0)
  }
}
val (value, index) = array
  .zipWithIndex
  .map{case (v, i) => (v, i.U)}
  .reduceTree[(UInt, UInt)]{case ((val1, idx1), (val2, idx2)) =>
    val is1 = val1 >= idx2
    ( Mux(is1, val1, idx2),
      Mux(is1, idx1, idx2))
  }

FWIW this is designed for 7nm hardware; though I doubt this actually matters to my question.

2 Answers

This is pretty simple thing to do in chisel3. The constraint that results are returned in a single cycle will cause this to generate a whole bunch of hardware. It is probably a good idea to carefully evaluate what you need here.

Never-the-less it's an interesting problem and demonstrates some of the power of chisel. I have provided an runnable example Scastie example.

Here is the code presented with a very simple test suite

import chisel3._
import chisel3.util.log2Ceil
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec
import treadle.extremaOfUIntOfWidth

/** write only memory that continuously outputs the first index of the highest value in an array and
  * that highest value. It works by building a evaluation network of highest values
  *
  * It would be trivial to add ability to read the values in this memory
  *
  * @param depth
  * @param bitWidth
  */
class ArrayMax(val depth: Int, val bitWidth: Int) extends MultiIOModule {
  val writeEnable = IO(Input(Bool()))
  val writeAddress = IO(Input(UInt(log2Ceil(depth).W)))
  val writeData = IO(Input(UInt(bitWidth.W)))

  val maxValue = IO(Output(UInt(bitWidth.W)))
  val indexOfMaxValue = IO(Output(UInt(log2Ceil(depth).W)))

  val array = Reg(Vec(depth, UInt(bitWidth.W)))

  when(writeEnable) {
    array(writeAddress) := writeData
  }

  val valuesAndIndices = array.zipWithIndex.map { case (value, index) => (value, index.U)}.toList

  // Look through the array pair wise and return the value and index of the higher of the pair
  def compareAdjacentValues(valuesAndIndices: Seq[(UInt, UInt)]): Seq[(UInt, UInt)] = {
    val pairs = valuesAndIndices.sliding(2, 2)
    pairs.map {
      case (aValue, aIndex) :: (bValue, bIndex) :: Nil =>
        val (higherValue, higherIndex) = (Wire(UInt(bitWidth.W)), Wire(UInt(log2Ceil(depth).W)))

        when(aValue < bValue) {
          higherValue := bValue
          higherIndex := bIndex
        } otherwise {
          higherValue := aValue
          higherIndex := aIndex
        }
        (higherValue, higherIndex)
      case (aValue, aIndex) :: Nil =>
        (aValue, aIndex)
      case a =>
        throw new Exception("Cannot get here, sliding should return list of size 1 or 2, $a")
    }.toList
  }

  def reduceToOne(pairs: Seq[(UInt, UInt)]): (UInt, UInt) = {
    if(pairs.length == 1) {
      pairs.head
    } else {
      reduceToOne(compareAdjacentValues(pairs))
    }
  }

  val (highestValue, index) = reduceToOne(valuesAndIndices)

  maxValue := highestValue
  indexOfMaxValue := index
}

/** Pumps random values at random indices into the dut and buffer that models it
  * Checks to see that the first highest value (there can be multiple occurrences)
  * is returned and the first index where that value appears.
  *
  * A simpler model might try to keep a record of the highest value in the array
  * but that will break down if that value is replaced by something lower
  */
class ArrayMaxSpec extends AnyFreeSpec with ChiselScalatestTester {
  val rand = new scala.util.Random

  "should act like a little write only memory" in {
    test(new ArrayMax(depth = 256, bitWidth = 8)) { dut =>

      val testArray = Array.fill(dut.depth)(0)

      for(i <- 0 until dut.depth * 10) {
        val index = rand.nextInt(dut.depth)
        dut.writeEnable.poke(true.B)
        dut.writeAddress.poke(index.U)
        val newValue = rand.nextInt(extremaOfUIntOfWidth(dut.bitWidth)._2.toInt)
        dut.writeData.poke(newValue.U)
        testArray(index) = newValue

        dut.clock.step()

        dut.maxValue.expect(testArray.max.U)
        dut.indexOfMaxValue.expect(testArray.indexOf(testArray.max).U)
      }
    }
  }
}

I found an algorithm that works by inverting the array. Essentially, a new array is created with a slot for each possible value, each element of the original array is sorted into a slot depending on the value, and finally, the new array can be priority encoded to find the highest slot with an element.

The implementation looks something like the following:

// VALUE_SPACE: number of possible values
val exists = WireDefault(VecInit(Seq.fill(VALUE_SPACE)(false.B)))
val index = Wire(Vec(VALUE_SPACE, UInt(log2ceil(array.length).W)))
index := DontCare

array.zipWithIndex.foreach{case (v, i) =>
  exists(v) := true.B
  index(v) := i.U
}

maxVal := VALUE_SPACE.U - PriorityEncoder(exists.reverse)
maxidx := PriorityMux(exists.reverse, index.reverse)

This algorithm gave me approximately 2.5x speedup over the comparator tree when using an array with 256 4-bit elements. However, it used much more area, and it probably only has benefit if the number of possible values is very small.

Related