package final_project

import chisel3._
import chisel3.util._

class BigBankReadData(val w: Int) extends Bundle{
    val data = Vec(48, SInt(w.W))
}

class SmallBankReadData(val w: Int) extends Bundle{
    val data = Vec(8, SInt(w.W))
}

class PackReadDataBundle(val w: Int, val h_w: Int, val c_w: Int, val big_w: Int) extends TransBundle{
    val job = Input(new PackJob(w, h_w, c_w, big_w))
    val from_big = Input(Vec(2, new BigBankReadData(w)))
    val from_small = Input(Vec(2, Vec(4, new SmallBankReadData(w))))
    val output = Output(new PackedData(w))
}

class PackReadData(val w: Int, val h_w: Int, val c_w: Int, val big_w: Int) extends Module{
    val io = IO(new PackReadDataBundle(w, h_w, c_w, big_w))
    val cache = RegInit(0.U.asTypeOf(Vec(2, Vec(64, Vec(10, SInt(w.W))))))

    val cnt_ic = RegInit(0.U.asTypeOf(ACounter(c_w.W)))
    val cnt_swap = RegInit(0.U.asTypeOf(ACounter(big_w.W)))
    val cnt_x = RegInit(0.U.asTypeOf(ACounter(h_w.W)))
    val cnt_y = RegInit(0.U.asTypeOf(ACounter(h_w.W)))

    val output_reg = Wire((new PackedData(w)))
    io.output := output_reg

    val state = RegInit(0.U(1.W))

    def copy64(): Unit = {
        for(i <- 0 to 7){
            for(j <- 1 to 6)
                output_reg.mat(i*8+j) := io.from_big(state).data(i*6+j-1)
            output_reg.mat(i*8+0) := io.from_small(state)(1).data(i)
            output_reg.mat(i*8+7) := io.from_small(state)(2).data(i)
        }
    }

    copy64()
    when(cnt_y.ccnt===0.U){
        for(i <- 2 to 7)
            output_reg.up(i) := io.from_big(state).data(6+i-2)
        output_reg.up(1) := io.from_small(state)(1).data(1)
        output_reg.up(8) := io.from_small(state)(2).data(1)
        when(cnt_x.ccnt===0.U){
            output_reg.up(0) := io.from_big(state).data(6)
        }.otherwise{
            output_reg.up(0) := io.from_small(state)(0).data(1)
        }
        when(cnt_x.ccnt===cnt_x.cend){
            output_reg.up(9) := io.from_big(state).data(6+5)
        }.otherwise{
            output_reg.up(9) := io.from_small(state)(3).data(1)
        }
    }.otherwise{
        output_reg.up := cache(~state)(0)
    }

    var x = 1
    for(i <- 0 to 63){
        if(i==(1<<x)-1){
            val t = (1<<x)-1

                cache(~state)(t) := cache(~state)(0)
            }.otherwise{
                if(t<63)
                    cache(~state)(t) := cache(~state)(t+1)
            }
            x += 1
        }
    }
    x = 1
    for(i <- 0 to 63){
        if(i==(1<<x)-1){
            x += 1
        } else {
            for(j <- 0 to 1){
                cache(j)(i) := cache(j)(i+1)
            }
        }
    }

    when(cnt_y.ccnt===cnt_y.cend){
        for(i <- 2 to 7)
            output_reg.down(i) := io.from_big(state).data(48-12+i-2)
        output_reg.down(1) := io.from_small(state)(1).data(6)
        output_reg.down(8) := io.from_small(state)(2).data(6)
        when(cnt_x.ccnt===0.U){
            output_reg.down(0) := io.from_big(state).data(48-12)
        }.otherwise{
            output_reg.down(0) := io.from_small(state)(0).data(6)
        }
        when(cnt_x.ccnt===cnt_x.cend){
            output_reg.down(9) := io.from_big(state).data(48-12+5)
        }.otherwise{
            output_reg.down(9) := io.from_small(state)(3).data(6)
        }
    }.otherwise{
        for(i <- 2 to 7)
            output_reg.down(i) := io.from_big(~state).data(i-2)
        output_reg.down(1) := io.from_small(~state)(1).data(0)
        output_reg.down(8) := io.from_small(~state)(2).data(0)
        when(cnt_x.ccnt===0.U){
            output_reg.down(0) := io.from_big(~state).data(0)
        }.otherwise{
            output_reg.down(0) := io.from_small(~state)(0).data(0)
        }
        when(cnt_x.ccnt===cnt_x.cend){
            output_reg.down(9) := io.from_big(~state).data(5)
        }.otherwise{
            output_reg.down(9) := io.from_small(~state)(3).data(0)
        }
    }

    when(cnt_x.ccnt===0.U){
        for(i <- 0 to 7)
            output_reg.left(i) := io.from_big(state).data(i*6)
    }.otherwise{
        output_reg.left := io.from_small(state)(0).data
    }
    when(cnt_x.ccnt===cnt_x.cend){
        for(i <- 0 to 7)
            output_reg.right(i) := io.from_big(state).data(i*6+5)
    }.otherwise{
        output_reg.right := io.from_small(state)(3).data
    }

    val nxt_up = Wire(Vec(10, SInt(w.W)))
    for(i <- 2 to 7)
        nxt_up(i) := io.from_big(state).data(48-6+i-2)
    nxt_up(1) := io.from_small(state)(1).data(7)
    nxt_up(8) := io.from_small(state)(2).data(7)
    when(cnt_x.ccnt===0.U){
        nxt_up(0) := io.from_big(state).data(48-6)
    }.otherwise{
        nxt_up(0) := io.from_small(state)(0).data(7)
    }
    when(cnt_x.ccnt===cnt_x.cend){
        nxt_up(9) := io.from_big(state).data(48-6+5)
    }.otherwise{
        nxt_up(9) := io.from_small(state)(3).data(7)
    }
    x = 1
    for(i <- 0 to 63){
        if(i==(1<<x)-1){
            val t = (1<<x)-1
            when(if(x==6) cnt_ic.cend(x-1) else ~cnt_ic.cend(x)&&cnt_ic.cend(x-1)){
                cache(state)(t) := nxt_up
            }.otherwise{
                if(t<63)
                    cache(state)(t) := cache(state)(t+1)
            }
            x += 1
        }
    }
    io.valid_out := io.valid_in
    when(io.flag_job){
        cnt_x.set(io.job.cnt_x_end)
        cnt_y.set(io.job.cnt_y_end)
        cnt_ic.set(io.job.in_chan)
        cnt_swap.set(io.job.cnt_swap_end)
        state := 0.U
    }.elsewhen(io.valid_in){
        cnt_ic.inc()
        when(cnt_swap.inc()){
            state := ~state
            when(cnt_y.inc()){
                cnt_x.inc()
            }
        }

    }
}


package final_project

import chisel3._
import chisel3.util._

    val now_addr = RegInit(0.U.asTypeOf(RCounter(addr_w)))
    io.addr := now_addr.ccnt
    io.valid_out := false.B
    when(io.flag_job){
        now_addr.set(io.begin_addr, io.addr_end)
    }.elsewhen(io.valid_in){
        now_addr.inc()
    }
}

class WeightReaderBundle(val w: Int, val addr_w: Int) extends TransBundle{
    val addr_end = Input(UInt(addr_w.W))
    val begin_addr = Input(UInt(addr_w.W))
    val addr = Output(UInt(addr_w.W))
}


package final_project

import chisel3._
import chisel3.util._

class BigBankWriteData(val w: Int) extends Bundle{
    val data = Vec(48, SInt(w.W))
}

class SmallBankWriteData(val w: Int) extends Bundle{
    val data = Vec(8, SInt(w.W))
}

// 2 groups, each group has:
// 4 small banks
// 1 big bank

// h_w: (w*h)(height)
// c_w: w(num of channels)

class Writer(val w: Int, val addr_w: Int, val h_w: Int, val c_w: Int, val id_w: Int) extends Module{
    val io = IO(new Bundle{
        val in_from_quant = Input(new QuantedData(w))
        val valid_in = Input(Bool())
        val valid_out = Output(Bool())
        val flag_job = Input(Bool())
        val job = Input(new WriteJobs(addr_w, h_w, c_w, id_w))
        val to_bigbank = Output(new BigBankWriteData(w))
        val to_smallbank = Output(Vec(2, new SmallBankWriteData(w)))
        val to_banks = Output(new AddressWriteGroup(addr_w, id_w))
    })

    val gen = RegInit(0.U.asTypeOf(new WriteGroup(addr_w, h_w, c_w, id_w)))

    io.valid_out := false.B
    io.to_bigbank := 0.U.asTypeOf(io.to_bigbank)
    io.to_smallbank := 0.U.asTypeOf(io.to_smallbank)
    io.to_banks := 0.U.asTypeOf(io.to_banks)

    for(i <- 0 to 7)
        for(j <- 0 to 5)
            io.to_bigbank.data(i*6+j) := io.in_from_quant.mat(i*8+(j+1))
    for(i <- 0 to 7)
        io.to_smallbank(0).data(i) := io.in_from_quant.mat(i*8)
    for(i <- 0 to 7)
        io.to_smallbank(1).data(i) := io.in_from_quant.mat(i*8+7)
    io.to_banks := gen.output()
    when(io.flag_job){
        gen.set(io.job)
    }.otherwise{
        when(io.valid_in){
            io.valid_out := true.B
            gen.go()
        }
    }
}

class RealWriterBundle(val w: Int, val addr_w: Int, val h_w: Int, val c_w: Int, val id_w: Int) extends TransBundle{
    val in_from_quant = Input(new QuantedData(w))
    val job = Input(new RealWriteJobs(addr_w, h_w, c_w, id_w))
    val to_bigbank = Output(new BigBankWriteData(w))
    val to_smallbank = Output(Vec(2, new SmallBankWriteData(w)))
    val to_banks = Output(new AddressWriteGroup(addr_w, id_w))
}

class RealWriter(val w: Int, val addr_w: Int, val h_w: Int, val c_w: Int, val id_w: Int) extends Module{
    val io = IO(new RealWriterBundle(w, addr_w, h_w, c_w, id_w))

    val gen = RegInit(VecInit(Seq.fill(2)(0.U.asTypeOf(new WriteGroup(addr_w, h_w, c_w, id_w)))))
    val cnt_ic = RegInit(0.U.asTypeOf(ACounter(c_w)))
    val state = RegInit(0.U(1.W))

