[Relay] Add new IR pass CombineParallelDense

[soiferj]

This IR pass is similar to CombineParallelConv2D, but with the Dense operator. The pass takes multiple parallel dense operators and combines them into one batch_matmul operator. Element-wise operations after dense are also stacked and fused.

The code is refactored so CombineParallelConv2D and CombineParallelDense can share code, and make it easier to write a new "combine" pass later.

This is in opt level 4. In another PR, we can add tuning to decide whether or not to apply this pass.

This will help models like BERT, which have parallel Dense + BiasAdd at the start of each layer.

Discussed in this RFC: https://discuss.tvm.ai/t/discussion-new-ir-pass-proposal-combineparalleldense/3813

@vinx13 @icemelon9 would you be able to take a look?

[jroesch]

@vinx13 I will add my review can you as well?

[jroesch]

@MarisaKirisame can you add review as well

[jroesch]

Overall LGTM, thanks for you contribution! My comments are about documentation mostly.

We are making a push this release cycle to improve the docs. I am trying to encourage everyone to incrementally grow them with each PR to passes, etc.

I will take another pass after other reviews come in, and then merge.

[MarisaKirisame]

As most operators take in batched input, it is not hard to imagine different version of CombineXXX... Furthermore, if we do dynamic batching ala tensorflow fold (which I think we will do), we need a CombineXXX for every XXX operator. So, is it possible to merge the two combine into a generic pass?

[soiferj]

@MarisaKirisame that's a really good idea. The only tricky part is that IsOpSupported and AreOpsCompatible may need to change for each op. For example, with Dense, we need to make sure that the units attribute is empty.

How about this: we create a new class, CombineParallelOpBatch, that has a default implementation of IsOpSupported which returns true, and a default implementation of AreOpsCompatible which returns true if the shapes and dtypes of both ops are the same. From there, we can have CombineParallelDense inherit from CombineParallelOpBatch and override IsOpSupported and AreOpsCompatible. The constructor can take the name of the op and the min number of branches to combine. This way, we can easily add new IR passes by adding a line like CombineParallelOpBatch('add', 3).

I'm not sure it's possible to merge CombineParallelConv2D and CombineParallelDense because CombineParallelConv2D concatenates the inputs rather than stacking. The implementations of most methods would have to change.

What do you think? @jroesch any thoughts?

[MarisaKirisame]

@soiferj sounds good to me. maybe the default should be False, because op need manual batching to be supported (for example, changing dense to BatchMatMul). For the easy case, we can override DeclareUnaryOp and DeclareBinaryOp, which will do the job

[soiferj]

@MarisaKirisame, sounds good. For now, I'll add support for this functionality. I think we can decide how we actually want to enable it for all ops in another PR. Is that alright?

[MarisaKirisame]

@soiferj exactly.

[soiferj]

@MarisaKirisame do you think that the new base class, CombineParallelOpBatch, should handle inputs with any number of dimensions and convert them to a batched version of ndim + 1 (for example, (2, 3, 4) -> (1, 2, 3, 4)), or should we just handle inputs with 2 dimensions and convert them to 3 dimensions (for example, (3, 4) -> (1, 3, 4))?

The reason why I ask is that broadcasting logic for element-wise ops that follow the main op can get tricky if we support an arbitrary number of dimensions.

[MarisaKirisame]

@soiferj both are completely ok with me.
the main scenario I am afraid of, is you add CombineDense, someone add CombineRelu next week, another one add CombineSoftmax, etc etc etc... And after one year we see 10 seprate pass that could be unifiable. If we start with a better design, we save ourself a big refactor down the road.
It dont need to support much (more dimension, broadcasting, etc), the infrastructure itself working and allowing extension is the most important part.

[jroesch]

@soiferj sounds good to me in general, I am always in favor of slowly layering changes, etc but setting yourself up for future success.

I think Marisa's point is that we should be able to hook just a few methods and leave the analysis/traversal code generic between the passes. For example we could also overload how they are combined with another hook, i.e concat vs. stack.

If I understand the goal correctly is we have some operators like, op1(args, attrs) and op1(args, attrs) we should only need to add 3 methods, one to check if op1 is supported, should be an easy one liner, that the args and attrs match, and how to combine the values, i.e. call stack on a sequence of arguments.

[soiferj]

Just pushed a new set of changes that introduces CombineParallelOpBatch. There is also a little bit of special logic for combining bias_add, as it needs to be broadcasted. For now, I'm pretty happy with the code design. I think we can look at moving the stack and concatenate logic into the base class if we see a need to in the future. Let me know what you think!

[soiferj]

Also, I have a design question: after we support CombineParallelXXXX for many ops, is there a purpose to fusing elementwise operators in these classes? If we can really combine most parallel ops into batch ops, the regular fusion pass should do this for us, right? This will clean up the code a lot.

[vinx13]

@soiferj Likely some fused operations are not inlined. Can you check what sub-function after fusion raised this error during Relay compilation?

[soiferj]

The error says it's coming from split_dev_host_funcs in codegen/build_module.cc. The line is

CHECK(ir::VerifyMemory(x, target->device_type)) << "Direct host side access to device memory is detected in " << x->func_name() << ". Did you forget to bind?";

[vinx13]

We need to trace the Relay part of these functions (and see how their TOPI schedules are called)
You can use debug logging here to get more info
https://github.com/dmlc/tvm/blob/master/python/tvm/relay/backend/_backend.py#L51-L53

[soiferj]

Okay, I have the output. What exactly should I be looking for?

[vinx13]

@soiferj After fusion the Relay program is split into several parts like

%x = fn(...) {
  dense
  relu
}

Each part is passed to tvm.build_module independently. So the thing is to find exactly which part is broken

[soiferj]

@vinx13, nevermind, I made a mistake on my testing branch. I created a new branch, merged my changes with master, and can run the model now. Sorry about that.

[soiferj]

The CI failure doesn't look related to me...it is occurring in test_forward_conv_batch_norm. The output from the test actually looks correct too. Is this a known issue? Edit: issue has been resolved.

[soiferj]

Thanks everyone for your help and feedback :) This PR has been approved for the past few days - would someone be able to take a look at merging when they get a chance?

[icemelon]

@jroesch Could you take one more look at this PR?

[soiferj]

@jroesch would you mind taking a look when you have a chance?

[jroesch]

Sorry for dropping the ball on this one, been busy with life and other work things :) will merge tonight.

[jroesch]

I think everything looks good to me, thanks for contribution, especially all the extra work generalizing the interface