Chisel MuxN generator toy

Mux4.v

module MyMux(
  input         clock,
  input         reset,
  input  [1:0]  io_sel,
  input  [15:0] io_in_0,
  input  [15:0] io_in_1,
  input  [15:0] io_in_2,
  input  [15:0] io_in_3,
  output [15:0] io_out
);
  wire [15:0] Mux2_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_in1; // @[MyMux.scala 18:21]
  wire  Mux2_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_out; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_1_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_1_in1; // @[MyMux.scala 18:21]
  wire  Mux2_1_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_1_out; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_2_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_2_in1; // @[MyMux.scala 18:21]
  wire  Mux2_2_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_2_out; // @[MyMux.scala 18:21]
  Mux2 Mux2 ( // @[MyMux.scala 18:21]
    .in0(Mux2_in0),
    .in1(Mux2_in1),
    .sel(Mux2_sel),
    .out(Mux2_out)
  );
  Mux2 Mux2_1 ( // @[MyMux.scala 18:21]
    .in0(Mux2_1_in0),
    .in1(Mux2_1_in1),
    .sel(Mux2_1_sel),
    .out(Mux2_1_out)
  );
  Mux2 Mux2_2 ( // @[MyMux.scala 18:21]
    .in0(Mux2_2_in0),
    .in1(Mux2_2_in1),
    .sel(Mux2_2_sel),
    .out(Mux2_2_out)
  );
  assign io_out = Mux2_2_out; // @[MyMux.scala 48:10]
  assign Mux2_in0 = io_in_0; // @[MyMux.scala 19:16]
  assign Mux2_in1 = io_in_1; // @[MyMux.scala 20:16]
  assign Mux2_sel = io_sel[0]; // @[MyMux.scala 21:16]
  assign Mux2_1_in0 = io_in_2; // @[MyMux.scala 19:16]
  assign Mux2_1_in1 = io_in_3; // @[MyMux.scala 20:16]
  assign Mux2_1_sel = io_sel[0]; // @[MyMux.scala 21:16]
  assign Mux2_2_in0 = Mux2_out; // @[MyMux.scala 19:16]
  assign Mux2_2_in1 = Mux2_1_out; // @[MyMux.scala 20:16]
  assign Mux2_2_sel = io_sel[1]; // @[MyMux.scala 21:16]
endmodule

Mux6.v

module MyMux(
  input         clock,
  input         reset,
  input  [2:0]  io_sel,
  input  [15:0] io_in_0,
  input  [15:0] io_in_1,
  input  [15:0] io_in_2,
  input  [15:0] io_in_3,
  input  [15:0] io_in_4,
  input  [15:0] io_in_5,
  output [15:0] io_out
);
  wire [15:0] Mux2_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_in1; // @[MyMux.scala 18:21]
  wire  Mux2_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_out; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_1_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_1_in1; // @[MyMux.scala 18:21]
  wire  Mux2_1_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_1_out; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_2_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_2_in1; // @[MyMux.scala 18:21]
  wire  Mux2_2_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_2_out; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_3_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_3_in1; // @[MyMux.scala 18:21]
  wire  Mux2_3_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_3_out; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_4_in0; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_4_in1; // @[MyMux.scala 18:21]
  wire  Mux2_4_sel; // @[MyMux.scala 18:21]
  wire [15:0] Mux2_4_out; // @[MyMux.scala 18:21]
  Mux2 Mux2 ( // @[MyMux.scala 18:21]
    .in0(Mux2_in0),
    .in1(Mux2_in1),
    .sel(Mux2_sel),
    .out(Mux2_out)
  );
  Mux2 Mux2_1 ( // @[MyMux.scala 18:21]
    .in0(Mux2_1_in0),
    .in1(Mux2_1_in1),
    .sel(Mux2_1_sel),
    .out(Mux2_1_out)
  );
  Mux2 Mux2_2 ( // @[MyMux.scala 18:21]
    .in0(Mux2_2_in0),
    .in1(Mux2_2_in1),
    .sel(Mux2_2_sel),
    .out(Mux2_2_out)
  );
  Mux2 Mux2_3 ( // @[MyMux.scala 18:21]
    .in0(Mux2_3_in0),
    .in1(Mux2_3_in1),
    .sel(Mux2_3_sel),
    .out(Mux2_3_out)
  );
  Mux2 Mux2_4 ( // @[MyMux.scala 18:21]
    .in0(Mux2_4_in0),
    .in1(Mux2_4_in1),
    .sel(Mux2_4_sel),
    .out(Mux2_4_out)
  );
  assign io_out = Mux2_4_out; // @[MyMux.scala 48:10]
  assign Mux2_in0 = io_in_0; // @[MyMux.scala 19:16]
  assign Mux2_in1 = io_in_1; // @[MyMux.scala 20:16]
  assign Mux2_sel = io_sel[0]; // @[MyMux.scala 21:16]
  assign Mux2_1_in0 = io_in_2; // @[MyMux.scala 19:16]
  assign Mux2_1_in1 = io_in_3; // @[MyMux.scala 20:16]
  assign Mux2_1_sel = io_sel[0]; // @[MyMux.scala 21:16]
  assign Mux2_2_in0 = io_in_4; // @[MyMux.scala 19:16]
  assign Mux2_2_in1 = io_in_5; // @[MyMux.scala 20:16]
  assign Mux2_2_sel = io_sel[0]; // @[MyMux.scala 21:16]
  assign Mux2_3_in0 = Mux2_out; // @[MyMux.scala 19:16]
  assign Mux2_3_in1 = Mux2_1_out; // @[MyMux.scala 20:16]
  assign Mux2_3_sel = io_sel[1]; // @[MyMux.scala 21:16]
  assign Mux2_4_in0 = Mux2_3_out; // @[MyMux.scala 19:16]
  assign Mux2_4_in1 = Mux2_2_out; // @[MyMux.scala 20:16]
  assign Mux2_4_sel = io_sel[2]; // @[MyMux.scala 21:16]
endmodule

MuxN.scala

package mux

import chisel3._
import chisel3.util._

class Mux2[T<:Data](gen: => T) extends BlackBox with HasBlackBoxResource {
  val io = IO(new Bundle {
    val in0 = Input(gen)
    val in1 = Input(gen)
    val sel = Input(Bool())
    val out = Output(gen)
  })
  addResource("/Mux2.v")
}

object Mux2 {
  def apply[T<:Data](in0: T, in1: T, sel: Bool): T = {
    val mux = Module(new Mux2(in0.cloneType))
    mux.io.in0 := in0
    mux.io.in1 := in1
    mux.io.sel := sel
    mux.io.out
  }
}

// Inspired by the chisel-builtin Vec implementation
object MuxN {
  def apply[T<:Data](v: Vec[T], sel: UInt): T = {
    assert(log2Ceil(v.length) <= sel.getWidth, "Vec length is too large to be selected")
    var cntLayer = 0
    var curLayer = v
    while (curLayer.length > 1) {
      curLayer = VecInit(curLayer.grouped(2).map( x =>
        if (x.length == 1) x(0) else Mux2(x(0), x(1), sel(cntLayer))
      ).toSeq)
      cntLayer = cntLayer + 1
    }
    curLayer(0)
  }
}

class MyMux(n: Int) extends Module {
  val io = IO(new Bundle{
    val sel = Input(UInt(log2Ceil(n).W))
    val in  = Input(Vec(n, UInt(16.W)))
    val out = Output(UInt(16.W))
  })
  io.out := MuxN(io.in, io.sel)
}


object MuxMain extends App {
  Driver.execute(args, () => new MyMux(6))
}