    for(i <- 0 to 7)

            io.to_bigbank.data(i*6+j) := io.in_from_quant.mat(i*8+(j+1))
    for(i <- 0 to 7)
        io.to_smallbank(0).data(i) := io.in_from_quant.mat(i*8)
    for(i <- 0 to 7)
        io.to_smallbank(1).data(i) := io.in_from_quant.mat(i*8+7)

    io.to_banks := gen(state).output()
    io.valid_out := false.B
    when(io.flag_job){
        gen(0).set(io.job.job(0))
        gen(1).set(io.job.job(1))
        cnt_ic.set(io.job.out_chan)
        state := 0.U
    }.otherwise{
        when(io.valid_in){
            io.valid_out := true.B
            gen(state).go()
            when(cnt_ic.inc()){
                state := ~state
            }
        }
    }

}

class CacheWriter(val w: Int, val para_num: Int) extends Module{
    val io = IO(new Bundle{
        val in_from_accumu = Input(Vec(para_num, new AccumuRawData(w)))
        val valid_in = Input(UInt(para_num.W))
        val out = Output(new AccumuRawData(w))
        val valid_out = Output(Bool())
    })

    val cache_valid = RegInit(0.U(para_num.W))
    cache_valid := (cache_valid >> 1.U)|io.valid_in

    val cache = RegInit(0.U.asTypeOf(Vec(para_num, new AccumuRawData(w))))
    for(i <- 0 to para_num-2)
        cache(i) := Mux(io.valid_in(i), io.in_from_accumu(i), cache(i+1))
    cache(para_num-1) := Mux(io.valid_in(para_num-1), io.in_from_accumu(para_num-1), 0.U.asTypeOf(cache(para_num-1)))

    io.out := cache(0)
    io.valid_out := cache_valid(0)

}

/*
class CacheWriter(val w: Int, val addr_w: Int, val num: Int, val h_w: Int, val c_w: Int, val id_w: Int) extends Module{
    val io = IO(new Bundle{
        val in_from_quant = Input(Vec(num, new QuantedData(w)))
        val valid_in = Input(Vec(num, Bool()))
        val valid_out = Output(Bool())
        val flag_job = Input(Bool())
        val jobs = Input(Vec(num, new WriteJobs(addr_w, h_w, c_w, id_w)))
        val to_bigbank = Output(new BigBankWriteData(w))
        val to_smallbank = Output(Vec(2, new SmallBankWriteData(w)))
        val to_banks = Output(new AddressWriteGroup(addr_w, id_w))
    })

    val writers = VecInit(Seq.fill(num)(Module(new Writer(w, addr_w, h_w, c_w, id_w)).io))
    val cache = RegInit(0.U.asTypeOf(VecInit(Seq.fill(num)(new QuantedData(w)))))
    val cache_valid = RegInit(0.U.asTypeOf(Vec(num, Bool())))
    val now = PriorityMux(Seq.tabulate(num){i => cache_valid(i)}, Seq.tabulate(num){i => i.U})

    io.to_bigbank := 0.U.asTypeOf(io.to_bigbank)
    io.to_smallbank := 0.U.asTypeOf(io.to_smallbank)
    io.to_banks := 0.U.asTypeOf(io.to_banks)
    io.valid_out := false.B
    for(i <- 0 to num-1)
        writers(i) := 0.U.asTypeOf(writers(i))

    when(io.flag_job){
        cache_valid := 0.U.asTypeOf(cache_valid)
        for(i <- 0 to num-1){
            writers(i).job := io.jobs(i)
            writers(i).flag_job := true.B
        }
    }.otherwise{
        for(i <- 0 to num-1)
            cache := io.in_from_quant(i)
        for(i <- 0 to num-1)
            cache_valid(i) := cache_valid(i)|io.valid_in(i)
        when(cache_valid(now)){
            writers(now).valid_in := true.B
            writers(now).in_from_quant := cache(now)
            io.to_bigbank := writers(now).to_bigbank
            io.to_smallbank := writers(now).to_smallbank
            io.to_banks := writers(now).to_banks
            cache_valid(now) := io.valid_in(now)
        }
    }

}*/


package final_project

import chisel3._
import chisel3.util._

class WRJob(addr_w: Int, h_w: Int, c_w: Int, id_w: Int) extends Bundle{
    val begin_addr = UInt(addr_w.W)
    val max_addr = UInt(addr_w.W)
    val min_addr = UInt(addr_w.W)
    val bank_id = UInt(id_w.W)
    val block_size = UInt(addr_w.W)
    val column_size = UInt(addr_w.W)
    val cnt_x_end = UInt(h_w.W)
    val cnt_y_end = UInt(h_w.W)
    val cnt_ic_end = UInt(c_w.W)
}

class WriteBigJob(addr_w: Int, h_w: Int, c_w: Int, id_w: Int) extends WRJob(addr_w, h_w, c_w, id_w){
    val a = UInt(c_w.W)
}

class WriteSmallJob(addr_w: Int, h_w: Int, c_w: Int, id_w: Int) extends WRJob(addr_w, h_w, c_w, id_w){
    val a = UInt(c_w.W)
    val ano_bank_id = UInt(id_w.W)
}

class ReadBigJob(addr_w: Int, h_w: Int, c_w: Int, id_w: Int) extends WRJob(addr_w, h_w, c_w, id_w){
    val cnt_loop_end = UInt((2*c_w).W)
    val begin_loop = UInt((2*c_w).W)
    val cnt_ups_end = UInt(1.W)
}

class ReadSmallJob(addr_w: Int, h_w: Int, c_w: Int, id_w: Int, big_w: Int) extends WRJob(addr_w, h_w, c_w, id_w){
    val ano_bank_id = UInt(id_w.W)
    val cnt_loop_end = UInt((2*c_w).W)
    val begin_loop = UInt((2*c_w).W)
    val cnt_invalid_end = UInt(big_w.W)
    val cnt_ups_end = UInt(1.W)
    val begin_swap = UInt(1.W)
}

class ReadJobs(addr_w: Int, h_w: Int, c_w: Int, id_w: Int, big_w: Int) extends Bundle{
    val big = new ReadBigJob(addr_w, h_w, c_w, id_w)
    val small = Vec(4, new ReadSmallJob(addr_w, h_w, c_w, id_w, big_w))
}

class WriteJobs(addr_w: Int, h_w: Int, c_w: Int, id_w: Int) extends Bundle{
    val big = new WriteBigJob(addr_w, h_w, c_w, id_w)
    val small = Vec(2, new WriteSmallJob(addr_w, h_w, c_w, id_w))
}

class RealWriteJobs(addr_w: Int, h_w: Int, c_w: Int, id_w: Int) extends Bundle{
    val job = Vec(2, new WriteJobs(addr_w, h_w, c_w, id_w))
    val out_chan = UInt(c_w.W)
}

object GenConvReadJob{
    def gen(addr_w: Int, h_w: Int, c_w: Int, id_w: Int, big_w: Int, loop_num: Int, begin_loop_num: Int, loop_h: Int,
            h: Int, w: Int, in_chan: Int, begin_addr: Int, min_addr: Int, max_addr: Int,
            small_begin_addr: Int, small_min_addr: Int, small_max_addr: Int,
            bank_id_big: Int, bank_id_small: Array[Int]): ReadJobs = {
        val ret = Wire(new ReadJobs(addr_w, h_w, c_w, id_w, big_w))
        ret.big.begin_addr := begin_addr.U
        ret.big.min_addr := min_addr.U
        ret.big.max_addr := max_addr.U
        ret.big.bank_id := bank_id_big.U
        ret.big.block_size := in_chan.U
        ret.big.column_size := (in_chan*h).U
        ret.big.cnt_x_end := (w-1).U
        ret.big.cnt_y_end := (h-1).U
        ret.big.cnt_ic_end := (in_chan-1).U
        ret.big.cnt_loop_end := (loop_num-1).U
        ret.big.cnt_ups_end := (loop_h-1).U
        ret.big.begin_loop := begin_loop_num.U
        for(i <- 0 to 3){
            ret.small(i).begin_addr := small_begin_addr.U
            ret.small(i).min_addr := small_min_addr.U
            ret.small(i).max_addr := small_max_addr.U
            ret.small(i).block_size := in_chan.U
            ret.small(i).column_size := (in_chan*h).U
            ret.small(i).cnt_x_end := (w/2-1).U
            ret.small(i).cnt_y_end := (h-1).U
            ret.small(i).cnt_ic_end := (in_chan-1).U
            ret.small(i).cnt_loop_end := (loop_num-1).U

            ret.small(i).cnt_ups_end := (loop_h-1).U
            ret.small(i).begin_loop := begin_loop_num.U
            if(i==0){
                ret.small(i).bank_id := bank_id_small(2).U
                ret.small(i).ano_bank_id := bank_id_small(0).U
            } else {
                ret.small(i).bank_id := bank_id_small(i).U
                ret.small(i).ano_bank_id := bank_id_small(i^2).U
            }
            if(i==3){
                ret.small(i).begin_swap := (1).U
            } else {
                ret.small(i).begin_swap := 0.U
            }

        }
        return ret
    }
}

object GenMaxpReadJob{
    def gen(addr_w: Int, h_w: Int, c_w: Int, id_w: Int, big_w: Int, loop_num: Int, begin_loop_num: Int, loop_h: Int,
            h: Int, w: Int, in_chan: Int, begin_addr: Int, min_addr: Int, max_addr: Int,
            small_begin_addr: Int, small_min_addr: Int, small_max_addr: Int,
            bank_id_big: Int, bank_id_small: Array[Int]): ReadJobs = {
        val ret = Wire(new ReadJobs(addr_w, h_w, c_w, id_w, big_w))
        ret := 0.U.asTypeOf(new ReadJobs(addr_w, h_w, c_w, id_w, big_w))
        ret.big.begin_addr := begin_addr.U
        ret.big.min_addr := min_addr.U
        ret.big.max_addr := max_addr.U
        ret.big.bank_id := bank_id_big.U
        ret.big.block_size := in_chan.U
        ret.big.column_size := (in_chan*h).U
        ret.big.cnt_x_end := (w/2-1).U
        ret.big.cnt_y_end := (h-1).U
        ret.big.cnt_ic_end := (in_chan-1).U
        for(i <- 0 to 1){
            ret.small(i).begin_addr := small_begin_addr.U
            ret.small(i).min_addr := small_min_addr.U
            ret.small(i).max_addr := small_max_addr.U
            if(i==0){
                ret.small(i).bank_id := bank_id_small(2).U
                ret.small(i).ano_bank_id := bank_id_small(0).U
            } else {
                ret.small(i).bank_id := bank_id_small(i).U
                ret.small(i).ano_bank_id := bank_id_small(i^2).U
            }

        }
        return ret
    }
}

