import "primitives/std.lib";

component mac_pe(top: 32, left: 32) -> (down: 32, right: 32, out: 32) {
  cells {
    // Storage
    acc = prim std_reg(32);
    mul_reg = prim std_reg(32);
    // Computation
    add = prim std_add(32);
    mul = prim std_mult_pipe(32);
  }

  wires {

    group do_mul {
      mul.left = top;
      mul.right = left;
      mul.go = !mul.done ? 1'd1;
      mul_reg.in = mul.done ? mul.out;
      mul_reg.write_en = mul.done ? 1'd1;
      do_mul[done] = mul_reg.done;
    }

    group do_add {
      add.left = acc.out;
      add.right = mul_reg.out;
      acc.in = add.out;
      acc.write_en = 1'd1;
      do_add[done] = acc.done;
    }

    out = acc.out;
    down = top;
    right = left;
  }

  control {
    seq { do_mul; do_add; }
  }
}

component main() -> () {
    cells {
      t0_idx = prim std_reg(2);
      t0_add = prim std_add(2);
      t0 = prim std_mem_d1(32, 3, 2);
      l0_idx = prim std_reg(2);
      l0_add = prim std_add(2);
      l0 = prim std_mem_d1(32, 3, 2);
      @external(1) out_mem = prim std_mem_d2(32, 1, 1, 1, 1);
      pe_00 = mac_pe;
      top_00_read = prim std_reg(32);
      left_00_read = prim std_reg(32);
    }
    wires {

      group t0_idx_init {
          t0_idx.in = 2'd3;
          t0_idx.write_en = 1'd1;
          t0_idx_init[done] = t0_idx.done;
      }

      group t0_idx_update {
          t0_add.left = 2'd1;
          t0_add.right = t0_idx.out;
          t0_idx.in = t0_add.out;
          t0_idx.write_en = 1'd1;
          t0_idx_update[done] = t0_idx.done;
      }

      group t0_move {
          t0.addr0 = t0_idx.out;
          top_00_read.in = t0.read_data;
          top_00_read.write_en = 1'd1;
          t0_move[done] = top_00_read.done;
      }

      group l0_idx_init {
          l0_idx.in = 2'd3;
          l0_idx.write_en = 1'd1;
          l0_idx_init[done] = l0_idx.done;
      }

      group l0_idx_update {
          l0_add.left = 2'd1;
          l0_add.right = l0_idx.out;
          l0_idx.in = l0_add.out;
          l0_idx.write_en = 1'd1;
          l0_idx_update[done] = l0_idx.done;
      }

      group l0_move {
          l0.addr0 = l0_idx.out;
          left_00_read.in = l0.read_data;
          left_00_read.write_en = 1'd1;
          l0_move[done] = left_00_read.done;
      }

      group pe_00_compute {

        pe_00.go = !pe_00.done ? 1'd1;
        pe_00.top = top_00_read.out;
        pe_00.left = left_00_read.out;

        pe_00_compute[done] = pe_00.done ? 1'd1;
      }


      group pe_00_out_write {
          out_mem.addr0 = 1'd0;
          out_mem.addr1 = 1'd0;
          out_mem.write_data = pe_00.out;
          out_mem.write_en = 1'd1;
          pe_00_out_write[done] = out_mem.done;
      }
    }
    control {

      seq {

        par {
            t0_idx_init; l0_idx_init;
        }
        par {
            t0_idx_update; l0_idx_update;
        }
        par {
            t0_move; l0_move;
        }
        par {
            t0_idx_update; l0_idx_update; pe_00_compute;
        }
        par {
            t0_move; l0_move;
        }
        par {
            t0_idx_update; l0_idx_update; pe_00_compute;
        }
        par {
            t0_move; l0_move;
        }
        par {
            pe_00_compute;
        }
        seq {
            pe_00_out_write;
        }
      }
    }
}