package magmasi.components

import chisel3._
import chisel3.util._

class flexdpecom2(IN_DATA_TYPE : Int = 32 ,DATA_TYPE: Int = 32, NUM_PES: Int = 32, LOG2_PES: Int = 5) extends Module {
    implicit val config = MagmasiConfig()
  val io = IO(new Bundle {
    val i_vn = Input(Vec(NUM_PES, UInt(LOG2_PES.W)))
    val i_stationary = Input(Bool())
    val i_data_valid = Input(Bool())

     val i_mux_bus   = Input(Vec(2 * (LEVELS - 2) * NUM_PES + NUM_PES, Bool()))
  })
    val LEVELS   : Int = (2 * (math.log(NUM_PES) / math.log(2))).toInt + 1

     val r_mult = RegInit(VecInit(Seq.fill(NUM_PES)(0.U((DATA_TYPE-1).W))))

    val r_stationary_ff = Reg(Bool())
    val r_stationary_ff2 = Reg(Bool())

    val r_data_valid_ff = Reg(Bool())
    val r_data_valid_ff2 = Reg(Bool())

    r_stationary_ff := io.i_stationary
    r_stationary_ff2 := r_stationary_ff

     r_data_valid_ff := io.i_data_valid
     r_data_valid_ff2 := r_data_valid_ff

    val my_controller = Module(new fancontrol(32,32,5))

    my_controller.io.i_vn := io.i_vn
    my_controller.io.i_stationary := 0.B
    my_controller.io.i_data_valid := io.i_data_valid
    val w_reduction_add = my_controller.io.o_reduction_add
    val w_reduction_cmd = my_controller.io.o_reduction_cmd
    val w_reduction_sel = my_controller.io.o_reduction_sel
    val w_reduction_valid = my_controller.io.o_reduction_valid

    val my_Benes = Module(new Benes())

    my_Benes.io.i_data_bus1 := io.i_data_bus
    my_Benes.io.i_data_bus2 := io.i_data_bus2
    my_Benes.io.i_mux_bus := io.i_mux_bus

    val w_dist_bus1 = my_Benes.io.o_dist_bus1
    val w_dist_bus2 = my_Benes.io.o_dist_bus2

     val buffer_mult = Module(new buffer_multiplication())

      buffer_mult.io.buffer1 := w_dist_bus1
      buffer_mult.io.buffer2 := w_dist_bus2

      r_mult := buffer_mult.io.out

    val my_fan_network = Module(new FanNetworkcom(32,32))

    my_fan_network.io.i_valid := w_reduction_valid
    my_fan_network.io.i_data_bus :=   r_mult
    my_fan_network.io.i_add_en_bus := w_reduction_add
    my_fan_network.io.i_cmd_bus := w_reduction_cmd
    my_fan_network.io.i_sel_bus := w_reduction_sel
    io.o_valid := my_fan_network.io.o_valid
    io.o_data_bus := my_fan_network.io.o_data_bus
    io.o_adder :=  my_fan_network.io.o_adder
}

    val i_data_bus = Input(Vec(NUM_PES, UInt(DATA_TYPE.W)))
     val i_data_bus2  = Input(Vec(NUM_PES, UInt(DATA_TYPE.W)))

    val o_valid = Output(Vec(NUM_PES, UInt(1.W)))


package magmasi.components
import scala.math

case class MagmasiConfig (
    val DATA_TYPE: Int = 16,
    val NUM_PES  : Int = 4,

    val rowsA :Int = 4,
    val colsA :Int = 4,
    val rowsB :Int = 4,
    val colsB :Int = 4
) {
    val LEVELS   : Int = (2 * (math.log(NUM_IN) / math.log(2))).toInt + 1
}

    val NUM_IN  : Int = 4,
    val SEL_IN: Int = 2,


package magmasi.components
import chisel3._
import chisel3.util._

class MultSwitch() extends Module {
  val io = IO(new Bundle {

    val i_valid = Input(Bool())
    val i_data = Input(UInt(16.W))
    val i_stationary = Input(Bool())
    val o_valid = Output(Bool())
    val o_data = Output(UInt(32.W))
  })

  val r_buffer = RegInit(0.U(16.W))
  val r_buffer_valid = RegInit(false.B)

   // Concatenate the temporary signal with 16 zeros and assign it to io.o_data

  // Logic to store correct value into the stationary buffer

  when(io.i_stationary && io.i_valid) {
    r_buffer := io.i_data // Latch the stationary value into the switch buffer
    r_buffer_valid := true.B // Validate buffer
  }

  val w_A = Mux(r_buffer_valid && io.i_valid, io.i_data, 0.U) // Streaming
  val w_B = Mux(r_buffer_valid && io.i_valid, r_buffer, 0.U) // Stationary

  // Logic to generate correct output valid signal
  val o_validReg = RegInit(false.B)

  io.o_valid := o_validReg
   val multiplier_output = Wire(UInt(32.W))
  // Instantiate multiplier
  val multiplier_module = Module(new SimpleMultiplier)
  multiplier_module.io.A := w_A
  multiplier_module.io.B := w_B
  // Create a temporary signal to hold multiplier output
  multiplier_output := multiplier_module.io.O // Assign the value from the multiplier to the temporary signal

  io.o_data := Cat(0.U(16.W), multiplier_output(15, 0))