object GenWriteConvJob{
    def gen(addr_w: Int, h_w: Int, c_w: Int, id_w: Int,
            h: Int, w: Int, out_chan: Int, begin_addr: Int, min_addr: Int, max_addr: Int,
            small_begin_addr: Int, small_min_addr: Int, small_max_addr: Int,
            bank_id_big: Int, bank_id_small: Array[Int], al: Int, ar: Int): WriteJobs = {
        val ret = Wire(new WriteJobs(addr_w, h_w, c_w, id_w))
        ret.big.begin_addr := begin_addr.U
        ret.big.min_addr := min_addr.U
        ret.big.max_addr := max_addr.U
        ret.big.bank_id := bank_id_big.U
        ret.big.block_size := out_chan.U
        ret.big.column_size := (out_chan*h).U
        ret.big.cnt_x_end := (w-1).U
        ret.big.cnt_y_end := (h-1).U
        ret.big.cnt_ic_end := (ar-al).U
        ret.big.a := (out_chan-ar+al).U
        for(i <- 0 to 1){
            ret.small(i).begin_addr := small_begin_addr.U
            ret.small(i).min_addr := small_min_addr.U
            ret.small(i).max_addr := small_max_addr.U
            ret.small(i).block_size := out_chan.U
            ret.small(i).column_size := (out_chan*h).U
            ret.small(i).cnt_x_end := (w/2-1).U
            ret.small(i).cnt_y_end := (h-1).U
            ret.small(i).cnt_ic_end := (ar-al).U
            ret.small(i).a := (out_chan-ar+al).U
            ret.small(i).bank_id := bank_id_small(i).U
            ret.small(i).ano_bank_id := bank_id_small(i+2).U
        }
        return ret
    }
}

class PackJob(addr_w: Int, h_w: Int, c_w: Int, big_w: Int) extends Bundle{
    val cnt_x_end = UInt(h_w.W)
    val cnt_y_end = UInt(h_w.W)
    val in_chan = UInt(c_w.W)
    val cnt_swap_end = UInt(big_w.W)
}

object GenPackJob{
    def gen(addr_w: Int, h_w: Int, c_w: Int, big_w: Int, cnt_x_end: Int, cnt_y_end: Int, cnt_swap_end: Int, in_chan: Int): PackJob = {
        val ret = Wire(new PackJob(addr_w, h_w, c_w, big_w))
        ret.cnt_x_end := cnt_x_end.U
        ret.cnt_y_end := cnt_y_end.U
        ret.in_chan := in_chan.U
        ret.cnt_swap_end := cnt_swap_end.U
        return ret
    }
}

ret.small(i).cnt_invalid_end := (if(i==0) (loop_num*in_chan*h*loop_h).U else 0.U)


// See README.md for license details.

package final_project

import chisel3._
import chisel3.tester._
import org.scalatest._
import chisel3.experimental.BundleLiterals._
import chisel3.iotesters._

class Calc8x8Tester(dut: Calc8x8) extends PeekPokeTester(dut) {

