Proposal: Expose LoopPartitionConfig and other Pass Config Class in Public Header

Development
3

Thanks a lot!

I got it. We can implement the legalize function with the C++ api as follow:

  Map<String, ObjectRef> pass_config_str;
  pass_config_str.Set("partition_const_loop", tvm::Bool(true));

  auto *reflection = ReflectionVTable::Global();

  ///@brief Legalize the pass config, which converts the `Map<String,Bool>` to
  ///`LoopPartitionConfig`
  ///@param obj The pass config, which is a `Map<String,Bool>`
  ///@param type_key struct `LoopPartitionConfigNode`'s `_type_key`
  ///@return The legalized pass config, which is a `LoopPartitionConfig`
  ///@sa tvm/src/tir/transforms/loop_partition.cc
  auto legalization = [=](const ObjectRef &obj) -> ObjectRef {
    return reflection->CreateObject("tir.transform.LoopPartitionConfig",
                                    Downcast<Map<String, ObjectRef>>(obj));
  };

  auto pass_config = legalization(pass_config_str);

  Map<String, ObjectRef> config;
  config.Set("tir.LoopPartition", pass_config);

  PassContext pass_ctx = PassContext::Create();
  pass_ctx->config = config;

  With<PassContext> scope(pass_ctx);
  IRModule partitionedMod = LoopPartition()(mod);

Above approach enables calling LoopPartition via the C++ API without altering the TVM source code. Once again, thank you for your response.

1 Like