[BYOC] Running pass on functions for specific 3rd party codegen

lhutton1 · May 4, 2020, 1:12pm

I need to run a transform pass that is specific to my 3rd party codegen and the functions it receives. Currently I’m doing this in the codegen object that receives a module with multiple functions for a backend. I’ve done some digging and found this in transform.cc:

bool FunctionPassNode::SkipFunction(const Function& func) const {
  return (func->GetAttr<String>(attr::kCompiler).defined()) ||
  func->GetAttr<Integer>(attr::kSkipOptimization, 0) != 0;
}

This is preventing me from running a pass on these functions. Is there any reason to skip a function that has been assigned an external compiler?

I also noticed this PR was closed: https://github.com/apache/incubator-tvm/pull/5028 which I believe was a solution to the problem, but the problem still persists for me. Is there anyway around this?

comaniac · May 4, 2020, 4:42pm

The reason is that most Relay passes (e.g., op fusion, op simplifying) may change external function structure and result in unrecognized ops or patterns that a customized codegen cannot handle. If you have a customized Relay pass specific for your codegen, you can still let the pass visit external functions after the partition graph pass by checking kCompiler attribute.

Although the PR you pointed out has been closed, its functionality has been covered in another PR, so that you can find the kCompiler attribute in the master upstream:

github.com

apache/incubator-tvm/blob/master/src/relay/transforms/partition_graph.cc#L326


  global_region_func = Function(params, tuple, tuple->checked_type_, {}, DictAttrs());
}


std::string target = end_node->attrs.as<CompilerAttrs>()->compiler;
std::string name = target + "_" + std::to_string(region->GetID());


global_region_func =
    WithAttr(std::move(global_region_func), tvm::attr::kGlobalSymbol, runtime::String(name));
global_region_func = WithAttr(std::move(global_region_func), attr::kPrimitive, tvm::Integer(1));
global_region_func =
    WithAttr(std::move(global_region_func), attr::kCompiler, tvm::runtime::String(target));
global_region_func = WithAttr(std::move(global_region_func), attr::kInline, tvm::Integer(1));


// Constant propagation
if (!params_bind.empty()) {
  global_region_func = Downcast<Function>(relay::Bind(global_region_func, params_bind));
}


std::string fname = name;
CHECK(!module_->ContainGlobalVar(fname)) << "Global function " << fname << " already exists";
// Create a global function and add it to the IRModule for the region.

lhutton1 · May 4, 2020, 4:57pm

Thanks, I see what you’re saying about a custom pass. Is there anyway to do this for a generic pass though? For my use-case it’ll be ConvertLayout followed by FoldConstant.

The reason this worked in the PR above was because we reset kCompiler to default, with the current implementation my function remains annotated with Compiler="extern_compiler"

comaniac · May 4, 2020, 4:57pm

The current implementation lifts all external functions to global functions and inlines them back after most passes so that they won’t be touched. In this case we turn off all passes to external functions without adding the logic of checking kCompiler to each single pass. We do not have a mechanism to let users enable certain passes to external functions yet.

One suggested solution is running those passes manually before partitioning the graph. We had the same requirement for FoldConstant and this solution works fine with it.

lhutton1 · May 4, 2020, 5:30pm

Makes sense, although for my use-case this won’t be possible. I’ll be using convert layout to convert the kernel layout of convolutions within functions for my external codegen (HWIO -> OHWI). If I run this manually before the partitioning takes place it will also alter the kernel layout for convolutions offloaded to normal TVM (we can’t expect all convolutions are offloaded to our external codegen), which TVM won’t like.

Alternatively writing my own custom pass to do something similar to ConvertLayout only for my external codegen doesn’t really make sense.

Since the functions have been outlined and will no longer be touched by the passes, does it make sense to apply the PR above to enable this use-case for the time being? I agree that this will need to be revisited with a more robust mechanism.

comaniac · May 4, 2020, 5:44pm

Your use case makes sense to me. Here are my two cents:

I personally don’t really think it’s improper to have a custom pass doing something similar to an existing pass, because you may need more customizations for your target. However, it’s really case by case. If the current ConvertLayout can perfectly meet your requirement, then I agree that we should not make another one.
It’s still not necessary to apply the PR. As the following comment mentioned, we now treat the value of all unset attributes as default. You can achieve that by safely checking if the attribute is defined or not.

lhutton1 · May 4, 2020, 6:01pm

Thanks, apologies, for point number 2 I mean’t after the comments had been addressed.

mbaret · May 5, 2020, 1:10pm

I think the point here is that it’s removing the kCompiler attribute from ‘external IR modules’ - that is, modules which only contains functions designated for a particular compiler. Once we’ve created an external IR module, there’s no reason to keep the restriction that ‘normal’ passes like ConvertLayout can’t run on them so we should probably remove the kCompiler attribute from all the global functions.