    //val rnd = new scala.util.Random
    val T = 10
    for(t <- 0 until T){
        var A = Array.tabulate(10, 10) {
            (i, j) => rnd.nextInt(5)-32768
        }
        for(i <- 0 to 63){
            poke(dut.io.input.mat(i), A(i/8+1)(i%8+1))
        }
        for(i <- 0 to 9) {
            poke(dut.io.input.up(i), A(0)(i))
        }
        for(i <- 0 to 7) {
            poke(dut.io.input.left(i), A(i+1)(0))
        }
        for(i <- 0 to 7) {
            poke(dut.io.input.right(i), A(i+1)(9))
        }
        for(i <- 0 to 9) {
            poke(dut.io.input.down(i), A(9)(i))
        }

        val B = Array.fill[Int](4, 3, 3)(0)
        for(t <- 0 to 3)
            for(i <- 0 to 8) B(t)(i/3)(i%3) = rnd.nextInt(32768)-16384
            //rnd.nextInt(5)-32768

        val std = Array.tabulate(4, 8, 8) {
            (t, i, j) => {
                var sum = BigInt(0)
                for(ii <- i to i+2)
                    for(jj <- j to j+2){
                        var a = A(ii)(jj)
                        var b = B(t)(ii-i)(jj-j)
                        sum = sum+BigInt(a)*b
                    }

                sum
            }
        }
        for(t <- 0 to 3)
            for(i <- 0 to 8)
                poke(dut.io.weight(t).real(i), B(t)(i/3)(i%3))

        poke(dut.io.flag, 3)
        poke(dut.io.mask, 15)
        poke(dut.io.valid_in, true)
        var stop = false
        while(!stop){
            step(1)
            poke(dut.io.valid_in, false)
            if(peek(dut.io.valid_out)==15){
                println("[valid: " + peek(dut.io.valid_out).toString + "] Result:")
                for(t <- 0 to 3)


package final_project

import chisel3._
import chisel3.util._

import chisel3._
import chisel3.tester._
import org.scalatest._
import chisel3.experimental.BundleLiterals._
import chisel3.iotesters._

class Global1(_w: Int, h_w: Int, c_w: Int, id_w: Int, big_w: Int, addr_w: Int, bias_w: Int, num: Int) extends Module{
    val io = IO(new Bundle{
        val rd_addr1 = Output(new AddressReadGroup(addr_w, id_w))
        val rd_addr2 = Output(new AddressReadGroup(addr_w, id_w))
        val rd_big = Input(Vec(2, new BigBankReadData(_w)))
        val rd_small = Input(Vec(2, Vec(4, new SmallBankReadData(_w))))
        val rd_valid_in = Input(Bool())
        val rd_valid_out = Output(Bool())
        val weight_addr = Output(UInt(addr_w.W))
        val weight_in = Input(Vec(9, SInt(16.W)))
        val wr_addr = Output(new AddressWriteGroup(addr_w, id_w))
        val bias_in = Input(SInt(bias_w.W))
        val bias_addr = Output(UInt(addr_w.W))
        val packed = Output(new PackedData(_w))
        //val wr_small = Output(Vec(4, UInt((_w*8).W)))
        //val wr_big = Output(Vec(1, UInt((_w*48).W)))
        val quanted = Output(new QuantedData(_w))
        val wr_valid_out = Output(Bool())
        val test_valid_out = Output(Bool())
        val to_bigbank = Output(new BigBankWriteData(_w))
        val to_smallbank = Output(Vec(2, new SmallBankWriteData(_w)))
    })

    val reader1 = Module(new GraphReader(addr_w, h_w, c_w, id_w, big_w)).io
    val reader2 = Module(new GraphReader(addr_w, h_w, c_w, id_w, big_w)).io
    val read_pack = Module(new PackReadData(_w, h_w, c_w, big_w)).io
    val read_switch = Module(new ReadSwitch(_w, h_w, c_w, 1)).io
    val read_weight = Module(new WeightReader(_w, addr_w)).io
    val calc8x8 = Module(new Calc8x8(_w, 1)).io
    val accu = Module(new Accumu(_w, addr_w, bias_w, 1)).io
    val quant = Module(new Quant(_w)).io
    val write_switch = Module(new WriteSwitch(_w, 3)).io
    val writer = Module(new RealWriter(_w, addr_w, h_w, c_w, id_w)).io
    val maxpool = Module(new Maxpool(_w)).io

    io.rd_addr1 := 0.U.asTypeOf(io.rd_addr1)
    io.rd_addr2 := 0.U.asTypeOf(io.rd_addr2)
    io.rd_valid_out := false.B
    io.weight_addr := 0.U.asTypeOf(io.weight_addr)

    io.rd_addr1 := reader1.to_banks
    reader1.flag_job := false.B
    reader1.valid_in := false.B
    reader1.job := 0.U.asTypeOf(reader1.job)
    reader1.job2 := 0.U.asTypeOf(reader1.job2)
    reader1.signal := false.B
    reader1.job_type := 0.U.asTypeOf(reader1.job_type)

    io.rd_valid_out := reader2.valid_out
    io.rd_addr2 := reader2.to_banks
    reader2.flag_job := false.B
    reader2.valid_in := false.B
    reader2.job := 0.U.asTypeOf(reader2.job)
    reader2.job_type := 0.U.asTypeOf(reader2.job_type)
    reader2.job2 := 0.U.asTypeOf(reader1.job2)
    reader2.signal := false.B

    read_weight.valid_in := false.B
    read_weight.addr_end := 0.U
    read_weight.begin_addr := 0.U
    io.weight_addr := read_weight.addr
    read_weight.flag_job := false.B

    read_pack.valid_in := io.rd_valid_in
    read_pack.flag_job := false.B
    read_pack.job := 0.U.asTypeOf(read_pack.job)
    read_pack.from_big := io.rd_big
    read_pack.from_small := io.rd_small

    read_switch.flag_job := false.B
    read_switch.job := 0.U.asTypeOf(read_switch.job)
    read_switch.valid_in := read_pack.valid_out
    read_switch.from := read_pack.output
    read_switch.from_weight(0) := io.weight_in
    read_switch.in_h := 0.U
    read_switch.in_chan := 0.U

    calc8x8.input := read_switch.to_calc8x8
    calc8x8.flag := 0.U.asTypeOf(calc8x8.flag)
    calc8x8.weight := read_switch.to_weight
    calc8x8.valid_in := read_switch.valid_out_calc8x8
    calc8x8.mask := 0.U
    calc8x8.multmap := read_switch.to_multmap

    accu.in_from_calc8x8 := calc8x8.output
    accu.valid_in := calc8x8.valid_out
    accu.flag_job := false.B
    accu.csum := 0.U
    accu.bias_begin_addr := 0.U
    accu.bias_end_addr := 0.U
    io.bias_addr := accu.bias_addr
    accu.bias_in(0) := io.bias_in
    accu.is_in_use := false.B

    quant.in_from_accumu := accu.result(0)
    quant.valid_in := accu.valid_out(0)
    quant.flag_job := false.B
    quant.quant_in := 0.U.asTypeOf(quant.quant_in)

    maxpool.from_big := io.rd_big
    for(t <- 0 to 1){
        maxpool.from_small(t)(0) := io.rd_small(t)(1)
        maxpool.from_small(t)(1) := io.rd_small(t)(2)
    }

    maxpool.valid_in := read_switch.valid_out_maxp

    //write_switch.valid_in(1) := false.B
    //write_switch.input(1) := 0.U.asTypeOf(write_switch.input(1))

    write_switch.valid_in(0) := quant.valid_out // conv
    write_switch.valid_in(1) := maxpool.valid_out // maxpool
    write_switch.valid_in(2) := read_switch.valid_out_copy // go through

    write_switch.valid_in(1) := false.B
    write_switch.valid_in(2) := false.B

    write_switch.input(0) := quant.result
    write_switch.input(1) := maxpool.result
    write_switch.input(2) := read_switch.to_copy

    writer.in_from_quant := write_switch.output
    writer.valid_in := write_switch.valid_out
    writer.flag_job := false.B
    writer.job := 0.U.asTypeOf(writer.job)
    io.wr_addr := writer.to_banks
    io.to_bigbank := writer.to_bigbank
    io.to_smallbank := writer.to_smallbank
    io.wr_valid_out := writer.valid_out

    /*
    io.wr_addr := 0.U.asTypeOf(io.wr_addr)
    io.to_bigbank := 0.U.asTypeOf(io.to_bigbank)
    io.to_smallbank := 0.U.asTypeOf(io.to_smallbank)
    io.wr_valid_out := false.B
    io.bias_addr := 0.U*/

    io.quanted := quant.result
    io.test_valid_out := calc8x8.valid_out
    // io.quanted := read_pack.output

    io.packed := 0.U.asTypeOf(io.packed)
    io.packed.mat := calc8x8.output(0).mat
    /*
        // w, h, wr_w, wr_h
        // in out chan
        // para
        // big addr
        // wr big addr
        // small addr
        // wr small addr
    */
    val paras = GenAllPara(addr_w, h_w, c_w, id_w, big_w,
        2, 2, 2, 2,
        2, 2,
        1, 1,
        0, 1023, 0,
        1020, 1023, 0,
        0, 1023, 0,
        1020, 1023, 0,
        0, 0, 17, 16
    )

    val state = RegInit(0.U(3.W))
    switch(state){
        is(0.U){
            paras.set_conv_reader1(reader1)
            paras.set_conv_reader2(reader2)
            paras.set_read_weight(read_weight)
            paras.set_conv_read_pack(read_pack)
            paras.set_conv_accu(accu)
            paras.set_quant(quant)
            paras.set_write(writer)
            paras.set_conv_read_switch(read_switch)

            calc8x8.flag := CalcType.calcConv
            calc8x8.mask := paras.mask.U
            state := 1.U
        }
        is(1.U){
            reader1.valid_in := true.B
            reader2.valid_in := true.B
            read_weight.valid_in := true.B
            calc8x8.flag := CalcType.calcConv
            calc8x8.mask := paras.mask.U
        }
    }
}

class ConvTester(dut: Global1, _w: Int, h_w: Int, c_w: Int, id_w: Int, big_w: Int, addr_w: Int, bias_w: Int, num: Int) extends PeekPokeTester(dut){

    var bank_id_big = Array(0, 1)
    var bank_id_small = Array(Array(3, 0, 1, 2), Array(7, 4, 5, 6))
    var (w, h) = (2, 2)
    var (in_chan, out_chan) = (2, 2)
    var A = Array.tabulate(8*8*w*h, in_chan){
        (i, ic) => i*in_chan+ic
    }
    var small_bank = Array.fill[Int](10, (h/2)*(w/2)*8*in_chan+64*in_chan)(0)
    var big_bank = Array.fill[Int](10, (h/2)*w*48*in_chan+64*in_chan)(0)
    for(i <- 0 to w-1)
        for(j <- 0 to h-1){
            var t1 = bank_id_small(j%2)(if(i%2==1) 3 else 1)

            for(ic <- 0 to in_chan-1){
                for(k <- 0 to 7){
                    small_bank(t1)((i/2)*(h/2)*8*in_chan+(j/2)*8*in_chan+ic*8+k) = A((i*h+j)*64+k*8)(ic)
                    small_bank(t2)((i/2)*(h/2)*8*in_chan+(j/2)*8*in_chan+ic*8+k) = A((i*h+j)*64+k*8+7)(ic)
                }
                for(k <- 0 to 7)
                    for(l <- 0 to 5){
                        big_bank(bank_id_big(j%2))(i*(h/2)*48*in_chan+(j/2)*48*in_chan+ic*48+(k*6+l)) = A((i*h+j)*64+k*8+l+1)(ic)
                    }
            }

        }
    var weight_bank = Array.fill[Int](9*in_chan*out_chan)(0)
    var bias_bank = Array.tabulate(out_chan){
        i => i
    }
    for(i <- 0 to out_chan-1)
        for(j <- 0 to in_chan-1){
            val B = Array.tabulate(3, 3) {
                (i, j) => i*3+j
            }
            for(k <- 0 to 8)
                weight_bank((i*in_chan+j)*9+k) = B(k/3)(k%3)
        }
    var read_reg = 0
    var T = 0
    var stop = 0
    while(T<20&&(stop==0)){
        T += 1
        val rd_bias = peek(dut.io.bias_addr).toInt
        poke(dut.io.bias_in, bias_bank(rd_bias))
        if(peek(dut.io.rd_valid_out)==1){
            poke(dut.io.rd_valid_in, true)
            val rd_w = (peek(dut.io.weight_addr)).toInt
            val vis = Array.fill[Int](10, 10)(0)
            val vis_id = Array.fill[Int](10, 10)(0)
            def get_addr(i: Int, j: Int): Int = {
                if(vis(i)(j)!=0) return vis(i)(j)
                //println("("+i+","+j+") addr: "+(if(i==1) (peek(dut.io.rd_addr1.addrs(j).addr)).toInt else (peek(dut.io.rd_addr2.addrs(j).addr)).toInt))
                vis(i)(j) = (if(i==1) (peek(dut.io.rd_addr1.addrs(j).addr)).toInt else (peek(dut.io.rd_addr2.addrs(j).addr)).toInt)
                return vis(i)(j)
            }
            def get_id(i: Int, j: Int): Int = {
                if(vis_id(i)(j)!=0) return vis_id(i)(j)
                //return 0
                //println("("+i+","+j+") addr: "+(if(i==1) (peek(dut.io.rd_addr1.addrs(j).bank_id)).toInt else (peek(dut.io.rd_addr2.addrs(j).bank_id)).toInt))
                vis_id(i)(j) = (if(i==1) (peek(dut.io.rd_addr1.addrs(j).bank_id)).toInt else (peek(dut.io.rd_addr2.addrs(j).bank_id)).toInt)
                return vis_id(i)(j)
            }
            for(i <- 1 to 4){
                for(j <- 0 to 7)
                    poke(dut.io.rd_small(0)(i-1).data(j), small_bank(get_id(1, i))(get_addr(1, i)*8+j))
            }
            for(j <- 0 to 47)
                poke(dut.io.rd_big(0).data(j), big_bank(get_id(1, 0))(get_addr(1, 0)*48+j))
            for(i <- 1 to 4){
                for(j <- 0 to 7)
                    poke(dut.io.rd_small(1)(i-1).data(j), small_bank(get_id(2, i))(get_addr(2, i)*8+j))
            }
            for(j <- 0 to 47)
                poke(dut.io.rd_big(1).data(j), big_bank(get_id(2, 0))(get_addr(2, 0)*48+j))
            for(i <- 0 to 8)
                poke(dut.io.weight_in(i), weight_bank(rd_w*9+i))

            var addr_str = ""

            for(i <- 0 to 4)
                addr_str += "("+get_addr(1, i)+", "+get_id(1, i)+"), "
            println("rd_addr1: "+addr_str)
            addr_str = ""
            for(i <- 0 to 4)
                addr_str += "("+get_addr(2, i)+", "+get_id(2, i)+"), "
            println("rd_addr2: "+addr_str)

            read_reg = 0
        } else {
            poke(dut.io.rd_valid_in, false)
        }
        if(peek(dut.io.wr_valid_out)==1){
            //stop = 1
            var addr_str = ""

            for(i <- 0 to 2)
                addr_str += "("+peek(dut.io.wr_addr.addrs(i).addr)+", "+peek(dut.io.wr_addr.addrs(i).bank_id)+"), "
            println("wr_addr1: "+addr_str)

            println("Written Result:")
            for(j <- 0 to 7){
                var str = peek(dut.io.to_smallbank(0).data(j)).toString()+" "
                for(k <- 0 to 5){
                    var x = peek(dut.io.to_bigbank.data(j*6+k)).toString()
                    //expect(result.mat(j*8+k), std(j)(k))
                    str = str+x+" "
                }
                str = str+peek(dut.io.to_smallbank(1).data(j)).toString()+" "
                println(str)
            }

        }
        if(peek(dut.io.test_valid_out)==1){
            println("Switch Result:")
            for(j <- 0 to 9){
                var str = ""
                if(j==0){
                    for(k <- 0 to 9){
                        var x = peek(dut.io.packed.up(k)).toString()
                        //expect(result.mat(j*8+k), std(j)(k))
                        str = str+x+" "
                    }
                } else if(j>=1&&j<=8) {
                    str = str+peek(dut.io.packed.left(j-1)).toString()+" "
                    for(k <- 0 to 7){
                        var x = peek(dut.io.packed.mat((j-1)*8+k)).toString()
                        //expect(result.mat(j*8+k), std(j)(k))
                        str = str+x+" "
                    }
                    str = str+peek(dut.io.packed.right(j-1)).toString()+" "
                } else {
                    for(k <- 0 to 9){
                        var x = peek(dut.io.packed.down(k)).toString()
                        //expect(result.mat(j*8+k), std(j)(k))
                        str = str+x+" "
                    }
                }

                println(str)
            }

        }
        step(1)
    }
}

class ConvSpec extends FlatSpec with Matchers {
    /*
    it should "Conv should pass" in {
        val w = 16
        val h_w = 16
        val c_w = 16
        val id_w = 16
        val addr_w = 16
        val bias_w = 36
        val big_w = 16
        val num = 1
        chisel3.iotesters.Driver(() => new Global1(w, h_w, c_w, id_w, big_w, addr_w, bias_w, num)) { c =>
            new ConvTester(c, w, h_w, c_w, id_w, big_w, addr_w, bias_w, num)
        } should be (true)
    }*/
}

var t2 = bank_id_small(j%2)(if(i%2==1) 0 else 2)


// See README.md for license details.

package final_project

import chisel3._
import chisel3.tester._
import org.scalatest._
import chisel3.experimental.BundleLiterals._
import chisel3.iotesters._

class LeakyReLUTester(dut: Calc8x8) extends PeekPokeTester(dut) {

    //val rnd = new scala.util.Random
    var A = Array.tabulate(8, 8) {
        (i, j) => rnd.nextInt(65536)-32768
    }
    for(i <- 0 to 63){
        poke(dut.io.input.mat(i), A(i/8)(i%8))
    }

    val std = Array.tabulate(8, 8) {
        (i, j) => {

        }
    }

    poke(dut.io.flag, 1)
    poke(dut.io.mask, 1)
    poke(dut.io.valid_in, true)
    var stop = false
    var fuck = 0
    while(!stop&&fuck<=20){
        step(1)
        fuck += 1
        println("[valid: " + peek(dut.io.valid_out).toString + "] Result:")
        if(peek(dut.io.valid_out)==1){
            for(j <- 0 to 7){
                var str = ""
                for(k <- 0 to 7){
                    var x = peek(dut.io.output(0).mat(j*8+k)).toString()
                    expect(dut.io.output(0).mat(j*8+k), std(j)(k))
                    str = str+x+" "
                }
                println(str)
            }

            stop = true
        }

    }

}

class LeakyReLUSpec extends FlatSpec with Matchers {

  it should "LeakyReLU should pass" in {
        chisel3.iotesters.Driver(() => new Calc8x8(16, 4)) { c =>
            new LeakyReLUTester(c)
        } should be (true)
    }
}


package final_project

import chisel3._
import chisel3.util._

import chisel3._
import chisel3.tester._
import org.scalatest._
import chisel3.experimental.BundleLiterals._
import chisel3.iotesters._

class Global2(_w: Int, h_w: Int, c_w: Int, id_w: Int, big_w: Int, addr_w: Int, bias_w: Int, num: Int) extends Module{
    val io = IO(new Bundle{
        val rd_addr1 = Output(new AddressReadGroup(addr_w, id_w))
        val rd_addr2 = Output(new AddressReadGroup(addr_w, id_w))
        val rd_big = Input(Vec(2, new BigBankReadData(_w)))
        val rd_small = Input(Vec(2, Vec(4, new SmallBankReadData(_w))))
        val rd_valid_in = Input(Bool())
        val rd_valid_out = Output(Bool())
        val weight_addr = Output(UInt(addr_w.W))
        val weight_in = Input(Vec(9, SInt(16.W)))
        val wr_addr = Output(new AddressWriteGroup(addr_w, id_w))
        val bias_in = Input(SInt(bias_w.W))
        val bias_addr = Output(UInt(addr_w.W))
        val packed = Output(new PackedData(_w))
        //val wr_small = Output(Vec(4, UInt((_w*8).W)))
        //val wr_big = Output(Vec(1, UInt((_w*48).W)))
        val quanted = Output(new QuantedData(_w))
        val wr_valid_out = Output(Bool())
        val test_valid_out = Output(Bool())
        val to_bigbank = Output(new BigBankWriteData(_w))
        val to_smallbank = Output(Vec(2, new SmallBankWriteData(_w)))
    })

    val reader1 = Module(new GraphReader(addr_w, h_w, c_w, id_w, big_w)).io
    val reader2 = Module(new GraphReader(addr_w, h_w, c_w, id_w, big_w)).io
    val read_pack = Module(new PackReadData(_w, h_w, c_w, big_w)).io
    val read_switch = Module(new ReadSwitch(_w, h_w, c_w, 1)).io
    val read_weight = Module(new WeightReader(_w, addr_w)).io
    val calc8x8 = Module(new Calc8x8(_w, 1)).io
    val accu = Module(new Accumu(_w, addr_w, bias_w, 1)).io
    val quant = Module(new Quant(_w)).io
    val write_switch = Module(new WriteSwitch(_w, 3)).io
    val writer = Module(new RealWriter(_w, addr_w, h_w, c_w, id_w)).io
    val maxpool = Module(new Maxpool(_w)).io

    io.rd_addr1 := 0.U.asTypeOf(io.rd_addr1)
    io.rd_addr2 := 0.U.asTypeOf(io.rd_addr2)
    io.rd_valid_out := false.B
    io.weight_addr := 0.U.asTypeOf(io.weight_addr)

    io.rd_addr1 := reader1.to_banks
    reader1.flag_job := false.B
    reader1.valid_in := false.B
    reader1.job := 0.U.asTypeOf(reader1.job)
    reader1.job_type := 0.U.asTypeOf(reader1.job_type)

    io.rd_valid_out := reader2.valid_out
    io.rd_addr2 := reader2.to_banks
    reader2.flag_job := false.B
    reader2.valid_in := false.B
    reader2.job := 0.U.asTypeOf(reader2.job)
    reader2.job_type := 0.U.asTypeOf(reader2.job_type)

    read_weight.valid_in := false.B
    read_weight.addr_end := 0.U
    read_weight.begin_addr := 0.U
    io.weight_addr := read_weight.addr
    read_weight.flag_job := false.B

    read_pack.valid_in := io.rd_valid_in
    read_pack.flag_job := false.B
    read_pack.job := 0.U.asTypeOf(read_pack.job)
    read_pack.from_big := io.rd_big
    read_pack.from_small := io.rd_small

    read_switch.flag_job := false.B
    read_switch.job := 0.U.asTypeOf(read_switch.job)
    read_switch.valid_in := read_pack.valid_out
    read_switch.from := read_pack.output
    read_switch.from_weight(0) := io.weight_in
    read_switch.in_h := 0.U
    read_switch.in_chan := 0.U

    calc8x8.input := read_switch.to_calc8x8
    calc8x8.flag := 0.U.asTypeOf(calc8x8.flag)
    calc8x8.weight := read_switch.to_weight
    calc8x8.valid_in := read_switch.valid_out_calc8x8
    calc8x8.mask := 0.U
    calc8x8.multmap := read_switch.to_multmap

    accu.in_from_calc8x8 := calc8x8.output
    accu.valid_in := calc8x8.valid_out
    accu.flag_job := false.B
    accu.csum := 0.U
    accu.bias_begin_addr := 0.U
    accu.bias_end_addr := 0.U
    io.bias_addr := accu.bias_addr
    accu.bias_in(0) := io.bias_in
    accu.is_in_use := false.B

    quant.in_from_accumu := accu.result(0)
    quant.valid_in := accu.valid_out(0)
    quant.flag_job := false.B
    quant.quant_in := 0.U.asTypeOf(quant.quant_in)

    maxpool.from_big := io.rd_big
    for(t <- 0 to 1){
        maxpool.from_small(t)(0) := io.rd_small(t)(0)
        maxpool.from_small(t)(1) := io.rd_small(t)(1)
    }

    maxpool.valid_in := read_switch.valid_out_maxp

    //write_switch.valid_in(1) := false.B
    //write_switch.input(1) := 0.U.asTypeOf(write_switch.input(1))

    write_switch.valid_in(0) := quant.valid_out // conv
    write_switch.valid_in(1) := maxpool.valid_out // maxpool
    write_switch.valid_in(2) := read_switch.valid_out_copy // go through

    write_switch.input(0) := quant.result
    write_switch.input(1) := maxpool.result
    write_switch.input(2) := read_switch.to_copy

    writer.in_from_quant := write_switch.output
    writer.valid_in := write_switch.valid_out
    writer.flag_job := false.B
    writer.job := 0.U.asTypeOf(writer.job)
    io.wr_addr := writer.to_banks
    io.to_bigbank := writer.to_bigbank
    io.to_smallbank := writer.to_smallbank
    io.wr_valid_out := writer.valid_out

    /*
    io.wr_addr := 0.U.asTypeOf(io.wr_addr)
    io.to_bigbank := 0.U.asTypeOf(io.to_bigbank)
    io.to_smallbank := 0.U.asTypeOf(io.to_smallbank)
    io.wr_valid_out := false.B
    io.bias_addr := 0.U*/

    io.quanted := quant.result
    io.test_valid_out := calc8x8.valid_out
    // io.quanted := read_pack.output

    io.packed := 0.U.asTypeOf(io.packed)
    io.packed.mat := calc8x8.output(0).mat
    /*
        // w, h, wr_w, wr_h
        // in out chan
        // para
        // big addr
        // wr big addr
        // small addr
        // wr small addr
    */
    val paras = GenAllPara(addr_w, h_w, c_w, id_w, big_w,
        2, 2, 2, 2,
        2, 2,
        1, 1,
        0, 1023, 0,
        1020, 1023, 0,
        0, 1023, 0,
        1020, 1023, 0,
        0, 0, 17, 16
    )

    val state = RegInit(0.U(3.W))
    switch(state){
        is(0.U){
            paras.set_maxp_reader1(reader1)
            paras.set_maxp_reader2(reader2)
            paras.set_write(writer)
            paras.set_maxpool_switch(read_switch)

            state := 1.U
        }
        is(1.U){
            reader1.valid_in := true.B
            reader2.valid_in := true.B

        }
    }
}

class MaxpoolTester(dut: Global2, _w: Int, h_w: Int, c_w: Int, id_w: Int, addr_w: Int, bias_w: Int, num: Int) extends PeekPokeTester(dut){

    var bank_id_big = Array(0, 1)
    var bank_id_small = Array(Array(3, 0, 1, 2), Array(7, 4, 5, 6))
    var (w, h) = (2, 2)
    var (in_chan, out_chan) = (2, 2)
    var A = Array.tabulate(8*8*w*h, in_chan){
        (i, ic) => i*in_chan+ic
    }
    var small_bank = Array.fill[Int](10, (h/2)*(w/2)*8*in_chan+64*in_chan)(0)
    var big_bank = Array.fill[Int](10, (h/2)*w*48*in_chan+64*in_chan)(0)
    for(i <- 0 to w-1)
        for(j <- 0 to h-1){
            var t1 = bank_id_small(j%2)(if(i%2==1) 3 else 1)
            var t2 = bank_id_small(j%2)(if(i%2==1) 0 else 2)
            for(ic <- 0 to in_chan-1){
                for(k <- 0 to 7){
                    small_bank(t1)((i/2)*(h/2)*8*in_chan+(j/2)*8*in_chan+ic*8+k) = A((i*h+j)*64+k*8)(ic)
                    small_bank(t2)((i/2)*(h/2)*8*in_chan+(j/2)*8*in_chan+ic*8+k) = A((i*h+j)*64+k*8+7)(ic)
                }
                for(k <- 0 to 7)
                    for(l <- 0 to 5){
                        big_bank(bank_id_big(j%2))(i*(h/2)*48*in_chan+(j/2)*48*in_chan+ic*48+(k*6+l)) = A((i*h+j)*64+k*8+l+1)(ic)
                    }
            }

        }
    var weight_bank = Array.fill[Int](9*in_chan*out_chan)(0)
    var bias_bank = Array.tabulate(out_chan){
        i => i
    }
    for(i <- 0 to out_chan-1)
        for(j <- 0 to in_chan-1){
            val B = Array.tabulate(3, 3) {
                (i, j) => i*3+j
            }
            for(k <- 0 to 8)
                weight_bank((i*in_chan+j)*9+k) = B(k/3)(k%3)
        }
    var read_reg = 0
    var T = 0
    var stop = 0

    val std = Array.tabulate(8, 8, 2){
        (i, j, t) => {
            var mx = 0
            for(x <- i*2 to i*2+1)
                for(y <- j*2 to j*2+1)
                    if(mx < A(((y/8)*h+(x/8))*64+(x%8)*8+(y%8))(t))
                        mx = A(((y/8)*h+(x/8))*64+(x%8)*8+(y%8))(t)
            mx
        }
    }

    var tt = 0
    while(T<20&&(stop==0)){
        T += 1
        val rd_bias = peek(dut.io.bias_addr).toInt
        poke(dut.io.bias_in, bias_bank(rd_bias))
        if(peek(dut.io.rd_valid_out)==1){
            poke(dut.io.rd_valid_in, true)
            val rd_w = (peek(dut.io.weight_addr)).toInt
            val vis = Array.fill[Int](10, 10)(0)
            val vis_id = Array.fill[Int](10, 10)(0)
            def get_addr(i: Int, j: Int): Int = {
                if(vis(i)(j)!=0) return vis(i)(j)
                //println("("+i+","+j+") addr: "+(if(i==1) (peek(dut.io.rd_addr1.addrs(j).addr)).toInt else (peek(dut.io.rd_addr2.addrs(j).addr)).toInt))

                return vis(i)(j)
            }
            def get_id(i: Int, j: Int): Int = {
                if(vis_id(i)(j)!=0) return vis_id(i)(j)
                //return 0
                //println("("+i+","+j+") addr: "+(if(i==1) (peek(dut.io.rd_addr1.addrs(j).bank_id)).toInt else (peek(dut.io.rd_addr2.addrs(j).bank_id)).toInt))
                vis_id(i)(j) = (if(i==1) (peek(dut.io.rd_addr1.addrs(j).bank_id)).toInt else (peek(dut.io.rd_addr2.addrs(j).bank_id)).toInt)
                return vis_id(i)(j)
            }
            for(i <- 1 to 4){
                for(j <- 0 to 7)
                    poke(dut.io.rd_small(0)(i-1).data(j), small_bank(get_id(1, i))(get_addr(1, i)*8+j))
            }
            for(j <- 0 to 47)
                poke(dut.io.rd_big(0).data(j), big_bank(get_id(1, 0))(get_addr(1, 0)*48+j))
            for(i <- 1 to 4){
                for(j <- 0 to 7)
                    poke(dut.io.rd_small(1)(i-1).data(j), small_bank(get_id(2, i))(get_addr(2, i)*8+j))
            }
            for(j <- 0 to 47)
                poke(dut.io.rd_big(1).data(j), big_bank(get_id(2, 0))(get_addr(2, 0)*48+j))
            for(i <- 0 to 8)
                poke(dut.io.weight_in(i), weight_bank(rd_w*9+i))

            var addr_str = ""

            for(i <- 0 to 4)
                addr_str += "("+get_addr(1, i)+", "+get_id(1, i)+"), "
            println("rd_addr1: "+addr_str)
            addr_str = ""
            for(i <- 0 to 4)
                addr_str += "("+get_addr(2, i)+", "+get_id(2, i)+"), "
            println("rd_addr2: "+addr_str)

            read_reg = 0
        } else {
            poke(dut.io.rd_valid_in, false)
        }
        if(peek(dut.io.wr_valid_out)==1){
            //stop = 1
            var addr_str = ""

            for(i <- 0 to 2)
                addr_str += "("+peek(dut.io.wr_addr.addrs(i).addr)+", "+peek(dut.io.wr_addr.addrs(i).bank_id)+"), "
            println("wr_addr1: "+addr_str)

            println("Written Result:")
            for(j <- 0 to 7){
                var str = peek(dut.io.to_smallbank(0).data(j)).toString()+" "
                expect(dut.io.to_smallbank(0).data(j), std(j)(0)(tt))
                for(k <- 0 to 5){
                    var x = peek(dut.io.to_bigbank.data(j*6+k)).toString()
                    expect(dut.io.to_bigbank.data(j*6+k), std(j)(k+1)(tt))
                    str = str+x+" "
                }
                str = str+peek(dut.io.to_smallbank(1).data(j)).toString()+" "
                expect(dut.io.to_smallbank(1).data(j), std(j)(7)(tt))
                println(str)
            }

            tt = tt+1
        }
        step(1)
    }
}

class MaxpoolSpec extends FlatSpec with Matchers {
    /*
    it should "Maxpool should pass" in {
        val w = 16
        val h_w = 16
        val c_w = 16
        val id_w = 16
        val addr_w = 16
        val bias_w = 36
        val big_w = 16
        val num = 1
        chisel3.iotesters.Driver(() => new Global2(w, h_w, c_w, id_w, big_w, addr_w, bias_w, num)) { c =>
            new MaxpoolTester(c, w, h_w, c_w, id_w, addr_w, bias_w, num)
        } should be (true)
    }*/
}

vis(i)(j) = (if(i==1) (peek(dut.io.rd_addr1.addrs(j).addr)).toInt else (peek(dut.io.rd_addr2.addrs(j).addr)).toInt)


// See README.md for license details.

package final_project

import chisel3._
import chisel3.util._
import chisel3.tester._
import org.scalatest._
import chisel3.experimental.BundleLiterals._
import chisel3.iotesters._

class ReadConvTestModule(addr_w: Int, h_w: Int, c_w: Int, id_w: Int, loop_num: Int, loop_h: Int,
            h: Int, w: Int, in_chan: Int, begin_addr: Int, min_addr: Int, max_addr: Int,
            small_begin_addr: Int, small_min_addr: Int, small_max_addr: Int,
            bank_id_big: Int, bank_id_small: Array[Int]) extends Module{
    val io = IO(new Bundle{
        val start = Input(Bool())
        val valid_out = Output(Bool())
        val to_banks = Output(new AddressReadGroup(addr_w, id_w))
    })
    val r = Module(new GraphReader(addr_w, h_w, c_w, id_w, 16))
    val state = RegInit(0.U(2.W))

    r.io.flag_job := false.B
    r.io.job := 0.U.asTypeOf(r.io.job)
    r.io.job2 := 0.U.asTypeOf(r.io.job2)
    r.io.signal := false.B
    r.io.job_type := 0.U.asTypeOf(r.io.job_type)
    r.io.valid_in := false.B
    io.to_banks := r.io.to_banks
    io.valid_out := r.io.valid_out
    when(io.start){
        state := 1.U
    }.otherwise{
        switch(state){
            is(1.U){
                r.io.flag_job := true.B
                r.io.job_type := ReadType.toConvOrCopy
                r.io.job := GenConvReadJob.gen(
                    addr_w, h_w, c_w, id_w, 16, loop_num, 0, loop_h,
                    h, w, in_chan, begin_addr, max_addr, min_addr, small_begin_addr, small_min_addr, small_max_addr,
                    bank_id_big, bank_id_small
                )
                r.io.valid_in := true.B
                state := 2.U
            }
            is(2.U){
                r.io.flag_job := false.B
                r.io.valid_in := true.B
                r.io.job_type := ReadType.toConvOrCopy
                state := 2.U
            }
        }
    }
}

class ReadConvTester(dut: ReadConvTestModule, std: Array[Array[Tuple2[Int, Int]]]) extends PeekPokeTester(dut) {
    poke(dut.io.start, true.B)
    step(1)
    poke(dut.io.start, false.B)
    step(1)
    for(i <- 0 to 10000){
        var str = ""
        str = str+"["+peek(dut.io.valid_out)+"]"
        for(j <- 0 to 4){
            str = str+"("+peek(dut.io.to_banks.addrs(j).bank_id)+","+peek(dut.io.to_banks.addrs(j).addr)+"), "
            if(std(i)(j)._1!=(-1)&&peek(dut.io.valid_out)==1){
                expect(dut.io.to_banks.addrs(j).bank_id, std(i)(j)._1)
                expect(dut.io.to_banks.addrs(j).addr, std(i)(j)._2)
            }
        }
        if(i%100==0)
            println(str)
        step(1)
    }
}

class ReadConvSpec extends FlatSpec with Matchers {

    it should "ReadConv should pass" in {
        var (h, w) = (4, 4)
        var in_chan = 16
        var (begin_addr, max_addr, min_addr) = (990, 1010, 12)
        var (small_begin_addr, small_max_addr, small_min_addr) = (540, 600, 29)
        var bank_id_big = 0
        var bank_id_small = Array(4, 1, 2, 3)
        var loop_num = 32
        var loop_h = 1

        var std = Array.tabulate(100000, 5){
            (i, j) => (0, 0)
        }

        def f(x: Int, L: Int, R: Int): Int = {
            return (x-L)%(R-L+1)+L
        }

        var x = 0
        for(i <- 0 to w-1)
            for(c <- 0 to loop_h-1)
                for(j <- 0 to h-1)
                    for(t <- 0 to loop_num-1)
                        for(k <- 0 to in_chan-1){
                            std(x)(1) = (
                                (if(i==0) -1 else bank_id_small(if(i%2==0) 0 else 2)),
                                f((((i-1)/2)*h+j)*in_chan+k+small_begin_addr, small_min_addr, small_max_addr)
                            )
                            std(x)(2) = (

                                f(((i/2)*(h)+j)*in_chan+k+small_begin_addr, small_min_addr, small_max_addr)
                            )
                            std(x)(3) = (
                                bank_id_small(if(i%2==0) 2 else 0),
                                f(((i/2)*(h)+j)*in_chan+k+small_begin_addr, small_min_addr, small_max_addr)
                            )
                            std(x)(4) = (
                                (if(i==w-1) -1 else bank_id_small(if(i%2==0) 3 else 1)),
                                f((((i+1)/2)*(h)+j)*in_chan+k+small_begin_addr, small_min_addr, small_max_addr)
                            )
                            std(x)(0) = (
                                bank_id_big,
                                f((i*(h)+j)*in_chan+k+begin_addr, min_addr, max_addr)
                            )
                            x += 1
                        }

        chisel3.iotesters.Driver(() => new ReadConvTestModule(10, 3, 6, 4, loop_num, loop_h,
            h, w, in_chan, begin_addr, max_addr, min_addr, small_begin_addr, small_min_addr, small_max_addr,
            bank_id_big, bank_id_small
        )) { c =>
            new ReadConvTester(c, std)
        } should be (true)
    }
}


// See README.md for license details.

package final_project

import chisel3._
import chisel3.util._
import chisel3.tester._
import org.scalatest._
import chisel3.experimental.BundleLiterals._
import chisel3.iotesters._

class WriteConvTestModule(addr_w: Int, h_w: Int, c_w: Int, id_w: Int,
            h: Int, w: Int, in_chan: Int, begin_addr: Int, min_addr: Int, max_addr: Int,
            small_begin_addr: Int, small_min_addr: Int, small_max_addr: Int,
            bank_id_big: Int, bank_id_small: Array[Int], al: Int, ar: Int) extends Module{
    val io = IO(new Bundle{
        val start = Input(Bool())
        val valid_out = Output(Bool())
        val to_banks = Output(new AddressWriteGroup(addr_w, id_w))
    })
    val r = Module(new Writer(16, addr_w, h_w, c_w, id_w))
    val state = RegInit(0.U(2.W))

    r.io.flag_job := false.B
    r.io.job := 0.U.asTypeOf(r.io.job)
    r.io.valid_in := false.B
    r.io.in_from_quant := 0.U.asTypeOf(r.io.in_from_quant)
    io.to_banks := r.io.to_banks
    io.valid_out := r.io.valid_out
    when(io.start){
        state := 1.U
    }.otherwise{
        switch(state){
            is(1.U){
                r.io.flag_job := true.B
                r.io.job := GenWriteConvJob.gen(
                    addr_w, h_w, c_w, id_w,
                    h, w, in_chan, begin_addr, max_addr, min_addr, small_begin_addr, small_min_addr, small_max_addr,
                    bank_id_big, bank_id_small, al, ar
                )
                r.io.valid_in := true.B
                state := 2.U
            }
            is(2.U){
                r.io.flag_job := false.B
                r.io.valid_in := true.B
                state := 2.U
            }
        }
    }
}

class WriteConvTester(dut: WriteConvTestModule, std: Array[Array[Tuple2[Int, Int]]]) extends PeekPokeTester(dut) {
    poke(dut.io.start, true.B)
    step(1)
    poke(dut.io.start, false.B)
    step(1)
    for(i <- 0 to 3000){
        var str = ""
        str = str+"["+peek(dut.io.valid_out)+"]"
        for(j <- 0 to 2){
            str = str+"("+peek(dut.io.to_banks.addrs(j).bank_id)+","+peek(dut.io.to_banks.addrs(j).addr)+"), "
            if(std(i)(j)._1!=(-1)&&peek(dut.io.valid_out)==1){
                expect(dut.io.to_banks.addrs(j).bank_id, std(i)(j)._1)
                expect(dut.io.to_banks.addrs(j).addr, std(i)(j)._2)
            }
        }
        //if(i%100==0)
        println(str)
        step(1)
    }
}

class WriteConvSpec extends FlatSpec with Matchers {

    it should "WriteConv should pass" in {
        var (h, w) = (8, 8)
        var in_chan = 64
        var (begin_addr, max_addr, min_addr) = (990, 1010, 12)
        var (small_begin_addr, small_max_addr, small_min_addr) = (540, 600, 29)
        var bank_id_big = 0
        var bank_id_small = Array(1, 2, 3, 4)
        var al = 32
        var ar = 63

        var std = Array.tabulate(10000, 3){
            (i, j) => (-1, 0)
        }

        def f(x: Int, L: Int, R: Int): Int = {
            return (x-L)%(R-L+1)+L
        }

        var x = 0
        for(i <- 0 to w-1)
            for(j <- 0 to h-1)
                for(k <- al to ar){
                    std(x)(1) = (

                        f((((i)/2)*h+j)*in_chan+k+small_begin_addr, small_min_addr, small_max_addr)
                    )
                    std(x)(2) = (
                        bank_id_small(if(i%2==0) 1 else 3),
                        f(((i/2)*(h)+j)*in_chan+k+small_begin_addr, small_min_addr, small_max_addr)
                    )
                    std(x)(0) = (
                        bank_id_big,
                        f((i*(h)+j)*in_chan+k+begin_addr, min_addr, max_addr)
                    )
                    x += 1
                }

        chisel3.iotesters.Driver(() => new WriteConvTestModule(10, 3, 6, 3,
            h, w, in_chan, f(begin_addr+al, min_addr, max_addr), max_addr, min_addr, f(small_begin_addr+al, small_min_addr, small_max_addr), small_min_addr, small_max_addr,
            bank_id_big, bank_id_small, al, ar
        )) { c =>
            new WriteConvTester(c, std)
        } should be (true)
    }
}


package magmasi.components

import chisel3._
import chisel3.util._

class AdderSwitch(DATA_TYPE: Int = 32, NUM_IN: Int = 2, SEL_IN: Int = 2,NUM_OUT: Int =2) extends Module {
  val io = IO(new Bundle {
    val i_valid = Input(Bool())
    val i_data_bus = Input(Vec(NUM_IN, UInt(DATA_TYPE.W)))
    val i_add_en = Input(UInt(3.W))
    val i_cmd = Input(UInt(3.W))
    val i_sel = Input(UInt(SEL_IN.W))
    val o_vn_valid = Output(UInt(NUM_OUT.W))
    val o_vn = Output(UInt((2 * DATA_TYPE).W))
    val o_adder = Output(Vec(NUM_OUT, UInt(DATA_TYPE.W)))
  })

  val clk = clock
  val rst = reset.asBool

  val reductionMux = Module(new ReductionMux(W = DATA_TYPE, NUM_IN = NUM_IN, SEL_IN = SEL_IN))
  reductionMux.io.i_data := io.i_data_bus
  reductionMux.io.i_sel := io.i_sel

  val adder32 = Module(new SimpleAdder)
  adder32.io.A := reductionMux.io.o_data(1)
  adder32.io.B := reductionMux.io.o_data(0)

  val r_adder = RegInit(0.U((2 * DATA_TYPE).W))
  val r_vn = RegInit(0.U((2 * DATA_TYPE).W))
  val r_vn_valid = RegInit(0.U(2.W))
  val r_add_en = RegNext(io.i_add_en, init = 0.U)

  when (rst) {
    r_adder := 0.U
    r_vn := 0.U
    r_vn_valid := 0.U
  } .elsewhen (io.i_valid) {
    switch (io.i_cmd) {
      is ("b001".U) {
        r_adder := Cat(reductionMux.io.o_data)
        r_vn_valid := 0.U
      }
      is ("b011".U) {
        r_adder := Cat(0.U((DATA_TYPE - 1).W), reductionMux.io.o_data(1))
        r_vn := io.i_data_bus(0)
        r_vn_valid := 1.U
      }
      is ("b100".U) {
        r_adder := Cat(reductionMux.io.o_data(0), 0.U((DATA_TYPE - 1).W))
        r_vn := io.i_data_bus(1)
        r_vn_valid := 2.U
      }

    }
  }

  when (rst) {
    io.o_adder := VecInit(Seq.fill(2)(0.U(DATA_TYPE.W)))
    io.o_vn := 0.U
    io.o_vn_valid := 0.U
  } .otherwise {
    when (!r_add_en) {
      io.o_adder := VecInit(Seq(r_adder(DATA_TYPE - 1, 0), r_adder(2 * DATA_TYPE - 1, DATA_TYPE)))
    } .otherwise {
      io.o_adder := VecInit(Seq.fill(2)(adder32.io.O))
    }
    io.o_vn := r_vn
    io.o_vn_valid := r_vn_valid
  }
}

      is ("b101".U) {
        r_vn := Cat(reductionMux.io.o_data)
        r_vn_valid := 3.U
      }


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

import chisel3._
import chisel3.util._

/*
Sample Indexing Diagram with 4 PES example

Table:
 O   --> switch (regular)
 $	 --> input switch
 #   --> output switch

-----------------------------------------------------------------------------------------------
Data IO Diagram:

r_data_bus_ff[Element 0] -->  $     0     0     0     0     0     #  --> w_dist_bus[Element 0]

r_data_bus_ff[Element 1] -->  $     0     0     0     0     0     #  --> w_dist_bus[Element 1]

r_data_bus_ff[Element 2] -->  $     0     0     0     0     0     #  --> w_dist_bus[Element 2]

r_data_bus_ff[Element 3] -->  $     0     0     0     0     0     #  --> w_dist_bus[Element 3]

-----------------------------------------------------------------------------------------------
Horizontal Internal Wires Diagram: (between each switch)

$ w_internal[0]  O w_internal[2]  O w_internal[4]  O w_internal[6] #

$ w_internal[8] O w_internal[10] O w_internal[12] O w_internal[14] #

$ w_internal[16] O w_internal[18] O w_internal[20] O w_internal[22] #

$ w_internal[24] O w_internal[26] O w_internal[28] O w_internal[30] #

-----------------------------------------------------------------------------------------------
Diagonal Internal Wires Diagram: (between each switch)

$ w_internal[1]  O w_internal[3]  O w_internal[5]  O w_internal[7] #

$ w_internal[9] O w_internal[11] O w_internal[13] O w_internal[15] #

$ w_internal[17] O w_internal[19] O w_internal[21] O w_internal[23] #

$ w_internal[25] O w_internal[27] O w_internal[29] O w_internal[31] #

-----------------------------------------------------------------------------------------------
Mux Select Signals Diagram (inputs to each switch)
	* input switch does not require any control signals --> value will go to both horizontal and diagonal
	* output switch only requires one control bit

NA  r_mux_bus_ff[0,1]      r_mux_bus_ff[2,3]      r_mux_bus_ff[4,5]     r_mux_bus_ff[24]
NA  r_mux_bus_ff[6,7]      r_mux_bus_ff[8,9]      r_mux_bus_ff[10,11]   r_mux_bus_ff[25]
NA  r_mux_bus_ff[12,13]    r_mux_bus_ff[14,15]    r_mux_bus_ff[16,17]	  r_mux_bus_ff[26]
NA  r_mux_bus_ff[18,19]	   r_mux_bus_ff[20,21]	  r_mux_bus_ff[22,23]   r_mux_bus_ff[27]

-----------------------------------------------------------------------------------------------
*/

class Benes(implicit val config: MagmasiConfig) extends Module {
  val LEVELS   : Int = (2 * (math.log(config.NUM_PES) / math.log(2))).toInt + 1
  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(2 * (LEVELS - 2) * config.NUM_PES + config.NUM_PES, Bool()))
    val o_dist_bus2  = Output(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))
    val o_dist_bus1  = Output(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))

  })
  io.o_dist_bus1 := io.i_data_bus1
  val clk = clock
  val rst = reset.asBool

  val r_data_bus_ff   = RegInit(VecInit(Seq.fill(config.NUM_PES)(0.U(config.DATA_TYPE.W))))
  val r_mux_bus_ff    = RegInit(VecInit(Seq.fill(2 * (LEVELS - 2) * config.NUM_PES + config.NUM_PES)(false.B)))
  val w_dist_bus      = Wire(Vec(config.NUM_PES, UInt(config.DATA_TYPE.W)))
  val w_internal      = Wire(Vec(2 * config.NUM_PES * (LEVELS - 1), UInt(config.DATA_TYPE.W))) // 8 x 2 = 16

  dontTouch(w_internal)

  r_data_bus_ff := Mux(rst, VecInit(Seq.fill(config.NUM_PES)(0.U(config.DATA_TYPE.W))), io.i_data_bus2)
  r_mux_bus_ff  := Mux(rst, VecInit(Seq.fill(2 * (LEVELS - 2) * config.NUM_PES + config.NUM_PES)(false.B)), io.i_mux_bus)
  io.o_dist_bus2 := Mux(rst, VecInit(Seq.fill(config.NUM_PES)(0.U(config.DATA_TYPE.W))), w_dist_bus)

  for (i <- 0 until config.NUM_PES) {

    val in_switch = Module(new InputSwitch(config.DATA_TYPE))
    in_switch.io.in := r_data_bus_ff(i)
    w_internal(2 * i * (LEVELS - 1))      := in_switch.io.y
    w_internal(2 * i * (LEVELS - 1) + 1)  := in_switch.io.z

  }

  for (i <- 0 until config.NUM_PES) {

    val out_switch = Module(new OutputSwitch(config.DATA_TYPE))
    out_switch.io.in0 := w_internal(2 * i * (LEVELS - 1) + (2 * (LEVELS - 2)))

    out_switch.io.sel := r_mux_bus_ff(2 * config.NUM_PES * (LEVELS - 2) + i)
    w_dist_bus(i)     := out_switch.io.y

  }

  for (i <- 0 until config.NUM_PES) {
    for (j <- 1 until (LEVELS - 1)) {
      val imm_switch = Module(new Switch(config.DATA_TYPE))
      imm_switch.io.in0 := w_internal(2 * i * (LEVELS - 1) + 2 * (j - 1))
      if (j <= (LEVELS - 1) / 2) {
        if (i % math.pow(2, j).toInt < math.pow(2, j - 1).toInt) {
          imm_switch.io.in1 := w_internal(2 * (i + math.pow(2, j - 1).toInt) * (LEVELS - 1) + 2 * (j - 1) + 1)
        } else {
          imm_switch.io.in1 := w_internal(2 * (i - math.pow(2, j - 1).toInt) * (LEVELS - 1) + 2 * (j - 1) + 1)
        }
      } else {
        if (i % math.pow(2, LEVELS - j).toInt < math.pow(2, LEVELS - j - 1).toInt) {
          imm_switch.io.in1 := w_internal(2 * (i + math.pow(2, LEVELS - j - 1).toInt) * (LEVELS - 1) + 2 * (j - 1) + 1)
        } else {
          imm_switch.io.in1 := w_internal(2 * (i - math.pow(2, LEVELS - j - 1).toInt) * (LEVELS - 1) + 2 * (j - 1) + 1)
        }
      }
      imm_switch.io.sel0 := r_mux_bus_ff(2 * (LEVELS - 2) * i + 2 * (j - 1))
      imm_switch.io.sel1 := r_mux_bus_ff(2 * (LEVELS - 2) * i + 2 * (j - 1) + 1)
      w_internal(2 * i * (LEVELS - 1) + 2 * j)      := imm_switch.io.y
      w_internal(2 * i * (LEVELS - 1) + 2 * j + 1)  := imm_switch.io.z
    }
  }
}

