[BYOC][runtime] JSON runtime for BYOC

Another alternate might be having a compilation flag to indicate whether to put this hash map into runtime

Would Array of NDArray be sufficient?

Yeah, I think I didn’t make it very clear. The problem was because we may have multiple subgraphs, each of them may have “var_name: NDarray” pairs. I was trying to just have one ModuleInitWrapper to take charge of the initialization of engines for all subgraphs so that users don’t need to orchestrate from the Python side. This means we need to have a field like Map<String, Map<String, NDArray>> metadata_, to save the mapping of each subgraph symbol to the var_name: ndarray pairs. This makes passing by array a bit more complicated.

Another alternative approach could be having multiple ModuleInitWrapper. Each of them only takes care of one engine. Then we would only have Map<var_name, ndarray>, and then we could pass this by two arrays Array<String> variables_ and Array<NDArray> metadata_. This would need users to be aware of that there would be multiple modules and they need to initiate multiple ModuleInitWrapper for different engines. This is okay for CSourceModule stuff because we would usually only have one such a module. For execution engines like TRT, there would be multiple of it depending on the number of subgraphs.

Maybe I can give a try to the alternative approach. @tqchen @junrushao how do you think about the effort from frontend?

Indeed it is a tradeoff, and in this case there can certainly be multiple choices. The key problem we want to answer is how do we want to expose pass “symbol” of meta data to each engine. This is an important topi as it can affect the serialization convention of our future packages, let us think a bit more about it.

To faciliate the discussion, let us think the following scenario

fn mymodule(%y) {
   let subfunc0 = (%x: Tensor[(10, 10)]) {
      %1 = meta[Constant][0] 
      %2 = meta[Constant][1] 
      (%x + %1) * %2 
  } 
   let subfunc1 = (%x: Tensor[(10, 10)]) {
      %2 = meta[Constant][0] 
      %x * %2 
  } 
   let subfunc2 = (%x: Tensor[(10, 10)]) {
      %3 = meta[Constant][1] 
      %x / %3 
  } 
  %1 = %subfunc0(%y)
  %2 = %subfunc1(%1)
  %3 = %subfunc2(%2)
  %3
}

The resulting module can be seen as follows:

metadatamodule (stores constant 0, 1, 2 and need to pass them to be module)
    - sourcemodule0(contains subfunc0, subfunc1)
    - sourcemodule1(contains subfunc2)

Our design choices include different ways to pass in these constants(metadata). The choices include:

  • C0: Uniquely name the “symbol” of each constant globally in the module
    • Pass by calling set_module_param(String name, NDArray)
  • C1: Uniquely name the “function” of each module, and pass initializers for each function.
    • init_mod0_subfunc0({const0, const1}), init_mod1_subfunc1({const0})
  • C2: Uniquely name each module(module0, module1)
    • call init_module0(Map<String, Object/NDArray>)

Each of the choices have certain pros and cons in terms of the following aspect:

  • API simplicity(Map does comes with an overhead).
  • The ability to see all required params in an engine(before calling setup functions),for example in subfunc0 may need to see const0, and const1 in order to setup its engine.
  • Consistency with the function naming convention: due to the restriction of DSO engine, we simply use one-level string to map from name to functions.

Finally, it is important to consider the following aspect of module combination. In the final linking stage, sourcemodule0 and sourcemodule1 goes into the same DSO library. This means we cannot use the same initialize function for them, and need to have a unique name for their initialization function.

cc @FrozenGene as it is also related to module exportation format

Yeah, I would prefer C1 or C2. C2 was pretty much what I was doing.

Do you think if it is possible to do C1? since it reduces the requirement for passing Map

Yeah, let me give it a try.

BTW, we will need to have the variables as well, i.e. %x1, %x2, %x3, something as I mentioned above. This is because we need to know which variable a ndarray should be assigned to.

Yah, I think it is fair to pass in names of each variable. Another way is to rely on the positional ordering of the constant themselves(so the variable name is implicit).

yeah, I thought about positional ordering as well. But it looks pass variables might be safer. For a CSourceModule external codegen we generate a wrapper like float* a = const_0; const_0 would need to be produced by the initializer later. So we would anyway need a name for it.

1 Like