Mux Tree Generation

[escou64]

Hi,

I need to constraint the way multiplexer trees are generated. Currently, with Chisel 6.5.0, the following code:

class MuxTree(nBit: Int) extends Module {
  val io = IO(new Bundle {
    val i_slct = Input(UInt(2.W))
    val i_in = Input(Vec(4, UInt(nBit.W)))   

    val o_out = Output(UInt(nBit.W))  
  })  

  io.o_out := io.i_in(io.i_slct)
}

leads to the following SystemVerilog:

// Generated by CIRCT firtool-1.62.0
module MuxTree(
  input        clock,
               reset,
  input  [1:0] io_i_slct,
  input  [3:0] io_i_in_0,
               io_i_in_1,
               io_i_in_2,
               io_i_in_3,
  output [3:0] io_o_out
);

  wire [3:0][3:0] _GEN = {{io_i_in_3}, {io_i_in_2}, {io_i_in_1}, {io_i_in_0}};
  assign io_o_out = _GEN[io_i_slct];
endmodule

In my case, I want the different multiplexers to be explicitly unrolled with intermediate signals, with the output of each 2:1 multiplexer. For example here:

// Manually written
module MuxTree(
  input        clock,
               reset,
  input  [1:0] io_i_slct,
  input  [3:0] io_i_in_0,
               io_i_in_1,
               io_i_in_2,
               io_i_in_3,
  output [3:0] io_o_out
);

wire [3:0] _GEN0_0;
wire [3:0] _GEN0_1;
wire [3:0] _GEN1_0;

assign _GEN0_0 = (io_i_slct[0] == 1'b1) ? io_i_in_1 : io_i_in_0;
assign _GEN0_1 = (io_i_slct[0] == 1'b1) ? io_i_in_3 : io_i_in_2;
assign _GEN1_0 = (io_i_slct[1] == 1'b1) ? _GEN0_1 : _GEN0_0;

assign io_o_out = _GEN1_0;
endmodule

Is it possible to force this kind of behavior in Chisel ? Thanks for your support.