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

Development
1

Subject: Request to Move LoopPartitionConfig to Public Header for Better C++ Integration

Dear TVM Community,

I’d like to propose a minor idea to the codebase by requesting that the LoopPartitionConfigNode and LoopPartitionConfig definitions be moved from src/tir/transforms/loop_partition.cc to a public header (loop_partition.h). This change would significantly improve the developer experience for those working with TVM’s transformation passes in C++ api.

Current Challenge

Currently, the LoopPartitionConfig class is defined exclusively within the loop_partition.cc implementation file. While this works for the pass itself, it creates significant friction for developers who need to configure this pass programmatically in their C++ applications.

I tried to config the PassContext as the TVM Python interface(I referred to :link:link)，but got the following error.:link:code

/// (partition_const_loop: true)
  Map<tvm::String, ObjectRef> inner_config;
  inner_config.Set("partition_const_loop", Bool(true));
  inner_config.Set("no_unroll_loop_with_extent_one", Bool(false));
  inner_config.Set("unroll_loop_with_partition_hint_no_interval", Bool(false));

  /// (tir.LoopPartition: inner_config)
  Map<String, ObjectRef> pass_config;
  pass_config.Set("tir.LoopPartition", inner_config);

  PassContext pass_ctx = PassContext::Current();
  pass_ctx->config = pass_config;
  {
    tvm::With<PassContext> scope(pass_ctx);
    PassContext current = PassContext::Current();
    auto pass = LoopPartition();
    /// @bug Aborted(core dumped)
    /// dumped at loop_partition.cc L801     auto cfg = ctx->GetConfig<LoopPartitionConfig>("tir.LoopPartition");
    
    auto partitionedMod = pass(mod);
  }

The program was dumped at tvm/src/tir/transforms/loop_partition.ccL801, for the reason that tvm::runtime::Map can’t be converted to LoopPartitionConfig.

Proposed Solution

By simply moving the existing class definitions to a public header loop_partition.h(I did it :link:link), we enable clean, type-safe configuration of PassContext.

#include <src/tir/transforms/loop_partition.h> // Proposed header

void configure_loop_partition_pass_context() {
 	auto configNode = make_object<LoopPartitionConfigNode>();
  	configNode->partition_const_loop = true;
    configNode->no_unroll_loop_with_extent_one = false;
    configNode->unroll_loop_with_partition_hint_no_interval = false;

    LoopPartitionConfig config(configNode);
    PassContext pass_ctx = PassContext::Create();
    Map<String, ObjectRef> pass_config;

    pass_config.Set("tir.LoopPartition", config);
    pass_ctx->config = pass_config;
    {
        tvm::With<PassContext> scope(pass_ctx);

        auto pass = LoopPartition();

        auto partitionedMod = pass(mod);
    }
}

By exposing the LoopPartitonConfigNode and LoopPartitonConfigNode in loop_partition.h, we can set the PassContext of Pass LoopPartition properly and concisely.

The passes including LoopPartition,UnrollLoop,InjectDoubleBuffer,RemoveNoOp,Simplify and so on, all need to retrieve their configurations from the PassContext. As a result, when developers wish to use these passes via the C++ API, they must instantiate the corresponding ConfigNode for each pass. Therefore, we need to place the declarations of each pass’s Config and ConfigNode in the .h file for developer access.

My expriment code

I have moved the defination of LoopPartitionNode and LoopPartition from loop_partition.cc into loop_partition.h in my :link:fork.

And my code of calling the pass in C++ api is here.:link:link

Looking forward to your discussion!

2

When possible, we recommend constructing pass context from the python side. This being said, we should be able to create a C++ API that mirrors some of the python API.

One of our principle is to reduce public surface area as a result the struct being in the cc file. But we should be able to create a Constructor of the pass context that mirrors python FFI call https://github.com/apache/tvm/blob/main/src/ir/transform.cc#L634 note the legalize here

so passing n dict should work

1 Like
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.