out_switch.io.in1 := w_internal(2 * (if(i % 2 == 0) i + 1 else i - 1) * (LEVELS - 1) + (2 * (LEVELS - 2)) + 1)


package magmasi.components

import chisel3._
import chisel3.util._

class Bitmap(implicit val config: MagmasiConfig) extends Module{
    val io = IO(new Bundle {
        val mat1 =Input(Vec(config.rowsA, Vec(config.colsA, UInt(config.DATA_TYPE.W))))
        val mat2 = Input(Vec(config.rowsB, Vec(config.colsB, UInt(config.DATA_TYPE.W))))
         val bitmap1 =Output(Vec(config.rowsA, Vec(config.colsA, UInt(config.DATA_TYPE.W))))
        val bitmap2 = Output(Vec(config.rowsB, Vec(config.colsB, UInt(config.DATA_TYPE.W))))
    })

    val matReg1 = RegInit(VecInit(Seq.fill(config.rowsA)(VecInit(Seq.fill(config.colsA)(0.U(config.DATA_TYPE.W))))))
    val matReg2 = RegInit(VecInit(Seq.fill(config.rowsB)(VecInit(Seq.fill(config.colsB)(0.U(config.DATA_TYPE.W))))))
    io.bitmap1 := matReg1
    io.bitmap2 := matReg2