  // Convert BF16 to FP32

}

  when (r_buffer_valid && io.i_valid) {
    o_validReg := true.B
  }.otherwise {


package magmasi.components

import chisel3._
import chisel3.util._

class ReductionMux(val W: Int = 32, val NUM_IN: Int = 4, val SEL_IN: Int = 2) extends Module {
  require(NUM_IN >= 2 && NUM_IN % 2 == 0, "NUM_IN has to be >= 2 and divisible by 2")
  require(SEL_IN % 2 == 0, "SEL_IN has to be divisible by 2")

  val SEL_IN_LEFT_END = SEL_IN / 2 - 1 // 2/1 =2
  val SEL_IN_RIGHT_START = SEL_IN / 2

  val io = IO(new Bundle {
    val i_data = Input(Vec(NUM_IN, UInt(W.W)))
    val i_sel = Input(UInt(SEL_IN.W))
    val o_data = Output(Vec(2, UInt(W.W)))
  })
  val w_data_left = Wire(Vec(NUM_IN/2,UInt(W.W)))
  for (i <- 0 until NUM_IN/2){
    w_data_left(i) := io.i_data(i)
  }
  val w_data_right = Wire(Vec(NUM_IN/2,UInt(W.W)))
  for (i <- 0 until NUM_IN/2  ){
    w_data_right(i) := io.i_data(i+NUM_IN/2)
  }

  val w_sel_in_left = (io.i_sel(SEL_IN_LEFT_END, 0))//210
  val w_sel_in_right = (io.i_sel(SEL_IN - 1, SEL_IN_RIGHT_START))

  when ((w_sel_in_left).asUInt < 2.U) {

    io.o_data(0) := w_data_left(w_sel_in_left)
  }.otherwise {


package magmasi.components

import chisel3._
import chisel3.util._

class Regor(rows:Int,cols:Int,DATA_TYPE:Int) extends Module{
        implicit val config = MagmasiConfig()

    val io = IO(new Bundle{
        val mat1 =Input(Vec(rows, Vec(cols, UInt(DATA_TYPE.W))))
        val mat2 = Input(Vec(rows, Vec(cols, UInt(DATA_TYPE.W))))
        val compressedBitmap =  Output(Vec(rows, Vec(cols, UInt(DATA_TYPE.W))))
    })
    val matReg1 = RegInit(VecInit(Seq.fill(rows)(VecInit(Seq.fill(cols)(0.U(DATA_TYPE.W))))))
    io.compressedBitmap := matReg1

    val bitmap = Module(new Bitmap())
    bitmap.io.mat1 := io.mat1
    bitmap.io.mat2 := io.mat2

    val reg = RegInit(VecInit(Seq.fill(rows)(false.B)))
    for (i <- 0 until rows) {
        reg(i) := io.mat2(i).map(_ =/= 0.U).reduce(_ || _)
    }
    dontTouch(reg)

    val i = RegInit(0.U(log2Ceil(rows).W))
    val j = RegInit(0.U(log2Ceil(cols).W))
    dontTouch(i)
    dontTouch(j)

    when (i < cols.U) {
        when ((reg(j) === 0.U) && (bitmap.io.bitmap1(i)(j) === 1.U)){
            matReg1(i)(j) := 0.U//bitmap.io.bitmap1(i)(j)

        }.otherwise{
            matReg1(i)(j) := bitmap.io.bitmap1(i)(j)
        }


package magmasi.components

import chisel3._
import chisel3.util._

class SimpleAdder extends Module{
    val io = IO(new Bundle{
        val A = Input(UInt(32.W))
        val B = Input(UInt(32.W))
        val O = Output(UInt(32.W))
    })


package magmasi.components

import chisel3._
import chisel3.util._

class SimpleMultiplier extends Module{
    val io = IO(new Bundle{
        val A = Input(UInt(16.W))
        val B = Input(UInt(16.W))
        val O = Output(UInt(16.W))
    })


/* These Modules are translated into CHISEL from https://github.com/georgia-tech-synergy-lab/SIGMA/blob/master/vmod/benes.v */
package magmasi.components

import chisel3._

// input_switch module
class InputSwitch(val width: Int) extends Module {
  val io = IO(new Bundle {
    val y = Output(UInt(width.W))
    val z = Output(UInt(width.W))
    val in = Input(UInt(width.W))
  })
  io.y := io.in
  io.z := io.in
}

// output_switch module
class OutputSwitch(val width: Int) extends Module {
  val io = IO(new Bundle {
    val y = Output(UInt(width.W))
    val in0 = Input(UInt(width.W))
    val in1 = Input(UInt(width.W))
    val sel = Input(Bool())
  })

  val mux0 = Module(new BenesMux(width))
  mux0.io.a := io.in0
  mux0.io.b := io.in1
  mux0.io.sel := io.sel
  io.y := mux0.io.o
}

// switch module
class Switch(val width: Int) extends Module {
  val io = IO(new Bundle {
    val y = Output(UInt(width.W))
    val z = Output(UInt(width.W))
    val in0 = Input(UInt(width.W))

  })

  val mux0 = Module(new BenesMux(width))
  val mux1 = Module(new BenesMux(width))
  mux0.io.a := io.in0
  mux0.io.b := io.in1
  mux0.io.sel := io.sel0
  mux1.io.a := io.in0
  mux1.io.b := io.in1
  mux1.io.sel := io.sel1
  io.y := mux0.io.o
  io.z := mux1.io.o
}

// benes_mux module is assumed to be available in the same scope
class BenesMux(W: Int) extends Module {
  val io = IO(new Bundle {
    val o = Output(UInt(W.W))
    val a = Input(UInt(W.W))
    val b = Input(UInt(W.W))
    val sel = Input(Bool())
  })

  io.o := Mux(io.sel, io.b, io.a)
}

    val in1 = Input(UInt(width.W))
    val sel0 = Input(Bool())


package  magmasi.components

import chisel3._
import chisel3.util._

class benseBuffers(implicit val config: MagmasiConfig) extends Module{

    // benes routing calling here to reach it's destination into the buffer1 and buffer2
    val benes = Module(new Benes())
    benes.io.i_data_bus1 <> io.i_data_bus1
    benes.io.i_data_bus2 <> io.i_data_bus2
    benes.io.i_mux_bus <> io.i_mux_bus

    //which buffer have get reaches their destination they will both get multiplication over here
    val buffer = Module(new buffer_multiplication())
    buffer.io.buffer1 <> benes.io.o_dist_bus1
    buffer.io.buffer2 <> benes.io.o_dist_bus2
    io.out := buffer.io.out
}

    val io = IO(new Bundle{
        val i_data_bus2 = Input(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))
        val i_data_bus1  = Input(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))
        val i_mux_bus   = Input(Vec(config.NUM_PES, UInt((config.LEVELS-1).W)))


package  magmasi.components

import chisel3._
import  chisel3.util._

class buffer_multiplication(implicit val config: MagmasiConfig) extends Module{
    val io = IO(new Bundle{
        val buffer1 = Input(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))
        val buffer2 = Input(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))
        val out = Output(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))
    })

    def multiply(vec1: Vec[UInt], vec2: Vec[UInt], index: Int = 0): Vec[UInt] = {
        if (index >= config.NUM_PES) {
            VecInit(Seq.fill(config.NUM_PES)(0.U(config.DATA_TYPE.W)))

        } else {
            val elementMul = vec1(index) * vec2(index)
            val result = multiply(vec1, vec2, index + 1)
            result(index) := elementMul
            result


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class AdderSwitchTester extends AnyFreeSpec with ChiselScalatestTester {
    "Adder Switch Test" in {
        implicit val config = MagmasiConfig()
        test(new AdderSwitch(32,2,2,2)).withAnnotations(Seq(VerilatorBackendAnnotation)) { dut =>
            dut.io.i_data_bus(0).poke(2.U)
            dut.io.i_data_bus(1).poke(5.U)
            // dut.io.i_data_bus(4).poke(2.U)
            // dut.io.i_data_bus(5).poke(5.U)
            // dut.io.i_data_bus(6).poke(4.U)

            // dut.io.i_data_bus(7).poke(6.U)
            dut.io.i_valid.poke(1.B)
            dut.io.i_add_en.poke(1.U)
            dut.io.i_cmd.poke(5.U)


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class BitmapTester extends AnyFreeSpec with ChiselScalatestTester {
    "Bitmap Table" in {
        implicit val config = MagmasiConfig()
        test(new Bitmap()).withAnnotations  (Seq(VerilatorBackendAnnotation)) { dut =>
      val inputData = Seq(
        Seq(0, 2, 3, 0),
        Seq(5, 6, 7, 0),
        Seq(9, 10, 11, 0),
        Seq(13, 14, 15, 0)
      )

      // Assign test values to KM_matrix
      for (i <- 0 until 4) {

        for (j <- 0 until 4) {
          dut.io.mat1(i)(j).poke(inputData(i)(j).U)
          dut.io.mat2(i)(j).poke(inputData(i)(j).U)
        }


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

            dut.io.i_data_bus(1).poke(4.U)
            // dut.io.i_data_bus(2).poke(5.U)
            // dut.io.i_data_bus(3).poke(6.U)
            // dut.io.i_data_bus(4).poke(3.U)
            // dut.io.i_data_bus(5).poke(4.U)
            // dut.io.i_data_bus(6).poke(5.U)
            // dut.io.i_data_bus(7).poke(6.U)
            dut.io.i_valid.poke(1.B)
            dut.io.i_add_en.poke(1.U)
            dut.io.i_cmd.poke(5.U)
            dut.io.i_sel.poke(1.U)
            dut.clock.step(3)
            //dut.io.o_adder.expect(10.U)

        }
    }
}

class EdgeAdderSwitchTester extends AnyFreeSpec with ChiselScalatestTester {
    "Edge Adder Switch Test" in {
        implicit val config = MagmasiConfig()


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class fancontroltest extends AnyFreeSpec with ChiselScalatestTester {
    "Fancontrolcom Test" in {
        test(new fancontrol(32,32,5)){ c =>
            c.io.i_vn(0).poke("b11111".U)
            c.io.i_vn(1).poke("b11111".U)
            c.io.i_vn(2).poke("b11111".U)
            c.io.i_vn(3).poke("b11111".U)
            c.io.i_vn(4).poke("b11111".U)
            c.io.i_vn(5).poke("b11111".U)
            c.io.i_vn(6).poke("b11111".U)
            c.io.i_vn(7).poke("b11111".U)
            c.io.i_vn(8).poke("b11111".U)
            c.io.i_vn(9).poke("b11111".U)


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class FanNetcomTester extends AnyFreeSpec with ChiselScalatestTester {
    "FanNetcomplete Test" in {
        test(new FanNetworkcom(32,32)){ c =>

            c.io.i_valid.poke(1.U)
            c.io.i_data_bus(0).poke(1.U)
            c.io.i_data_bus(1).poke(1.U)
            c.io.i_data_bus(2).poke(1.U)
            c.io.i_data_bus(3).poke(1.U)
            c.io.i_data_bus(4).poke(1.U)
            c.io.i_data_bus(5).poke(1.U)
            c.io.i_data_bus(6).poke(1.U)
            c.io.i_data_bus(7).poke(1.U)
            c.io.i_data_bus(8).poke(1.U)
            c.io.i_data_bus(9).poke(1.U)
            c.io.i_data_bus(10).poke(1.U)
            c.io.i_data_bus(11).poke(1.U)
            c.io.i_data_bus(12).poke(1.U)
            c.io.i_data_bus(13).poke(1.U)
            c.io.i_data_bus(14).poke(1.U)
            c.io.i_data_bus(15).poke(1.U)
            c.io.i_data_bus(16).poke(1.U)
            c.io.i_data_bus(17).poke(1.U)
            c.io.i_data_bus(18).poke(1.U)
            c.io.i_data_bus(19).poke(1.U)
            c.io.i_data_bus(20).poke(1.U)
            c.io.i_data_bus(21).poke(1.U)
            c.io.i_data_bus(22).poke(1.U)
            c.io.i_data_bus(23).poke(1.U)
            c.io.i_data_bus(24).poke(1.U)
            c.io.i_data_bus(25).poke(1.U)
            c.io.i_data_bus(26).poke(1.U)
            c.io.i_data_bus(27).poke(1.U)
            c.io.i_data_bus(28).poke(1.U)
            c.io.i_data_bus(29).poke(1.U)
            c.io.i_data_bus(30).poke(1.U)
            c.io.i_data_bus(31).poke(1.U)
            //c.io.i_add_en_bus.poke(2147483647.U)


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class flexdpecom2test extends AnyFreeSpec with ChiselScalatestTester {
    "flexdpe Test" in {
        test(new flexdpecom2()){ c =>

           // for input data base 1
            c.io.i_data_bus(0).poke(1.U)
            c.io.i_data_bus(1).poke(1.U)
            c.io.i_data_bus(2).poke(1.U)
            c.io.i_data_bus(3).poke(1.U)
          // for input data bus 2
            c.io.i_data_bus2(0).poke(0.U)
            c.io.i_data_bus2(1).poke(2.U)
            c.io.i_data_bus2(2).poke(2.U)
            c.io.i_data_bus2(3).poke(0.U)

          // for i vn
            c.io.i_vn(3).poke("b11101".U)
            c.io.i_vn(4).poke("b11101".U)

            c.io.i_vn(7).poke("b10111".U)
            c.io.i_vn(8).poke("b10111".U)

            c.io.i_vn(11).poke("b01111".U)
            c.io.i_vn(12).poke("b01111".U)

            c.io.i_vn(15).poke("b11110".U)
            c.io.i_vn(16).poke("b11110".U)

            c.io.i_vn(23).poke("b11111".U)

            c.io.i_vn(13).poke("b11000".U)

          // for muxes
            c.io.i_mux_bus(576).poke(1.B)
            c.io.i_mux_bus(579).poke(1.B)

          // stationary and valid
            c.io.i_stationary.poke(1.B)
            c.io.i_data_valid.poke(1.B)

        c.clock.step(20)
           // for input data base 1
            c.io.i_data_bus(0).poke(1.U)
            c.io.i_data_bus(1).poke(1.U)
            c.io.i_data_bus(2).poke(1.U)
            c.io.i_data_bus(3).poke(1.U)
          // for input data bus 2
            c.io.i_data_bus2(0).poke(0.U)
            c.io.i_data_bus2(1).poke(3.U)
            c.io.i_data_bus2(2).poke(3.U)
            c.io.i_data_bus2(3).poke(0.U)

          // for i vn
            c.io.i_vn(3).poke("b11101".U)
            c.io.i_vn(4).poke("b11101".U)

            c.io.i_vn(7).poke("b10111".U)
            c.io.i_vn(8).poke("b10111".U)

            c.io.i_vn(11).poke("b01111".U)
            c.io.i_vn(12).poke("b01111".U)

            c.io.i_vn(15).poke("b11110".U)
            c.io.i_vn(16).poke("b11110".U)

            c.io.i_vn(23).poke("b11111".U)
            c.io.i_vn(24).poke("b11111".U)

            c.io.i_vn(18).poke("b11111".U)
            c.io.i_vn(13).poke("b11000".U)

          // for muxes
            c.io.i_mux_bus(576).poke(1.B)
            c.io.i_mux_bus(579).poke(1.B)

          // stationary and valid
            c.io.i_stationary.poke(1.B)
            c.io.i_data_valid.poke(1.B)

        c.clock.step(40)
    }
  }
}


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class MultSwitchTester extends AnyFreeSpec with ChiselScalatestTester {
    "multiplier Switch Test" in {
        test(new MultSwitch()){ dut =>
            dut.io.i_valid.poke(1.B)

            dut.clock.step(2)
        }
    }
}

            dut.io.i_data.poke(2.U)
            dut.io.i_stationary.poke(1.U)
            dut.io.o_valid.expect(0.B)


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class MybenesTester extends AnyFreeSpec with ChiselScalatestTester {
    "Benes Test" in {
        implicit val config = MagmasiConfig()
        test(new Benes()).withAnnotations(Seq(VerilatorBackendAnnotation)) { dut =>
          val inputData = Seq(1,2,2,4)
          for (i <- 0 until 4) {
            dut.io.i_data_bus2(i).poke(inputData(i).U)
            //dut.io.i_data_bus2(i).poke(inputData(i).U)
          }
          dut.io.i_mux_bus(0).poke(0.B)
          dut.io.i_mux_bus(1).poke(0.B)
          dut.io.i_mux_bus(2).poke(0.B)
          dut.io.i_mux_bus(3).poke(0.B)

          dut.clock.step(5)
        }
    }
}

val LEVELS   : Int = (2 * (math.log(4) / math.log(2))).toInt + 1


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class ReductionMuxTester extends AnyFreeSpec with ChiselScalatestTester {
    "Reduction Mux Test" in {
        implicit val config = MagmasiConfig()
        test(new ReductionMux).withAnnotations(Seq(VerilatorBackendAnnotation)) { dut =>
            dut.io.i_data(0).poke(3.U)

            dut.io.i_data(1).poke(4.U)
            dut.io.i_data(2).poke(5.U)
            dut.io.i_data(3).poke(6.U)
            dut.io.i_sel.poke(3.U)
            dut.clock.step(1)


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class RegorTester extends AnyFreeSpec with ChiselScalatestTester {
    "Temporary Register step 3 Table" in {
        implicit val config = MagmasiConfig()
        test(new Regor(4,4,32)).withAnnotations  (Seq(VerilatorBackendAnnotation)) { dut =>
      val inputData = Seq(
        Seq(0, 2, 3, 1),
        Seq(5, 6, 7, 0),
        Seq(9, 10, 11, 0),
        Seq(1, 2, 3, 0)

        Seq(0, 0, 0, 0)
      )

      // Assign test values to KM_matrix
      for (i <- 0 until 4) {
        for (j <- 0 until 4) {
          dut.io.mat1(i)(j).poke(inputData(i)(j).U)
          dut.io.mat2(i)(j).poke(inputData1(i)(j).U)
        }
    }
          dut.clock.step(65)
        }
    }
}

      )
            val inputData1 = Seq(
        Seq(1, 2, 0, 0),


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class benesBuffersTester extends AnyFreeSpec with ChiselScalatestTester {
    "My Benes Top test" in {
        implicit val config = MagmasiConfig()
        test(new benseBuffers()).withAnnotations(Seq(VerilatorBackendAnnotation)) { dut =>

            val inputData = Seq(1, 2, 3, 4, 5, 6, 7, 8)

      // Connect input data to the DUT

        dut.io.i_mux_bus(0).poke(3.U)
        dut.io.i_mux_bus(1).poke(3.U)
        dut.io.i_mux_bus(2).poke(3.U)
        dut.io.i_mux_bus(3).poke(3.U)
        val LEVELS   : Int = (2 * (math.log(4) / math.log(2))).toInt + 1

      dut.clock.step(3)

    }
  }
}

      for (i <- 0 until 4) {
        dut.io.i_data_bus1(i).poke(inputData(i).U)
      }


package magmasi.components

import chisel3._
import chisel3.util._
import chiseltest._
import org.scalatest.freespec.AnyFreeSpec

class bufferMultTester extends AnyFreeSpec with ChiselScalatestTester {
    "buffer multiplication Test" in {
        val DATA_TYPE = 32
        val NUM_IN = 4
        implicit val config = MagmasiConfig()
        test(new buffer_multiplication()).withAnnotations(Seq(VerilatorBackendAnnotation)) { dut =>

        val inputData = Seq(1, 2, 3, 4, 5, 6, 7, 8)
        for (i <- 0 until NUM_IN) {
            dut.io.buffer1(i).poke(inputData(i).U)
            dut.io.buffer2(i).poke(inputData(i).U)


package alu

import chisel3._
import chisel3.util._
import statecode.CoreState
import statecode.AluOpCode

class Alu extends Module {
  val io = IO(new Bundle {
    val enable = Input(Bool())

    val core_state = Input(CoreState())

    val decoded_alu_op = new Bundle {
      val arithmetic_mux = Input(AluOpCode()) // Select arithmetic operation
      val output_mux     = Input(Bool())      // Select operation in ALU
    }

    val reg_in = new Bundle {
      val rs = Input(UInt(8.W))
      val rt = Input(UInt(8.W))
    }

    val alu_out = Output(UInt(8.W))
  })

  val alu_out_reg = RegInit(0.U(8.W))

  when(io.enable) {
    when(io.core_state === CoreState.EXECUTE) {

            alu_out_reg := io.reg_in.rs + io.reg_in.rt
          }
          is(AluOpCode.SUB) {
            alu_out_reg := io.reg_in.rs - io.reg_in.rt
          }
          is(AluOpCode.MUL) {
            alu_out_reg := io.reg_in.rs * io.reg_in.rt
          }
          is(AluOpCode.DIV) {
            alu_out_reg := io.reg_in.rs / io.reg_in.rt
          }
        }
      }
    }
  }

  when(reset.asBool) {
    io.alu_out := 0.U
  }.otherwise {
    io.alu_out := alu_out_reg
  }
}

      when(io.decoded_alu_op.output_mux) {
        // Set values to compare with NZP register in alu_out[2:0]
        val gt = io.reg_in.rs > io.reg_in.rt
        val eq = io.reg_in.rs === io.reg_in.rt
        val lt = io.reg_in.rs < io.reg_in.rt


package controller

import chisel3._
import chisel3.util._

import statecode.ControlState

// MEMORY CONTROLLER
// > Receives memory requests from all cores
// > Throttles requests based on limited external memory bandwidth
// > Waits for responses from external memory and distributes them back to cores
class Controller(
  AddrBits:     Int = 8,
  DataBits:     Int = 16,
  NumConsumers: Int = 4,
  NumChannels:  Int = 1,
  WriteEnable:  Boolean = true)
    extends Module {
  // TODO: Use decoupledIO interface for consumer and memory
  val io = IO(new Bundle {
    // Consumer Interface (Fetchers / LSUs)
    val consumer_read_addr_receiver = Vec(NumConsumers, Flipped(new DecoupledIO(UInt(AddrBits.W))))
    val consumer_read_data          = Output(Vec(NumConsumers, UInt(DataBits.W)))

    val consumer_write_receiver =
      if (WriteEnable)
        Some(
          Vec(
            NumConsumers,
            Flipped(DecoupledIO(new Bundle {
              val address = UInt(AddrBits.W)
              val data    = UInt(DataBits.W)
            }))
          )
        )
      else None

    // Memory Interface (Data / Program)
    val mem_read_sender = Vec(NumChannels, new DecoupledIO(UInt(AddrBits.W)))
    val mem_read_data   = Input(Vec(NumChannels, UInt(DataBits.W)))

    val mem_write_sender =
      if (WriteEnable)
        Some(
          Vec(
            NumChannels,
            DecoupledIO(new Bundle {
              val address = UInt(AddrBits.W)
              val data    = UInt(DataBits.W)
            })
          )
        )
      else None
  })

  val mem_read_valid   = RegInit(VecInit(Seq.fill(NumChannels)(false.B)))
  val mem_read_address = RegInit(VecInit(Seq.fill(NumChannels)(0.U(AddrBits.W))))

  val mem_write_valid   = RegInit(VecInit(Seq.fill(NumChannels)(false.B)))
  val mem_write_address = RegInit(VecInit(Seq.fill(NumChannels)(0.U(AddrBits.W))))
  val mem_write_data    = RegInit(VecInit(Seq.fill(NumChannels)(0.U(DataBits.W))))

  val consumer_read_ready  = RegInit(VecInit(Seq.fill(NumConsumers)(false.B)))
  val consumer_read_data   = RegInit(VecInit(Seq.fill(NumConsumers)(0.U(DataBits.W))))
  val consumer_write_ready = RegInit(VecInit(Seq.fill(NumConsumers)(false.B)))

  val current_consumer = RegInit(VecInit(Seq.fill(NumChannels)(0.U(log2Ceil(NumConsumers).W))))
  val controller_state = RegInit(VecInit(Seq.fill(NumChannels)(ControlState.IDLE)))

  val channel_serving_consumer = RegInit(VecInit(Seq.fill(NumConsumers)(false.B)))

  // // Add debug printing
  // printf("#### Module Internal State ####\n")
  // printf(cf"#controller_state: ${controller_state}\n")
  // printf(cf"#mem_read_valid: ${mem_read_valid}\n")
  // printf(cf"#mem_read_address: ${mem_read_address}\n")
  // printf(cf"#mem_write_valid: ${mem_write_valid}\n")
  // printf(cf"#mem_write_address: ${mem_write_address}\n")
  // printf(cf"#mem_write_data: ${mem_write_data}\n")
  // printf(cf"#consumer_read_ready: ${consumer_read_ready}\n")
  // printf(cf"#consumer_read_data: ${consumer_read_data}\n")
  // printf(cf"#consumer_write_ready: ${consumer_write_ready}\n")
  // printf(cf"#channel_serving_consumer: ${channel_serving_consumer}\n")
  // printf(cf"#current_consumer: ${current_consumer}\n")

  when(!reset.asBool) {
    // For each channel, we handle processing concurrently
    for (i <- 0 until NumChannels) {
      switch(controller_state(i)) {
        is(ControlState.IDLE) {
          // While this channel is idle, cycle through consumers looking for one with a pending request
          // Once we find a pending request, pick it up with this channel and stop looking for requests
          val read_signals  = Wire(Vec(NumConsumers, Bool()))
          // val write_signals = Wire(Vec(NumConsumers, Bool()))
          val write_signals = VecInit(Seq.fill(NumConsumers)(false.B))
          for (j <- 0 until NumConsumers) {
            read_signals(j) := io.consumer_read_addr_receiver(j).valid && !channel_serving_consumer(j)
            // write_signals(j) := io.consumer_write_receiver(j).valid && !channel_serving_consumer(j)
            io.consumer_write_receiver.map(write_receiver =>
              write_signals(j) := write_receiver(j).valid && !channel_serving_consumer(j)
            )
          }

          val first_read_idx  = PriorityEncoder(read_signals)
          val first_write_idx = PriorityEncoder(write_signals)
          // printf("#first_read_idx: %b, first_write_idx: %b\n", first_read_idx, first_write_idx)

          when(read_signals.asUInt > 0.U && first_read_idx <= first_write_idx) {
            val read_address = io.consumer_read_addr_receiver(first_read_idx).bits

            channel_serving_consumer(first_read_idx) := true.B
            current_consumer(i)                      := first_read_idx

            mem_read_valid(i)   := true.B
            mem_read_address(i) := read_address

            controller_state(i) := ControlState.READ_WAITING
          }.elsewhen(write_signals.asUInt > 0.U) {
            // val write_address = io.consumer_write_receiver(first_write_idx).bits.address
            // val write_data    = io.consumer_write_receiver(first_write_idx).bits.data
            val write_address = io.consumer_write_receiver.map(_.apply(first_write_idx).bits.address).getOrElse(0.U)
            val write_data    = io.consumer_write_receiver.map(_.apply(first_write_idx).bits.data).getOrElse(0.U)

            channel_serving_consumer(first_write_idx) := true.B

          }
        }
        is(ControlState.READ_WAITING) {
          // Wait for response from memory for pending read request
          when(io.mem_read_sender(i).ready) {
            mem_read_valid(i)                        := false.B
            consumer_read_ready(current_consumer(i)) := true.B
            consumer_read_data(current_consumer(i))  := io.mem_read_data(i)
            controller_state(i)                      := ControlState.READ_RELAYING
          }
        }
        is(ControlState.WRITE_WAITING) {
          // Wait for response from memory for pending write request
          val write_ready = io.mem_write_sender.map(_.apply(i).ready).getOrElse(false.B)
          // when(io.mem_write_sender(i).ready) {
          when(write_ready) {
            mem_write_valid(i)                        := false.B
            consumer_write_ready(current_consumer(i)) := true.B
            controller_state(i)                       := ControlState.WRITE_RELAYING
          }
        }
        is(ControlState.READ_RELAYING) {
          when(!io.consumer_read_addr_receiver(current_consumer(i)).valid) {
            channel_serving_consumer(current_consumer(i)) := false.B
            consumer_read_ready(current_consumer(i))      := false.B
            controller_state(i)                           := ControlState.IDLE
          }
        }
        is(ControlState.WRITE_RELAYING) {
          val write_valid = io.consumer_write_receiver.map(!_.apply(current_consumer(i)).valid).getOrElse(false.B)
          // when(!io.consumer_write_receiver(current_consumer(i)).valid) {
          when(write_valid) {
            channel_serving_consumer(current_consumer(i)) := false.B
            consumer_write_ready(current_consumer(i))     := false.B
            controller_state(i)                           := ControlState.IDLE
          }
        }
      }
    }
  }

  for (i <- 0 until NumConsumers) {
    io.consumer_read_addr_receiver(i).ready := consumer_read_ready(i)
    io.consumer_read_data(i)                := consumer_read_data(i)
    // io.consumer_write_receiver(i).ready     := consumer_write_ready(i)
    io.consumer_write_receiver.map(_.apply(i).ready := consumer_write_ready(i))
  }

  for (i <- 0 until NumChannels) {
    io.mem_read_sender(i).valid := mem_read_valid(i)
    io.mem_read_sender(i).bits  := mem_read_address(i)
    // io.mem_write_sender(i).valid        := mem_write_valid(i)
    // io.mem_write_sender(i).bits.address := mem_write_address(i)
    // io.mem_write_sender(i).bits.data    := mem_write_data(i)
    io.mem_write_sender.map(sender => {
      sender(i).valid        := mem_write_valid(i)
      sender(i).bits.address := mem_write_address(i)
      sender(i).bits.data    := mem_write_data(i)
    })
  }
}

            current_consumer(i)                       := first_write_idx

            mem_write_valid(i)   := true.B
            mem_write_address(i) := write_address


package core

import chisel3._
import chisel3.util._
import statecode.CoreState

import fetcher.Fetcher
import decoder.Decoder
import scheduler.Scheduler
import registers.RegisterFiles
import alu.Alu
import lsu.MemLoadStoreUnit
import pc.ProgramCounter

// TODO: Change each submodules' IO Bundle,
//       and then connect them in a more Chisel way

// COMPUTE CORE
// > Handles processing 1 block at a time
// > The core also has it's own scheduler to manage control flow
// > Each core contains 1 fetcher & decoder, and register files, ALUs, LSUs, PC for each thread
class Core(
  DataMemAddrBits:    Int = 8,
  DataMemDataBits:    Int = 8,
  ProgramMemAddrBits: Int = 8,
  ProgramMemDataBits: Int = 16,
  ThreadsPerBlock:    Int = 4)
    extends Module {
  val io = IO(new Bundle {
    // Kernel Execution
    val start = Input(Bool())
    val done  = Output(Bool())

    // Block Metadata
    val block_id     = Input(UInt(8.W))
    val thread_count = Input(UInt(log2Ceil(ThreadsPerBlock).W))

    // Program Memory
    val program_mem_read_address_sender = new DecoupledIO(UInt(ProgramMemAddrBits.W))
    val program_mem_read_data           = Input(UInt(ProgramMemDataBits.W))

    // Data Memory

    val data_mem_read_data           = Input(Vec(ThreadsPerBlock, UInt(DataMemDataBits.W)))
    val data_mem_read_address_sender = Vec(ThreadsPerBlock, new DecoupledIO(UInt(DataMemAddrBits.W)))
    val data_mem_write_sender        = Vec(
      ThreadsPerBlock,
      new DecoupledIO(new Bundle {
        val address = UInt(DataMemAddrBits.W)
        val data    = UInt(DataMemDataBits.W)
      })
    )

    // debug outputs
    val core_state = Output(CoreState())
    val current_pc = Output(UInt(8.W))
  })

  val fetcher   = Module(new Fetcher(ProgramMemAddrBits, ProgramMemDataBits))
  val decoder   = Module(new Decoder())
  val scheduler = Module(new Scheduler(ThreadsPerBlock))

  // val core_state = RegNext(scheduler.io.core_state)
  // val current_pc = RegNext(scheduler.io.current_pc)
  val core_state = Wire(CoreState())
  val current_pc = Wire(UInt(8.W))
  core_state := scheduler.io.core_state
  current_pc := scheduler.io.current_pc

  // printf(cf"--Core State: $core_state, current PC: $current_pc, Fetcher State: ${fetcher.io.fetcher_state}\n")

  // Fetcher inputs connections (3/3)
  fetcher.io.core_state    := core_state
  fetcher.io.current_pc    := current_pc
  fetcher.io.mem_read_data := io.program_mem_read_data

  // Decoder inputs connections (2/2)
  decoder.io.core_state  := scheduler.io.core_state
  decoder.io.instruction := fetcher.io.instruction

  // Scheduler inputs connections (5/7)
  scheduler.io.start         := io.start
  scheduler.io.mem_rw_enable := decoder.io.mem_rw_enable
  scheduler.io.decoded_ret   := decoder.io.decoded_ret
  scheduler.io.fetcher_state := fetcher.io.fetcher_state

  // TODO: just use a for loop
  val compute_units = Seq
    .tabulate(ThreadsPerBlock)(i => {
      val alu     = Module(new Alu())
      val lsu     = Module(new MemLoadStoreUnit())
      val regfile = Module(new RegisterFiles(ThreadsPerBlock, i, DataMemDataBits))
      val pc      = Module(new ProgramCounter(DataMemDataBits, DataMemAddrBits))

      val enable = (i.U < io.thread_count)

      // alu inputs connections (6/6)
      alu.io.enable         := enable
      alu.io.core_state     := scheduler.io.core_state
      alu.io.reg_in         := regfile.io.reg_out
      alu.io.decoded_alu_op := decoder.io.decoded_alu_op

      // lsu inputs connections (7/7)

      regfile.io.enable                   := enable
      regfile.io.block_id                 := io.block_id
      regfile.io.core_state               := scheduler.io.core_state
      regfile.io.decoded_reg_address      := decoder.io.decoded_reg_address
      regfile.io.decoded_reg_write_enable := decoder.io.decoded_reg_write_enable
      regfile.io.decoded_reg_input_mux    := decoder.io.decoded_reg_input_mux
      regfile.io.decoded_immediate        := decoder.io.decoded_immediate
      regfile.io.alu_out                  := alu.io.alu_out
      regfile.io.lsu_out                  := lsu.io.lsu_out

      // pc inputs connections (8/8)
      pc.io.enable                   := enable
      pc.io.decoded_nzp              := decoder.io.decoded_nzp
      pc.io.decoded_immediate        := decoder.io.decoded_immediate
      pc.io.decoded_nzp_write_enable := decoder.io.decoded_nzp_write_enable
      pc.io.decoded_pc_mux           := decoder.io.decoded_pc_mux
      pc.io.alu_out                  := alu.io.alu_out
      pc.io.current_pc               := scheduler.io.current_pc
      pc.io.core_state               := scheduler.io.core_state

      // Connect to scheduler input (7/7)
      scheduler.io.lsu_state(i) := lsu.io.lsu_state
      scheduler.io.next_pc(i)   := pc.io.next_pc

      // Connect to core module outputs (5/8)
      io.data_mem_read_address_sender(i) <> lsu.io.mem_read_address_sender
      io.data_mem_write_sender(i) <> lsu.io.mem_write_sender

      (alu, lsu, regfile, pc)
    })

  // Connect to module outputs (8/8)
  io.done := scheduler.io.done
  io.program_mem_read_address_sender <> fetcher.io.mem_read_address_sender

  // debug
  io.core_state := core_state
  io.current_pc := current_pc
}

      lsu.io.enable        := enable
      lsu.io.core_state    := scheduler.io.core_state
      lsu.io.mem_rw_enable := decoder.io.mem_rw_enable
      lsu.io.mem_read_data := io.data_mem_read_data(i)


package decoder

import chisel3._
import chisel3.util._

import statecode.CoreState
import statecode.AluOpCode
import statecode.RegInputOp
import statecode.DecoderState

class Decoder extends Module {
  val io = IO(new Bundle {
    val core_state  = Input(CoreState())
    val instruction = Input(UInt(16.W))

    // Instruction signals
    val decoded_reg_address = new Bundle {
      val rd = Output(UInt(4.W))
      val rs = Output(UInt(4.W))
      val rt = Output(UInt(4.W))
    }
    val decoded_nzp         = Output(UInt(3.W))
    val decoded_immediate   = Output(UInt(8.W))

    // Control signals
    val mem_rw_enable            = new Bundle {
      val read_enable  = Output(Bool()) // Enable reading from memory
      val write_enable = Output(Bool()) // Enable writing to memory
    }
    val decoded_reg_write_enable = Output(Bool())       // Enable writing to a register
    val decoded_nzp_write_enable = Output(Bool())       // Enable writing to NZP register
    val decoded_reg_input_mux    = Output(RegInputOp()) // Select input to register

    val decoded_alu_op = new Bundle {
      val arithmetic_mux = Output(AluOpCode()) // Select arithmetic operation
      val output_mux     = Output(Bool())      // Select operation in ALU
    }
    val decoded_pc_mux = Output(Bool()) // Select source of next PC

    // Return (finished executing thread)
    val decoded_ret = Output(Bool())
  })

  val decoded_rd_address         = RegInit(0.U(4.W))
  val decoded_rs_address         = RegInit(0.U(4.W))
  val decoded_rt_address         = RegInit(0.U(4.W))
  val decoded_nzp                = RegInit(0.U(3.W))
  val decoded_immediate          = RegInit(0.U(8.W))
  val decoded_reg_write_enable   = RegInit(false.B)
  val decoded_mem_read_enable    = RegInit(false.B)
  val decoded_mem_write_enable   = RegInit(false.B)
  val decoded_nzp_write_enable   = RegInit(false.B)
  val decoded_reg_input_mux      = RegInit(RegInputOp.ARITHMETIC)
  val decoded_alu_arithmetic_mux = RegInit(AluOpCode.ADD)
  val decoded_alu_output_mux     = RegInit(false.B)
  val decoded_pc_mux             = RegInit(false.B)
  val decoded_ret                = RegInit(false.B)

  when(!reset.asBool) {
    when(io.core_state === CoreState.DECODE) {
      // Get instruction signals from instruction every time
      decoded_rd_address := io.instruction(11, 8)
      decoded_rs_address := io.instruction(7, 4)
      decoded_rt_address := io.instruction(3, 0)
      decoded_immediate  := io.instruction(7, 0)
      decoded_nzp        := io.instruction(11, 9)

      // Control signals reset on every decode and set conditionally by instruction
      decoded_reg_write_enable   := false.B
      decoded_mem_read_enable    := false.B
      decoded_mem_write_enable   := false.B
      decoded_nzp_write_enable   := false.B
      decoded_reg_input_mux      := RegInputOp.ARITHMETIC
      decoded_alu_arithmetic_mux := AluOpCode.ADD
      decoded_alu_output_mux     := false.B
      decoded_pc_mux             := false.B
      decoded_ret                := false.B

      // Set the control signals for each instruction
      val (instruction_opcode, _) = DecoderState.safe(io.instruction(15, 12))
      switch(instruction_opcode) {
        is(DecoderState.NOP) {
          // no-op
        }
        is(DecoderState.BRnzp) {
          decoded_pc_mux := true.B
        }
        is(DecoderState.CMP) {
          decoded_alu_output_mux   := true.B
          decoded_nzp_write_enable := true.B
        }
        is(DecoderState.ADD) {
          decoded_reg_write_enable   := true.B
          decoded_reg_input_mux      := RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux := AluOpCode.ADD
        }
        is(DecoderState.SUB) {
          decoded_reg_write_enable   := true.B
          decoded_reg_input_mux      := RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux := AluOpCode.SUB
        }
        is(DecoderState.MUL) {
          decoded_reg_write_enable   := true.B
          decoded_reg_input_mux      := RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux := AluOpCode.MUL
        }
        is(DecoderState.DIV) {

          decoded_mem_read_enable  := true.B
        }
        is(DecoderState.STR) {
          decoded_mem_write_enable := true.B
        }
        is(DecoderState.CONST) {
          decoded_reg_write_enable := true.B
          decoded_reg_input_mux    := RegInputOp.CONSTANT
        }
        is(DecoderState.RET) {
          decoded_ret := true.B
        }
      }
    }
  }

  io.decoded_reg_address.rd        := decoded_rd_address
  io.decoded_reg_address.rs        := decoded_rs_address
  io.decoded_reg_address.rt        := decoded_rt_address
  io.decoded_nzp                   := decoded_nzp
  io.decoded_immediate             := decoded_immediate
  io.decoded_reg_write_enable      := decoded_reg_write_enable
  io.mem_rw_enable.read_enable     := decoded_mem_read_enable
  io.mem_rw_enable.write_enable    := decoded_mem_write_enable
  io.decoded_nzp_write_enable      := decoded_nzp_write_enable
  io.decoded_reg_input_mux         := decoded_reg_input_mux
  io.decoded_alu_op.arithmetic_mux := decoded_alu_arithmetic_mux
  io.decoded_alu_op.output_mux     := decoded_alu_output_mux
  io.decoded_pc_mux                := decoded_pc_mux
  io.decoded_ret                   := decoded_ret
}

          decoded_reg_write_enable   := true.B
          decoded_reg_input_mux      := RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux := AluOpCode.DIV
        }


package gcd

import chisel3._
import chisel3.util.Decoupled

class GcdInputBundle(val w: Int) extends Bundle {
  val value1 = UInt(w.W)
  val value2 = UInt(w.W)
}

class GcdOutputBundle(val w: Int) extends Bundle {
  val value1 = UInt(w.W)
  val value2 = UInt(w.W)
  val gcd    = UInt(w.W)
}

/** Compute Gcd using subtraction method. Subtracts the smaller from the larger until register y is zero. value input
  * register x is then the Gcd. Unless first input is zero then the Gcd is y. Can handle stalls on the producer or
  * consumer side
  */
class DecoupledGcd(width: Int) extends Module {
  val input  = IO(Flipped(Decoupled(new GcdInputBundle(width))))
  val output = IO(Decoupled(new GcdOutputBundle(width)))

  val xInitial    = Reg(UInt())
  val yInitial    = Reg(UInt())
  val x           = Reg(UInt())
  val y           = Reg(UInt())
  val busy        = RegInit(false.B)
  val resultValid = RegInit(false.B)

  input.ready  := !busy
  output.valid := resultValid
  output.bits  := DontCare

  when(busy) {
    when(x > y) {

      when(x === 0.U) {
        output.bits.gcd := y
      }.otherwise {
        output.bits.gcd := x
      }

      output.bits.value1 := xInitial
      output.bits.value2 := yInitial
      resultValid        := true.B

      when(output.ready && resultValid) {
        busy        := false.B
        resultValid := false.B
      }
    }
  }.otherwise {
    when(input.valid) {
      val bundle = input.deq()
      x        := bundle.value1
      y        := bundle.value2
      xInitial := bundle.value1
      yInitial := bundle.value2
      busy     := true.B
    }
  }
}

      x := x - y
    }.otherwise {
      y := y - x


package dcr

import chisel3._
import chisel3.util._

// > Used to configure high-level settings
// > In this minimal example, the DCR is used to configure the number of threads to run for the kernel
class DeviceControlRegister extends Module {
  val io = IO(new Bundle {
    val device_control_write_enable = Input(Bool())
    val device_control_data         = Input(UInt(8.W))
    val thread_count                = Output(UInt(8.W))
  })

  // Store device control data in dedicated register
  val device_control_register = RegInit(0.U(8.W))

  when(!reset.asBool && io.device_control_write_enable) {
    device_control_register := io.device_control_data
  }


package dispatch

import chisel3._
import chisel3.util._
import org.apache.commons.lang3.ThreadUtils.ThreadPredicate

class Dispatch(NumCores: Int = 2, ThreadsPerCore: Int = 4) extends Module {
  val ThreadCountWidth = log2Ceil(ThreadsPerCore) + 1
  val io               = IO(new Bundle {
    val start = Input(Bool())

    // Kernel Metadata
    val thread_count = Input(UInt(8.W))

    // Core States
    val core_done         = Input(Vec(NumCores, Bool()))
    val core_start        = Output(Vec(NumCores, Bool()))
    val core_reset        = Output(Vec(NumCores, Bool()))
    val core_block_id     = Output(Vec(NumCores, UInt(8.W)))
    val core_thread_count = Output(Vec(NumCores, UInt(ThreadCountWidth.W)))

    // Kernel Execution
    val done = Output(Bool())
  })

  printf(cf"Dispatch: start = ${io.start}, thread_count = ${io.thread_count}, core_done = ${io.core_done}\n")

  // Calculate the total number of blocks based on total threads & threads per block
  val total_blocks = (io.thread_count + ThreadsPerCore.U - 1.U) / ThreadsPerCore.U

  val done              = RegInit(false.B)
  val blocks_dispatched = RegInit(0.U(8.W))
  val blocks_done       = RegInit(0.U(8.W))
  val start_execution   = RegInit(false.B)
  val core_start        = RegInit(VecInit(Seq.fill(NumCores)(false.B)))
  val core_reset        = RegInit(VecInit(Seq.fill(NumCores)(true.B)))
  val core_block_id     = RegInit(VecInit(Seq.fill(NumCores)(0.U(8.W))))
  val core_thread_count = RegInit(VecInit(Seq.fill(NumCores)(ThreadsPerCore.U(ThreadCountWidth.W))))

  printf(

  )

  // printf(cf"--total_blocks = $total_blocks, thread_count = ${io.thread_count}, blocks_done = $blocks_done\n")

  // Keep track of how many blocks have been processed
  when(!reset.asBool && io.start) {
    when(!start_execution) {
      start_execution := true.B
      core_reset.foreach(_ := true.B)
    }

    // If the last block has finished processing, mark this kernel as done executing
    when(blocks_done === total_blocks) {
      done := true.B
    }

    for (i <- 0 until NumCores) {
      when(core_reset(i)) {
        core_reset(i) := false.B

        // If this core was just reset, check if there are more blocks to be dispatched
        when(blocks_dispatched < total_blocks) {
          // printf(cf"--i = $i, blocks_dispatched = $blocks_dispatched, total_blocks = $total_blocks\n")
          core_start(i)        := true.B
          core_block_id(i)     := blocks_dispatched
          core_thread_count(i) := Mux(
            blocks_dispatched === total_blocks - 1.U,
            io.thread_count - (blocks_dispatched * ThreadsPerCore.U),
            ThreadsPerCore.U
          )
          blocks_dispatched    := blocks_dispatched + 1.U
        }
      }

      // printf(cf"--core_start(i) = ${core_start(i)}, io.core_done(i) = ${io.core_done(i)}\n")
      when(core_start(i) && io.core_done(i)) {
        // If a core just finished executing it's current block, reset it
        core_reset(i) := true.B
        core_start(i) := false.B
        blocks_done   := blocks_done + 1.U
      }
    }
  }

  io.core_start        := core_start
  io.core_reset        := core_reset
  io.core_block_id     := core_block_id
  io.core_thread_count := core_thread_count
  io.done              := done
}

cf"Dispatch outputs: core_start=${core_start}, core_reset=${core_reset}, core_block_id=${core_block_id}, core_thread_count=${core_thread_count}, done=${done}\n"


object Elaborate extends App {
  val firtoolOptions = Array(
    "--lowering-options=" + List(
      // make yosys happy
      // see https://github.com/llvm/circt/blob/main/docs/VerilogGeneration.md
      "disallowLocalVariables",
      "disallowPackedArrays",
      "locationInfoStyle=wrapInAtSquareBracket"
    ).reduce(_ + "," + _)
  )
  circt.stage.ChiselStage.emitSystemVerilogFile(new gcd.GCD(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new alu.Alu(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new pc.ProgramCounter, args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new lsu.MemLoadStoreUnit(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new registers.RegisterFiles(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new scheduler.Scheduler(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new fetcher.Fetcher(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new decoder.Decoder(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new core.Core(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new dispatch.Dispatch(), args, firtoolOptions)

  circt.stage.ChiselStage.emitSystemVerilogFile(new dcr.DeviceControlRegister(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new controller.Controller(), args, firtoolOptions)
  circt.stage.ChiselStage.emitSystemVerilogFile(new gpu.Gpu(), args, firtoolOptions)


package fetcher

import chisel3._
import chisel3.util._

import statecode.CoreState
import statecode.FetcherState

class Fetcher(ProgramMemAddrBits: Int = 8, ProgramMemDataBits: Int = 16) extends Module {
  val io = IO(new Bundle {
    // Execution State
    val core_state = Input(CoreState())
    val current_pc = Input(UInt(8.W))

    // Program Memory
    val mem_read_address_sender = new DecoupledIO(UInt(ProgramMemAddrBits.W))
    val mem_read_data           = Input(UInt(ProgramMemDataBits.W))

    val fetcher_state = Output(FetcherState())
    val instruction   = Output(UInt(ProgramMemDataBits.W))
  })

  val fetcher_state    = RegInit(FetcherState.IDLE)
  val mem_read_valid   = RegInit(false.B)
  val mem_read_address = RegInit(0.U(ProgramMemAddrBits.W))
  val instruction      = RegInit(0.U(ProgramMemDataBits.W))

  when(!reset.asBool) {
    switch(fetcher_state) {
      is(FetcherState.IDLE) {
        when(io.core_state === CoreState.FETCH) {
          fetcher_state    := FetcherState.FETCHING
          mem_read_valid   := true.B
          mem_read_address := io.current_pc
        }
      }
      is(FetcherState.FETCHING) {
        when(io.mem_read_address_sender.ready) {
          fetcher_state  := FetcherState.FETCHED
          instruction    := io.mem_read_data
          mem_read_valid := false.B
        }
      }
      is(FetcherState.FETCHED) {

  }

  io.fetcher_state                 := fetcher_state
  io.mem_read_address_sender.valid := mem_read_valid
  io.mem_read_address_sender.bits  := mem_read_address
  io.instruction                   := instruction
}

        when(io.core_state === CoreState.DECODE) {
          fetcher_state := FetcherState.IDLE
        }


package gcd

import chisel3._

/** Compute GCD using subtraction method. Subtracts the smaller from the larger until register y is zero. value in
  * register x is then the GCD
  */
class GCD extends Module {
  val io = IO(new Bundle {
    val value1        = Input(UInt(16.W))
    val value2        = Input(UInt(16.W))
    val loadingValues = Input(Bool())
    val outputGCD     = Output(UInt(16.W))
    val outputValid   = Output(Bool())
  })

  val x = Reg(UInt())
  val y = Reg(UInt())

  when(io.loadingValues) {
    x := io.value1
    y := io.value2
  }

  io.outputGCD   := x
  io.outputValid := y === 0.U
}

when(x > y) { x := x - y }.otherwise { y := y - x }


package lsu

import chisel3._
import chisel3.util._
import statecode.CoreState
import statecode.LSUState

class MemLoadStoreUnit extends Module {
  val io = IO(new Bundle {
    val enable = Input(Bool())

    val core_state = Input(CoreState())

    // Decoder outputs
    val mem_rw_enable = new Bundle {
      val read_enable  = Input(Bool()) // Enable reading from memory
      val write_enable = Input(Bool()) // Enable writing to memory
    }

    // Registers
    val reg_in = new Bundle {
      val rs = Input(UInt(8.W))
      val rt = Input(UInt(8.W))
    }

    val mem_read_data           = Input(UInt(8.W))
    val mem_read_address_sender = new DecoupledIO(UInt(8.W))

    val mem_write_sender = DecoupledIO(new Bundle {
      val address = UInt(8.W)
      val data    = UInt(8.W)
    })

    // LSU Outputs
    val lsu_state = Output(LSUState())
    val lsu_out   = Output(UInt(8.W))
  })

  val lsu_state         = RegInit(LSUState.IDLE)
  val lsu_out           = RegInit(0.U(8.W))
  val mem_read_valid    = RegInit(false.B)
  val mem_read_address  = RegInit(0.U(8.W))
  val mem_write_valid   = RegInit(false.B)
  val mem_write_address = RegInit(0.U(8.W))
  val mem_write_data    = RegInit(0.U(8.W))

  when(io.enable) {
    // If memory read enable is triggered (LDR instruction)
    when(io.mem_rw_enable.read_enable) {
      switch(io.lsu_state) {
        is(LSUState.IDLE) {
          when(io.core_state === CoreState.REQUEST) {
            lsu_state := LSUState.REQUESTING
          }
        }
        is(LSUState.REQUESTING) {
          mem_read_valid   := true.B
          mem_read_address := io.reg_in.rs
          lsu_state        := LSUState.WAITING
        }
        is(LSUState.WAITING) {
          when(io.mem_read_address_sender.ready) {
            mem_read_valid := false.B
            lsu_out        := io.mem_read_data
            lsu_state      := LSUState.DONE
          }
        }
        is(LSUState.DONE) {
          when(io.core_state === CoreState.UPDATE) {
            lsu_state := LSUState.IDLE
          }
        }
      }
    }

    // If memory write enable is triggered (STR instruction)
    when(io.mem_rw_enable.write_enable) {
      switch(io.lsu_state) {
        is(LSUState.IDLE) {
          when(io.core_state === CoreState.REQUEST) {
            lsu_state := LSUState.REQUESTING
          }
        }
        is(LSUState.REQUESTING) {
          mem_write_valid   := true.B
          mem_write_address := io.reg_in.rs
          mem_write_data    := io.reg_in.rt
          lsu_state         := LSUState.WAITING
        }
        is(LSUState.WAITING) {
          when(io.mem_write_sender.ready) {
            mem_write_valid := false.B
            lsu_state       := LSUState.DONE
          }
        }
        is(LSUState.DONE) {

    }
  }

  io.lsu_state                     := lsu_state
  io.lsu_out                       := lsu_out
  io.mem_read_address_sender.valid := mem_read_valid
  io.mem_read_address_sender.bits  := mem_read_address
  io.mem_write_sender.valid        := mem_write_valid
  io.mem_write_sender.bits.address := mem_write_address
  io.mem_write_sender.bits.data    := mem_write_data
}


package pc

import chisel3._
import chisel3.util._
import statecode.CoreState
import scala.reflect.internal.Reporter.INFO
// PROGRAM COUNTER
// > Calculates the next PC for each thread to update to (but currently we assume all threads
//   update to the same PC and don't support branch divergence)
// > Currently, each thread in each core has it's own calculation for next PC
// > The NZP register value is set by the CMP instruction (based on >/=/< comparison) to
//   initiate the BRnzp instruction for branching
class ProgramCounter(DataMemWidth: Int = 8, MemAddrWidth: Int = 8) extends Module {
  val io = IO(new Bundle {
    val enable = Input(Bool())

    val core_state = Input(CoreState())

    val decoded_nzp              = Input(UInt(2.W))
    val decoded_immediate        = Input(UInt(DataMemWidth.W))
    val decoded_nzp_write_enable = Input(Bool())
    val decoded_pc_mux           = Input(Bool())

    val alu_out = Input(UInt(DataMemWidth.W))

    val current_pc = Input(UInt(MemAddrWidth.W))
    val next_pc    = Output(UInt(MemAddrWidth.W))
  })

  val nzp     = RegInit(0.U(2.W))
  val next_pc = RegInit(0.U(MemAddrWidth.W))

  when(io.enable) {
    when(io.core_state === CoreState.EXECUTE) {
      when(io.decoded_pc_mux) {
        when((nzp & io.decoded_nzp) =/= 0.U) {
          // On BRnzp instruction, branch to immediate if NZP case matches previous CMP
          next_pc := io.decoded_immediate
        }.otherwise {
          // Otherwise, just update to PC + 1 (next line)
          next_pc := io.current_pc + 1.U
        }
      }.otherwise {
        // By default update to PC + 1 (next line)
        next_pc := io.current_pc + 1.U
      }
    }

    // Store NZP when core_state = UPDATE

  }.otherwise {
    io.next_pc := next_pc
  }
}

    when(io.core_state === CoreState.UPDATE) {
      when(io.decoded_nzp_write_enable) {
        nzp := io.alu_out(2, 0)
      }
    }


package registers

import chisel3._
import chisel3.util._
import statecode.CoreState
import statecode.RegInputOp

class RegisterFiles(ThreadsPerBlk: Int = 4, ThreadId: Int = 0, DataBits: Int = 8) extends Module {
  val io = IO(new Bundle {
    val enable = Input(Bool())

    // Kernel Execution
    val block_id = Input(UInt(8.W))

    // State
    val core_state = Input(CoreState())

    // Instruction Signals
    val decoded_reg_address = new Bundle {
      val rd = Input(UInt(4.W))
      val rs = Input(UInt(4.W))
      val rt = Input(UInt(4.W))
    }

    // Control Signals
    val decoded_reg_write_enable = Input(Bool())
    val decoded_reg_input_mux    = Input(RegInputOp())
    val decoded_immediate        = Input(UInt(DataBits.W))

    // Thread Unit Outputs
    val alu_out = Input(UInt(DataBits.W))
    val lsu_out = Input(UInt(DataBits.W))

    // Register File Outputs
    val reg_out = new Bundle {
      val rs = Output(UInt(DataBits.W))
      val rt = Output(UInt(DataBits.W))
    }
  })

  def InitValByIndex(index: Int): UInt = {
    index match {
      case 14 => ThreadsPerBlk.U
      case 15 => ThreadId.U
      case _  => 0.U(8.W)
    }
  }

  val registers = RegInit(VecInit(Seq.tabulate(16)(i => InitValByIndex(i))))

  val rs = RegInit(0.U(DataBits.W))
  val rt = RegInit(0.U(DataBits.W))

  when(!reset.asBool && io.enable) {
    // [Bad Solution] Shouldn't need to set this every cycle
    // Update the block_id when a new block is issued from dispatcher
    registers(13) := io.block_id

    // Fill rs/rt when core_state = REQUEST
    when(io.core_state === CoreState.REQUEST) {
      rs := registers(io.decoded_reg_address.rs)
      rt := registers(io.decoded_reg_address.rt)
    }.elsewhen(io.core_state === CoreState.UPDATE) { // Store rd when core_state = UPDATE
      // Only allow writing to R0 - R12
      when(io.decoded_reg_write_enable && io.decoded_reg_address.rd < 13.U) {
        switch(io.decoded_reg_input_mux) {
          is(RegInputOp.ARITHMETIC) {
            registers(io.decoded_reg_address.rd) := io.alu_out

    }
  }

  // printf(
  //   "registers: [%d%d%d%d%d%d%d%d%d%d%d%d%d]\n",
  //   registers(0),
  //   registers(1),
  //   registers(2),
  //   registers(3),
  //   registers(4),
  //   registers(5),
  //   registers(6),
  //   registers(7),
  //   registers(8),
  //   registers(9),
  //   registers(10),
  //   registers(11),
  //   registers(12)
  // )

  io.reg_out.rs := rs
  io.reg_out.rt := rt
}

          }
          is(RegInputOp.MEMORY) {
            registers(io.decoded_reg_address.rd) := io.lsu_out
          }
          is(RegInputOp.CONSTANT) {


package scheduler

import chisel3._
import chisel3.util._
import statecode.CoreState
import statecode.LSUState
import statecode.FetcherState

// SCHEDULER
// > Manages the entire control flow of a single compute core processing 1 block
// 1. FETCH - Retrieve instruction at current program counter (PC) from program memory
// 2. DECODE - Decode the instruction into the relevant control signals
// 3. REQUEST - If we have an instruction that accesses memory, trigger the async memory requests from LSUs
// 4. WAIT - Wait for all async memory requests to resolve (if applicable)
// 5. EXECUTE - Execute computations on retrieved data from registers / memory
// 6. UPDATE - Update register values (including NZP register) and program counter
// > Each core has it's own scheduler where multiple threads can be processed with
//   the same control flow at once.
// > Technically, different instructions can branch to different PCs, requiring "branch divergence." In
//   this minimal implementation, we assume no branch divergence (naive approach for simplicity)
class Scheduler(ThreadsPerBlock: Int = 4) extends Module {
  val io = IO(new Bundle {
    val start = Input(Bool())

    // Control signals
    val mem_rw_enable = new Bundle {
      val read_enable  = Input(Bool()) // Enable reading from memory
      val write_enable = Input(Bool()) // Enable writing to memory
    }
    val decoded_ret   = Input(Bool())

    // Memory access state
    val fetcher_state = Input(FetcherState())
    val lsu_state     = Input(Vec(ThreadsPerBlock, LSUState()))

    // Current & Next PC
    val current_pc = Output(UInt(8.W))
    val next_pc    = Input(Vec(ThreadsPerBlock, UInt(8.W)))

    // Execution state
    val core_state = Output(CoreState())
    val done       = Output(Bool())
  })

  val current_pc = RegInit(0.U(8.W))
  val core_state = RegInit(CoreState.IDLE)
  val done       = RegInit(false.B)

  when(!reset.asBool) {
    switch(core_state) {
      is(CoreState.IDLE) {
        // Here after reset (before kernel is launched, or after previous block has been processed)
        when(io.start) {
          core_state := CoreState.FETCH
        }
      }

      is(CoreState.FETCH) {
        // Move on once fetcher_state = FETCHED

      is(CoreState.DECODE) {
        // Decode is synchronous so we move on after one cycle
        core_state := CoreState.REQUEST
      }

      is(CoreState.REQUEST) {
        // Request is synchronous so we move on after one cycle
        core_state := CoreState.WAIT
      }

      is(CoreState.WAIT) {
        // Wait for all LSUs to finish their request before continuing
        // Make sure no lsu_state = REQUESTING or WAITING
        val any_lsu_waiting =
          io.lsu_state.exists(lsu_state => lsu_state === LSUState.WAITING || lsu_state === LSUState.REQUESTING)

        // If no LSU is waiting for a response, move onto the next stage
        when(!any_lsu_waiting) {
          core_state := CoreState.EXECUTE
        }
      }

      is(CoreState.EXECUTE) {
        // Execute is synchronous so we move on after one cycle
        core_state := CoreState.UPDATE
      }

      is(CoreState.UPDATE) {
        when(io.decoded_ret) {
          // If we reach a RET instruction, this block is done executing
          done       := true.B
          core_state := CoreState.DONE
        }.otherwise {
          // TODO: Branch divergence. For now assume all next_pc converge
          current_pc := io.next_pc((ThreadsPerBlock - 1).U)
          core_state := CoreState.FETCH
        }
      }

      is(CoreState.DONE) {
        // no-op
      }
    }
  }

  io.current_pc := current_pc
  io.core_state := core_state
  io.done       := done
}

        when(io.fetcher_state === FetcherState.FETCHED) {
          core_state := CoreState.DECODE
        }


package statecode

import chisel3._

object CoreState extends ChiselEnum {
  val IDLE    = Value("b000".U)
  val FETCH   = Value("b001".U)
  val DECODE  = Value("b010".U)
  val REQUEST = Value("b011".U)
  val WAIT    = Value("b100".U)
  val EXECUTE = Value("b101".U)
  val UPDATE  = Value("b110".U)
  val DONE    = Value("b111".U)
}

object AluOpCode extends ChiselEnum {
  val ADD, SUB, MUL, DIV = Value
}

  val CMP   = Value("b0010".U)
  val ADD   = Value("b0011".U)
  val SUB   = Value("b0100".U)
  val MUL   = Value("b0101".U)
  val DIV   = Value("b0110".U)
  val LDR   = Value("b0111".U)
  val STR   = Value("b1000".U)
  val CONST = Value("b1001".U)
  val RET   = Value("b1111".U)
}

object ControlState extends ChiselEnum {
  val IDLE           = Value("b000".U)
  val READ_WAITING   = Value("b010".U)
  val WRITE_WAITING  = Value("b011".U)
  val READ_RELAYING  = Value("b100".U)
  val WRITE_RELAYING = Value("b101".U)
}

object FetcherState extends ChiselEnum {
  val IDLE     = Value("b000".U)
  val FETCHING = Value("b001".U)
  val FETCHED  = Value("b010".U)
}

object LSUState extends ChiselEnum {
  val IDLE, REQUESTING, WAITING, DONE = Value
}

object RegInputOp extends ChiselEnum {
  val ARITHMETIC = Value("b00".U)
  val MEMORY     = Value("b01".U)
  val CONSTANT   = Value("b10".U)
}


package alu

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.AluOpCode
import statecode.CoreState

class AluModel {
  var alu_out = 0

  def update(
    enable:         Boolean,
    core_state:     CoreState.Type,
    rs:             Int,
    rt:             Int,
    arithmetic_mux: AluOpCode.Type,
    output_mux:     Boolean
  ): Unit = {
    if (enable && core_state == CoreState.EXECUTE) {
      if (output_mux) {
        // Compare operation

        val eq = if (rs == rt) 1 else 0
        val lt = if (rs < rt) 1 else 0
        alu_out = (gt << 2) | (eq << 1) | lt
      } else {
        // Arithmetic operation
        alu_out = arithmetic_mux match {
          case AluOpCode.ADD => (rs + rt) & 0xff
          case AluOpCode.SUB => ((rs - rt) + 256) & 0xff
          case AluOpCode.MUL => (rs * rt) & 0xff
          case AluOpCode.DIV => if (rt == 0) 0 else rs / rt
        }
      }
    }
  }

  def reset(): Unit = {
    alu_out = 0
  }
}

class AluSpec extends AnyFreeSpec with Matchers {
  "Test ALU for all supported operations" in {
    simulate(new Alu) { dut =>
      def cmpResult(rs: Int, rt: Int): Int = {
        val gt = if (rs > rt) 1 else 0
        val eq = if (rs == rt) 1 else 0
        val lt = if (rs < rt) 1 else 0
        (gt << 2) | (eq << 1) | lt
      }

      // Calculate expected result based on operation
      def arithmeticResult(rs: Int, rt: Int, op: AluOpCode.Type): Int = {
        op match {
          case AluOpCode.ADD => (rs + rt) & 0xff
          case AluOpCode.SUB => ((rs - rt) + 256) & 0xff
          case AluOpCode.MUL => (rs * rt) & 0xff
          case AluOpCode.DIV => if (rt == 0) 0 else rs / rt
        }
      }

      val testValues = for {
        rs         <- 0 to 255
        rt         <- 0 to 255
        op         <- List(AluOpCode.ADD, AluOpCode.SUB, AluOpCode.MUL, AluOpCode.DIV)
        core_state <- List(CoreState.IDLE, CoreState.EXECUTE)
        is_cmp     <- List(true, false)
      } yield (rs, rt, op, core_state, is_cmp)

      dut.reset.poke(true.B)
      dut.clock.step()
      dut.reset.poke(false.B)
      dut.clock.step()

      var aluModel = new AluModel()
      // Test each combination
      for ((rs, rt, op, core_state, is_cmp) <- testValues) {
        // Set inputs
        dut.io.enable.poke(true.B)
        dut.io.core_state.poke(core_state)
        dut.io.decoded_alu_op.arithmetic_mux.poke(op)
        dut.io.decoded_alu_op.output_mux.poke(is_cmp)
        dut.io.reg_in.rs.poke(rs.U)
        dut.io.reg_in.rt.poke(rt.U)

        dut.clock.step()

        aluModel.update(
          enable = true,
          core_state = core_state,
          rs = rs,
          rt = rt,
          arithmetic_mux = op,
          output_mux = is_cmp
        )
        // Verify output
        dut.io.alu_out.expect(aluModel.alu_out.U)
      }
    }
  }
}


package controller

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.ControlState

class ControllerModel(AddrBits: Int = 8, DataBits: Int = 16, NumConsumers: Int = 4, NumChannels: Int = 1) {
  var controller_state         = Array.fill(NumChannels)(ControlState.IDLE)
  var mem_read_valid           = Array.fill(NumChannels)(false)
  var mem_read_address         = Array.fill(NumChannels)(0)
  var mem_write_valid          = Array.fill(NumChannels)(false)
  var mem_write_address        = Array.fill(NumChannels)(0)
  var mem_write_data           = Array.fill(NumChannels)(0)
  var consumer_read_ready      = Array.fill(NumConsumers)(false)
  var consumer_read_data       = Array.fill(NumConsumers)(0)
  var consumer_write_ready     = Array.fill(NumConsumers)(false)
  var channel_serving_consumer = Array.fill(NumConsumers)(false)
  var current_consumer         = Array.fill(NumChannels)(0)

  def update(
    consumer_read_valid:  Array[Boolean],
    consumer_read_addr:   Array[Int],
    consumer_write_valid: Array[Boolean],
    consumer_write_addr:  Array[Int],
    consumer_write_data:  Array[Int],
    mem_read_ready:       Array[Boolean],
    mem_read_data:        Array[Int],
    mem_write_ready:      Array[Boolean]
  ): Unit = {
    // Update model state for each channel
    for (i <- 0 until NumChannels) {
      controller_state(i) match {
        case ControlState.IDLE =>
          // Follow the original verilog implementation
          var break = false
          var j     = 0
          while (j < NumConsumers && !break) {
            if (consumer_read_valid(j) && !channel_serving_consumer(j)) {
              channel_serving_consumer(j) = true
              current_consumer(i) = j
              mem_read_valid(i) = true
              mem_read_address(i) = consumer_read_addr(j)
              controller_state(i) = ControlState.READ_WAITING
              break = true
            } else if (consumer_write_valid(j) && !channel_serving_consumer(j)) {
              channel_serving_consumer(j) = true
              current_consumer(i) = j
              mem_write_valid(i) = true
              mem_write_address(i) = consumer_write_addr(j)
              mem_write_data(i) = consumer_write_data(j)
              controller_state(i) = ControlState.WRITE_WAITING
              break = true
            }
            j += 1
          }

        case ControlState.READ_WAITING =>
          if (mem_read_ready(i)) {
            mem_read_valid(i) = false
            consumer_read_ready(current_consumer(i)) = true
            consumer_read_data(current_consumer(i)) = mem_read_data(i)
            controller_state(i) = ControlState.READ_RELAYING
          }

        case ControlState.WRITE_WAITING =>

            mem_write_valid(i) = false
            consumer_write_ready(current_consumer(i)) = true
            controller_state(i) = ControlState.WRITE_RELAYING
          }

        case ControlState.READ_RELAYING =>
          if (!consumer_read_valid(current_consumer(i))) {
            channel_serving_consumer(current_consumer(i)) = false
            consumer_read_ready(current_consumer(i)) = false
            controller_state(i) = ControlState.IDLE
          }

        case ControlState.WRITE_RELAYING =>
          if (!consumer_write_valid(current_consumer(i))) {
            channel_serving_consumer(current_consumer(i)) = false
            consumer_write_ready(current_consumer(i)) = false
            controller_state(i) = ControlState.IDLE
          }
      }
    }

    // // Debug printing
    // println("=== Model Internal State ===")
    // println(s"-controller_state: ${controller_state.mkString(", ")}")
    // println(s"-mem_read_valid: ${mem_read_valid.mkString(", ")}")
    // println(s"-mem_read_address: ${mem_read_address.mkString(", ")}")
    // println(s"-mem_write_valid: ${mem_write_valid.mkString(", ")}")
    // println(s"-mem_write_address: ${mem_write_address.mkString(", ")}")
    // println(s"-mem_write_data: ${mem_write_data.mkString(", ")}")
    // println(s"-consumer_read_ready: ${consumer_read_ready.mkString(", ")}")
    // println(s"-consumer_read_data: ${consumer_read_data.mkString(", ")}")
    // println(s"-consumer_write_ready: ${consumer_write_ready.mkString(", ")}")
    // println(s"-channel_serving_consumer: ${channel_serving_consumer.mkString(", ")}")
    // println(s"-current_consumer: ${current_consumer.mkString(", ")}")
    // println("==================\n")
  }

  // def reset(): Unit = {
  //   controller_state = Array.fill(3)(ControlState.IDLE)
  //   mem_read_valid = Array.fill(1)(false)
  //   mem_read_address = Array.fill(1)(0)
  //   mem_write_valid = Array.fill(1)(false)
  //   mem_write_address = Array.fill(1)(0)
  //   mem_write_data = Array.fill(1)(0)
  //   consumer_read_ready = Array.fill(4)(false)
  //   consumer_read_data = Array.fill(4)(0)
  //   consumer_write_ready = Array.fill(4)(false)
  //   channel_serving_consumer = Array.fill(4)(false)
  //   current_consumer = Array.fill(1)(0)
  // }
}

class ControllerSpec extends AnyFreeSpec with Matchers {
  "Test controller" - {
    "should match model behavior with random inputs" in {
      val AddrBits     = 8
      val DataBits     = 16
      val NumConsumers = 4
      val NumChannels  = 1

      simulate(new Controller(AddrBits, DataBits, NumConsumers, NumChannels)) { dut =>
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        dut.clock.step()

        var cnt             = 0
        val rng             = new scala.util.Random(42)
        val controllerModel = new ControllerModel()

        while (cnt < 10000) {
          // Generate random inputs
          val consumer_read_valid  = Array.fill(NumConsumers)(rng.nextBoolean())
          val consumer_read_addr   = Array.fill(NumConsumers)(rng.nextInt(AddrBits))
          val consumer_write_valid = Array.fill(NumConsumers)(rng.nextBoolean())
          val consumer_write_addr  = Array.fill(NumConsumers)(rng.nextInt(AddrBits))
          val consumer_write_data  = Array.fill(NumConsumers)(rng.nextInt(DataBits))
          val mem_read_ready       = Array.fill(NumChannels)(rng.nextBoolean())
          val mem_read_data        = Array.fill(NumChannels)(rng.nextInt(DataBits))
          val mem_write_ready      = Array.fill(NumChannels)(rng.nextBoolean())

          // Apply inputs to DUT
          for (i <- 0 until NumConsumers) {
            dut.io.consumer_read_addr_receiver(i).valid.poke(consumer_read_valid(i).B)
            dut.io.consumer_read_addr_receiver(i).bits.poke(consumer_read_addr(i).U)
            // dut.io.consumer_write_receiver(i).valid.poke(consumer_write_valid(i).B)
            // dut.io.consumer_write_receiver(i).bits.address.poke(consumer_write_addr(i).U)
            // dut.io.consumer_write_receiver(i).bits.data.poke(consumer_write_data(i).U)
            dut.io.consumer_write_receiver.map(receiver => {
              receiver(i).valid.poke(consumer_write_valid(i).B)
              receiver(i).bits.address.poke(consumer_write_addr(i).U)
              receiver(i).bits.data.poke(consumer_write_data(i).U)
            })
          }

          for (i <- 0 until NumChannels) {
            dut.io.mem_read_sender(i).ready.poke(mem_read_ready(i).B)
            dut.io.mem_read_data(i).poke(mem_read_data(i).U)
            // dut.io.mem_write_sender(i).ready.poke(mem_write_ready(i).B)
            dut.io.mem_write_sender.map(_.apply(i).ready.poke(mem_write_ready(i).B))
          }

          dut.clock.step()

          // println("\n\n=== Input Signals ===")
          // println(s"consumer_read_valid: ${consumer_read_valid.mkString(", ")}")
          // println(s"consumer_read_addr: ${consumer_read_addr.mkString(", ")}")
          // println(s"consumer_write_valid: ${consumer_write_valid.mkString(", ")}")
          // println(s"consumer_write_addr: ${consumer_write_addr.mkString(", ")}")
          // println(s"consumer_write_data: ${consumer_write_data.mkString(", ")}")
          // println(s"mem_read_ready: ${mem_read_ready.mkString(", ")}")
          // println(s"mem_read_data: ${mem_read_data.mkString(", ")}")
          // println(s"mem_write_ready: ${mem_write_ready.mkString(", ")}")

          // Update model
          controllerModel.update(
            consumer_read_valid,
            consumer_read_addr,
            consumer_write_valid,
            consumer_write_addr,
            consumer_write_data,
            mem_read_ready,
            mem_read_data,
            mem_write_ready
          )

          // Compare model with DUT
          for (i <- 0 until NumConsumers) {
            dut.io.consumer_read_addr_receiver(i).ready.expect(controllerModel.consumer_read_ready(i).B)
            dut.io.consumer_read_data(i).expect(controllerModel.consumer_read_data(i).U)
            // dut.io.consumer_write_receiver(i).ready.expect(controllerModel.consumer_write_ready(i).B)
            dut.io.consumer_write_receiver.map(_.apply(i).ready.expect(controllerModel.consumer_write_ready(i).B))
          }

          for (i <- 0 until NumChannels) {
            dut.io.mem_read_sender(i).valid.expect(controllerModel.mem_read_valid(i).B)
            dut.io.mem_read_sender(i).bits.expect(controllerModel.mem_read_address(i).U)
            // dut.io.mem_write_sender(i).valid.expect(controllerModel.mem_write_valid(i).B)
            // if (controllerModel.mem_write_valid(i)) {
            //   dut.io.mem_write_sender(i).bits.address.expect(controllerModel.mem_write_address(i).U)
            //   dut.io.mem_write_sender(i).bits.data.expect(controllerModel.mem_write_data(i).U)
            // }
            dut.io.mem_write_sender.map(sender => {
              sender(i).valid.expect(controllerModel.mem_write_valid(i).B)
              if (controllerModel.mem_write_valid(i)) {
                sender(i).bits.address.expect(controllerModel.mem_write_address(i).U)
                sender(i).bits.data.expect(controllerModel.mem_write_data(i).U)
              }
            })
          }

          cnt += 1
        }
      }
    }
  }
}


package core

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.CoreState
import statecode.FetcherState
import statecode.LSUState
import statecode.RegInputOp
import statecode.AluOpCode

class CoreModel(
  DataMemAddrBits:    Int = 8,
  DataMemDataBits:    Int = 8,
  ProgramMemAddrBits: Int = 8,
  ProgramMemDataBits: Int = 16,
  ThreadsPerBlock: Int = 4) {
  // State variables matching the Core module's submodules
  var fetcherModel   = new fetcher.FetcherModel()
  var decoderModel   = new decoder.DecoderModel()
  var schedulerModel = new scheduler.SchedulerModel(ThreadsPerBlock)

  // Arrays for per-thread components
  var aluModels     = Array.fill(ThreadsPerBlock)(new alu.AluModel())
  var lsuModels     = Array.fill(ThreadsPerBlock)(new lsu.LsuModel())
  var regfileModels = Array.fill(ThreadsPerBlock)(new registers.RegModel())
  var pcModels      = Array.fill(ThreadsPerBlock)(new pc.PCModel())

  class PrevCycleData {
    // Fetcher outputs
    var fetcher_state = FetcherState.IDLE
    var read_valid    = false
    var read_address  = 0
    var instruction   = 0
    // Fetcher end

    // Decoder outputs
    var decoded_rd_address         = 0
    var decoded_rs_address         = 0
    var decoded_rt_address         = 0
    var decoded_nzp                = 0
    var decoded_immediate          = 0
    var decoded_reg_write_enable   = false
    var decoded_mem_read_enable    = false
    var decoded_mem_write_enable   = false
    var decoded_nzp_write_enable   = false
    var decoded_reg_input_mux      = RegInputOp.ARITHMETIC
    var decoded_alu_arithmetic_mux = AluOpCode.ADD
    var decoded_alu_output_mux     = false
    var decoded_pc_mux             = false
    var decoded_ret                = false
    // Decoder end

    // Scheduler outputs
    var current_pc = 0
    var core_state = CoreState.IDLE
    var done       = false
    // Scheduler end

    // ALU outputs
    var alu_out = Array.fill(ThreadsPerBlock)(0)
    // ALU end

    // LSU outputs
    // var lsu_read_valid    = Array.fill(ThreadsPerBlock)(false)
    // var lsu_read_address  = Array.fill(ThreadsPerBlock)(0)
    // var lsu_write_valid   = Array.fill(ThreadsPerBlock)(false)
    // var lsu_write_address = Array.fill(ThreadsPerBlock)(0)
    // var lsu_write_data    = Array.fill(ThreadsPerBlock)(0)
    var lsu_state = Array.fill(ThreadsPerBlock)(LSUState.IDLE)
    var lsu_out   = Array.fill(ThreadsPerBlock)(0)
    // LSU end

    // Register Files outputs
    var regout_rs = Array.fill(ThreadsPerBlock)(0)
    var regout_rt = Array.fill(ThreadsPerBlock)(0)
    // Register end

    def save(model: fetcher.FetcherModel) = {
      fetcher_state = model.fetcher_state
      read_valid = model.read_valid
      read_address = model.read_address
      instruction = model.instruction
    }

    def save(model: decoder.DecoderModel) = {
      decoded_rd_address = model.decoded_rd_address
      decoded_rs_address = model.decoded_rs_address
      decoded_rt_address = model.decoded_rt_address
      decoded_nzp = model.decoded_nzp
      decoded_immediate = model.decoded_immediate
      decoded_reg_write_enable = model.decoded_reg_write_enable
      decoded_mem_read_enable = model.decoded_mem_read_enable
      decoded_mem_write_enable = model.decoded_mem_write_enable
      decoded_nzp_write_enable = model.decoded_nzp_write_enable
      decoded_reg_input_mux = model.decoded_reg_input_mux
      decoded_alu_arithmetic_mux = model.decoded_alu_arithmetic_mux
      decoded_alu_output_mux = model.decoded_alu_output_mux
      decoded_pc_mux = model.decoded_pc_mux
      decoded_ret = model.decoded_ret
    }

    def save(model: scheduler.SchedulerModel) = {
      core_state = model.core_state
      current_pc = model.current_pc
      done = model.done
    }

    def save(model: alu.AluModel, idx: Int) = {
      alu_out(idx) = model.alu_out
    }

    def save(model: pc.PCModel, idx: Int) = {
      next_pc(idx) = model.next_pc
    }

    def save(model: lsu.LsuModel, idx: Int) = {
      // lsu_read_valid(idx) = model.read_valid
      // lsu_read_address(idx) = model.read_address
      // lsu_write_valid(idx) = model.write_valid
      // lsu_write_address(idx) = model.write_address
      // lsu_write_data(idx) = model.write_data
      lsu_state(idx) = model.lsu_state
      lsu_out(idx) = model.output_data
    }

    def save(model: registers.RegModel, idx: Int) = {
      regout_rs(idx) = model.rs_out()
      regout_rt(idx) = model.rt_out()
    }
  }

  var prevCycleData = new PrevCycleData

  // core's 8 outputs
  var done = false

  var program_mem_read_address_valid = false
  var program_mem_read_address_bits  = 0

  var data_mem_read_address_valid = Array.fill(ThreadsPerBlock)(false)
  var data_mem_read_address_bits  = Array.fill(ThreadsPerBlock)(0)

  var data_mem_write_address_valid = Array.fill(ThreadsPerBlock)(false)
  var data_mem_write_address_bits  = Array.fill(ThreadsPerBlock)(0)
  var data_mem_write_data_bits     = Array.fill(ThreadsPerBlock)(0)

  // debug var
  var cycle = 0

  def update(
    // Kernel execution inputs
    start:        Boolean,
    block_id:     Int,
    thread_count: Int,

    // Memory inputs
    program_mem_read_ready: Boolean,
    program_mem_read_data:  Int,

    // Data Memory signals
    data_mem_read_data:   Array[Int],
    data_mem_read_ready:  Array[Boolean],
    data_mem_write_ready: Array[Boolean]
  ): Unit = {
    // Update fetcher
    fetcherModel.update(
      core_state = prevCycleData.core_state,
      current_pc = prevCycleData.current_pc,
      mem_read_ready = program_mem_read_ready,
      mem_read_data = program_mem_read_data
    )

    // Update decoder
    decoderModel.update(
      core_state = prevCycleData.core_state,
      instruction = prevCycleData.instruction
    )

    // Update scheduler
    schedulerModel.update(
      start = start,
      mem_read_enable = prevCycleData.decoded_mem_read_enable,
      mem_write_enable = prevCycleData.decoded_mem_write_enable,
      decoded_ret = prevCycleData.decoded_ret,
      fetcher_state = prevCycleData.fetcher_state,
      lsu_state = prevCycleData.lsu_state,
      next_pc = prevCycleData.next_pc
    )

    // Update per-thread components
    for (i <- 0 until ThreadsPerBlock) {
      val enable = i < thread_count

      val rd_in = prevCycleData.decoded_rd_address
      val rs_in = prevCycleData.decoded_rs_address
      val rt_in = prevCycleData.decoded_rt_address

      // Update ALU
      aluModels(i).update(
        enable = enable,
        core_state = prevCycleData.core_state,
        rs = prevCycleData.regout_rs(i),
        rt = prevCycleData.regout_rt(i),
        arithmetic_mux = prevCycleData.decoded_alu_arithmetic_mux,
        output_mux = prevCycleData.decoded_alu_output_mux
      )

      // Update LSU
      lsuModels(i).update(
        enable = enable,
        core_state = prevCycleData.core_state,
        read_enable = prevCycleData.decoded_mem_read_enable,
        write_enable = prevCycleData.decoded_mem_write_enable,
        rs = prevCycleData.regout_rs(i),
        rt = prevCycleData.regout_rt(i),
        mem_read_data = data_mem_read_data(i),
        mem_read_ready = data_mem_read_ready(i),
        mem_write_ready = data_mem_write_ready(i)
      )

      // Update Register Files
      regfileModels(i).update(
        enable = enable,
        block_id = block_id,
        core_state = prevCycleData.core_state,
        rs = rs_in,
        rd = rd_in,
        rt = rt_in,
        decoded_reg_write_enable = prevCycleData.decoded_reg_write_enable,
        decoded_reg_input_mux = prevCycleData.decoded_reg_input_mux,
        decoded_immediate = prevCycleData.decoded_immediate,
        alu_out = prevCycleData.alu_out(i),
        lsu_out = prevCycleData.lsu_out(i)
      )

      // Update PC
      pcModels(i).update(
        enable = enable,
        core_state = prevCycleData.core_state,
        decoded_nzp = prevCycleData.decoded_nzp,
        decoded_immediate = prevCycleData.decoded_immediate,
        decoded_nzp_write_enable = prevCycleData.decoded_nzp_write_enable,
        decoded_pc_mux = prevCycleData.decoded_pc_mux,
        alu_out = prevCycleData.alu_out(i),
        current_pc = prevCycleData.current_pc
      )

      prevCycleData.save(aluModels(i), i)
      prevCycleData.save(lsuModels(i), i)
      prevCycleData.save(regfileModels(i), i)
      prevCycleData.save(pcModels(i), i)
    }

    // update outputs
    done = schedulerModel.done
    program_mem_read_address_valid = fetcherModel.read_valid
    program_mem_read_address_bits = fetcherModel.read_address

    for (i <- 0 until ThreadsPerBlock) {
      data_mem_read_address_valid(i) = lsuModels(i).read_valid
      data_mem_read_address_bits(i) = lsuModels(i).read_address
      data_mem_write_address_valid(i) = lsuModels(i).write_valid
      data_mem_write_address_bits(i) = lsuModels(i).write_address
      data_mem_write_data_bits(i) = lsuModels(i).write_data
    }

    prevCycleData.save(fetcherModel)
    prevCycleData.save(decoderModel)
    prevCycleData.save(schedulerModel)

    // print and update debug info
    // println(
    //   "#Model After Cycle: " + cycle + ", CoreState: " + prevCycleData.core_state + ", PC: " + prevCycleData.current_pc + ", Fetcher State: " + fetcherModel.fetcher_state
    // )
    cycle += 1
  }

  // def reset(): Unit = {
  //   fetcherModel.reset()
  //   decoderModel.reset()
  //   schedulerModel.reset()
  //   aluModels.foreach(_.reset())
  //   lsuModels.foreach(_.reset())
  //   regfileModels.foreach(_.reset())
  //   pcModels.foreach(_.reset())
  // }

}

class CoreSpec extends AnyFreeSpec with Matchers {
  "Test Core" - {
    "should match model behavior with random inputs" in {
      val DataMemAddrBits:    Int = 8
      val DataMemDataBits:    Int = 8
      val ProgramMemAddrBits: Int = 8
      val ProgramMemDataBits: Int = 16
      val ThreadsPerBlock:    Int = 4
      simulate(new Core(DataMemAddrBits, DataMemDataBits, ProgramMemAddrBits, ProgramMemDataBits, ThreadsPerBlock)) {
        dut =>
          dut.reset.poke(true.B)
          dut.clock.step()
          dut.reset.poke(false.B)
          dut.clock.step()

          // println("########## Start ##########")

          var cnt       = 0
          val rng       = new scala.util.Random(42) // 42 is the seed for reproducibility
          val coreModel = new CoreModel()
          while (cnt < 10000) {
            val block_id             = rng.nextInt(13)
            val thread_count         = rng.nextInt(ThreadsPerBlock)
            val mem_read_ready       = rng.nextBoolean()
            val mem_read_data        = rng.nextInt(256)
            val data_mem_read_data   = Array.fill(ThreadsPerBlock)(rng.nextInt(256))
            val data_mem_read_ready  = Array.fill(ThreadsPerBlock)(rng.nextBoolean())
            val data_mem_write_ready = Array.fill(ThreadsPerBlock)(rng.nextBoolean())

            dut.io.start.poke(true.B)
            dut.io.block_id.poke(block_id.U)
            dut.io.thread_count.poke(thread_count.U)

            dut.io.program_mem_read_address_sender.ready.poke(mem_read_ready.B)
            dut.io.program_mem_read_data.poke(mem_read_data.U)

            for (i <- 0 until ThreadsPerBlock) {
              dut.io.data_mem_read_data(i).poke(data_mem_read_data(i).U)
              dut.io.data_mem_read_address_sender(i).ready.poke(data_mem_read_ready(i).B)
              dut.io.data_mem_write_sender(i).ready.poke(data_mem_write_ready(i).B)
            }

            dut.clock.step()

            // println(s"\n=== Random Values for Cycle $cnt ===")
            // println(s"Block ID: $block_id")
            // println(s"Thread Count: $thread_count")
            // println(s"Program Memory Read: [ready=$mem_read_ready, data=$mem_read_data]")
            // println(s"Read Data:    [${data_mem_read_data.mkString(", ")}]")
            // println(s"Read Ready:   [${data_mem_read_ready.mkString(", ")}]")
            // println(s"Write Ready:  [${data_mem_write_ready.mkString(", ")}]")
            // println()
            // println("--Aft Core State: " + dut.io.core_state.peekValue() + ", current_pc: " + dut.io.current_pc.peek())

            coreModel.update(
              start = true,
              block_id = block_id,
              thread_count = thread_count,
              // Memory inputs
              program_mem_read_ready = mem_read_ready,
              program_mem_read_data = mem_read_data,
              // Data Memory signals
              data_mem_read_data = data_mem_read_data,
              data_mem_read_ready = data_mem_read_ready,
              data_mem_write_ready = data_mem_write_ready
            )

            // verify all 8 outputs
            dut.io.done.expect(coreModel.done.B)
            dut.io.program_mem_read_address_sender.valid.expect(coreModel.program_mem_read_address_valid.B)
            dut.io.program_mem_read_address_sender.bits.expect(coreModel.program_mem_read_address_bits.U)

            for (i <- 0 until ThreadsPerBlock) {
              dut.io.data_mem_read_address_sender(i).valid.expect(coreModel.data_mem_read_address_valid(i).B)
              dut.io.data_mem_read_address_sender(i).bits.expect(coreModel.data_mem_read_address_bits(i).U)
              dut.io.data_mem_write_sender(i).valid.expect(coreModel.data_mem_write_address_valid(i).B)
              dut.io.data_mem_write_sender(i).bits.address.expect(coreModel.data_mem_write_address_bits(i).U)
              dut.io.data_mem_write_sender(i).bits.data.expect(coreModel.data_mem_write_data_bits(i).U)
            }

            cnt += 1
          }
      }
    }
  }
}

    // PC output
    var next_pc = Array.fill(ThreadsPerBlock)(0)


package decoder

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.CoreState
import statecode.AluOpCode
import statecode.RegInputOp

class DecoderModel {
  var decoded_rd_address         = 0
  var decoded_rs_address         = 0
  var decoded_rt_address         = 0
  var decoded_nzp                = 0
  var decoded_immediate          = 0
  var decoded_reg_write_enable   = false
  var decoded_mem_read_enable    = false
  var decoded_mem_write_enable   = false
  var decoded_nzp_write_enable   = false
  var decoded_reg_input_mux      = RegInputOp.ARITHMETIC
  var decoded_alu_arithmetic_mux = AluOpCode.ADD
  var decoded_alu_output_mux     = false
  var decoded_pc_mux             = false
  var decoded_ret                = false

  def update(
    core_state:  CoreState.Type,
    instruction: Int
  ): Unit = {
    if (core_state == CoreState.DECODE) {
      decoded_rd_address = (instruction & 0xf00) >> 8
      decoded_rs_address = (instruction & 0xf0) >> 4
      decoded_rt_address = instruction & 0x000f
      decoded_immediate = instruction & 0x00ff
      decoded_nzp = (instruction & 0xe00) >> 9

      // reset
      decoded_reg_write_enable = false
      decoded_mem_read_enable = false
      decoded_mem_write_enable = false
      decoded_nzp_write_enable = false
      decoded_reg_input_mux = RegInputOp.ARITHMETIC
      decoded_alu_arithmetic_mux = AluOpCode.ADD
      decoded_alu_output_mux = false
      decoded_pc_mux = false
      decoded_ret = false

      val instruction_opcode = (instruction & 0xf000) >> 12
      instruction_opcode match {
        case 0x0 => {
          // DecoderState.NOP
        }
        case 0x1 => {
          // DecoderState.BRnzp
          decoded_pc_mux = true
        }
        case 0x2 => {
          // DecoderState.CMP
          decoded_alu_output_mux = true
          decoded_nzp_write_enable = true
        }
        case 0x3 => {
          // DecoderState.ADD
          decoded_reg_write_enable = true
          decoded_reg_input_mux = RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux = AluOpCode.ADD
        }
        case 0x4 => {
          // DecoderState.SUB
          decoded_reg_write_enable = true
          decoded_reg_input_mux = RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux = AluOpCode.SUB
        }
        case 0x5 => {
          // DecoderState.MUL
          decoded_reg_write_enable = true
          decoded_reg_input_mux = RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux = AluOpCode.MUL
        }
        case 0x6 => {
          // DecoderState.DIV
          decoded_reg_write_enable = true
          decoded_reg_input_mux = RegInputOp.ARITHMETIC
          decoded_alu_arithmetic_mux = AluOpCode.DIV
        }
        case 0x7 => {
          // DecoderState.LDR
          decoded_reg_write_enable = true
          decoded_reg_input_mux = RegInputOp.MEMORY
          decoded_mem_read_enable = true
        }
        case 0x8 => {
          // DecoderState.STR
          decoded_mem_write_enable = true
        }
        case 0x9 => {
          // DecoderState.CONST
          decoded_reg_write_enable = true
          decoded_reg_input_mux = RegInputOp.CONSTANT
        }
        case 0xf => {
          // DecoderState.RET
          decoded_ret = true
        }

      }
    }

  }
}

class DecoderSpec extends AnyFreeSpec with Matchers {
  "Test decoder" - {
    "should match model behavior with random inputs" in {
      val ProgramMemAddrBits: Int = 8
      val ProgramMemDataBits: Int = 16
      simulate(new Decoder) { dut =>
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        dut.clock.step()

        var cnt          = 0
        val rng          = new scala.util.Random(42) // 42 is the seed for reproducibility
        val decoderModel = new DecoderModel()
        while (cnt < 10000) {
          val core_state  = CoreState.DECODE
          val instruction = rng.nextInt(65536) // 2^16 for 16-bit data

          //   println("instruction: " + f"${instruction.toBinaryString}%16s".replace(' ', '0').grouped(4).mkString(" "))

          dut.io.core_state.poke(core_state)
          dut.io.instruction.poke(instruction.U)

          dut.clock.step()

          decoderModel.update(core_state, instruction)
          dut.io.decoded_reg_address.rd.expect(decoderModel.decoded_rd_address.U)
          dut.io.decoded_reg_address.rs.expect(decoderModel.decoded_rs_address.U)
          dut.io.decoded_reg_address.rt.expect(decoderModel.decoded_rt_address.U)

          dut.io.decoded_nzp.expect(decoderModel.decoded_nzp.U)
          dut.io.decoded_immediate.expect(decoderModel.decoded_immediate.U)

          dut.io.mem_rw_enable.read_enable.expect(decoderModel.decoded_mem_read_enable.B)
          dut.io.mem_rw_enable.write_enable.expect(decoderModel.decoded_mem_write_enable.B)

          dut.io.decoded_reg_write_enable.expect(decoderModel.decoded_reg_write_enable.B)
          dut.io.decoded_nzp_write_enable.expect(decoderModel.decoded_nzp_write_enable.B)
          assert(dut.io.decoded_reg_input_mux.peekValue().asBigInt == decoderModel.decoded_reg_input_mux.litValue)

          assert(
            dut.io.decoded_alu_op.arithmetic_mux
              .peekValue()
              .asBigInt == decoderModel.decoded_alu_arithmetic_mux.litValue
          )
          dut.io.decoded_alu_op.output_mux.expect(decoderModel.decoded_alu_output_mux.B)

          dut.io.decoded_pc_mux.expect(decoderModel.decoded_pc_mux.B)
          dut.io.decoded_ret.expect(decoderModel.decoded_ret.B)

          cnt += 1
        }
      }
    }
  }
}