    val i = RegInit(0.U(log2Ceil(config.rowsA).W))
    val j = RegInit(0.U(log2Ceil(config.colsB).W))

    when (j < config.colsB.U) {
                    when (io.mat1(i)(j) =/= 0.U){
                matReg1(i)(j) := 1.U
            }.otherwise{
                matReg1(i)(j) := 0.U
            }

            when (io.mat2(i)(j) =/= 0.U){
                matReg2(i)(j) := 1.U
            }.otherwise{


package magmasi.components
import chisel3._
import chisel3.util._
class EdgeAdderSwitch(DATA_TYPE: Int = 32, NUM_IN: Int = 2, SEL_IN: Int = 2,NUM_OUT: Int = 2) extends Module {
  val io = IO(new Bundle {
    val i_valid = Input(Bool())
    val i_data_bus = Input(Vec(NUM_IN, UInt(64.W)))
    val i_add_en = Input(UInt(3.W))
    val i_cmd = Input(UInt(5.W))
    val i_sel = Input(UInt(SEL_IN.W))
    val o_vn = Output(UInt((2 * DATA_TYPE).W))
    val o_vn_valid = Output(UInt(2.W))
    val o_adder = Output(UInt(DATA_TYPE.W))
  })

  val clk = clock
  val rst = reset.asBool

  val reductionMux = Module(new ReductionMux(W = DATA_TYPE, NUM_IN = NUM_IN, SEL_IN = SEL_IN))
  reductionMux.io.i_data := io.i_data_bus
  reductionMux.io.i_sel := io.i_sel

  val adder32 = Module(new SimpleAdder)
  adder32.io.A := reductionMux.io.o_data(1)
  adder32.io.B := reductionMux.io.o_data(0)

  val r_valid = RegNext(io.i_valid, init = false.B)
  val r_adder = RegInit(0.U(DATA_TYPE.W))
  val r_vn = RegInit(0.U((2 * DATA_TYPE).W))
  val r_vn_valid = RegInit(0.U(2.W))
  val r_add_en = RegNext(io.i_add_en, init = 0.U)

        r_vn_valid := 1.U
      }
      is ("b100".U) {
        r_adder := reductionMux.io.o_data(0)
        r_vn := Cat(0.U(DATA_TYPE.W), io.i_data_bus(NUM_IN - 1))
        r_vn_valid := 2.U
      }
      is ("b101".U) {
        r_vn := Cat(reductionMux.io.o_data)
        r_vn_valid := 3.U
      }
      is (0.U) {
        r_vn_valid := 0.U
      }
    }
  }

  when (rst) {
    io.o_adder := 0.U
    io.o_vn := 0.U
    io.o_vn_valid := 0.U
  } .otherwise {
    when (!r_add_en) {
      io.o_adder := r_adder
    } .otherwise {
      io.o_adder := adder32.io.O
    }
    io.o_vn := r_vn
    io.o_vn_valid := r_vn_valid
  }
}

  when (rst) {
    r_adder := 0.U
    r_vn := 0.U
    r_vn_valid := 0.U
  } .elsewhen (r_valid) {