[RFC] DietCode: An Auto-Scheduler for Dynamic Tensor Programs

[ArmageddonKnight]

@comaniac @junrushao1994

[comaniac]

As one of the co-authors, it glad to see this RFC finally out :)
cc @masahi @ZihengJiang @Laurawly @Hzfengsy @MasterJH5574 @jinhongyii please help review and share your thoughts. Thanks.

[junrushao]

We are more than happy to collaborate on AutoTIR side to make good things happen :-)

[junrushao]

LGTM

[areusch]

hey guys, thanks for the RFC! would you guys be up for doing a short walkthrough of it at a Community Meeting?

[comaniac]

@ArmageddonKnight we might present/demo DietCode and introduce this RFC in a community meeting. Maybe in 5/25 or 6/1?

[ArmageddonKnight]

@areusch @comaniac Sure. I would be happy to. I will be traveling on 6/1 so 5/25 looks good.

[areusch]

@ArmageddonKnight @comaniac ok--let me do some agenda shuffling (we now have 3 topics for that meeting) and hopefully we can accommodate your travel schedule here.

[junrushao]

it’s definitely legit ask, but I would love to remind that as apache preferred approach, it would be better if technical discussion remains in an achievable way, so that the door is always open for all community participants to understand the technical discussion without having to appear at certain meeting at a certain time point (like early in the morning 😂) :-)

Perhaps it would be helpful if we take meeting notes and publicize it somewhere to help people who cannot join the meeting?

[areusch]

@junrushao1994 yep, we will post meeting notes here per the TVM Community Meeting guidelines.

[areusch]

@ArmageddonKnight just confirming, we have you on the agenda to discuss tomorrow

[ArmageddonKnight]

@areusch Thank you for the reminder. I have put that on my calendar.

[ArmageddonKnight]

@areusch BTW, may I ask whether there is anything I should prepare for the meeting? Alternatively, any particular thing that the community might be interested in? How much time am I given for the presentation?

[areusch]

adding a couple of questions here based on my understanding from the community meeting this morning

[areusch]

We discussed this at the TVM Community Meeting this morning. Here are notes (thanks to mbs-octoml for taking these):

• Most discussion was around whether the approach could help the related
  problem of long training time and a large number of distinct kernels
  for fully statically-shaped models. Eg resnet50 with its dozens of
  static shape instances for conv2d.

  It is useful to consider kernel compilation in two steps:

  • Start with the N instances of the same operator, with different
    shapes.
  • After tuning those independently, may end up with M << N schedules.
  • However, each schedule is instantiated for the original static shape
    so that no dimension variables are left behind at run time,
    thus we are back at N implementations.

  There's a significant perf impact to having only M kernel implementation
  which retain their parameterization on dimension variables. However, TVM
  could support not doing that inlining today.

  DietCode could be used to reduce M. The learned decision tree would dispatch
  from a fully shape polymorphic implementation to the specific schedule
  implementation. But at this stage the focus is on dynamic shape workloads.

• No changes to Relay are required. Though Relay's dynamic shape functions
  support data-dependent output shapes, currently DietCode does not search
  over output shape and assumes only input shapes need to be inspected. But
  sounds like this could be addressed in future work.

[ArmageddonKnight]

@areusch @cgerum Thank you for your feedback. I have addressed your feedback. Please review it again.

[areusch]

thanks for addressing my comments @ArmageddonKnight ! @cgerum could you take another look here? I think otherwise we should be good to merge

[cgerum]

@areusch It's good to merge from my side

[comaniac]

Thanks all for your valuable comments and insightful discussions. This RFC is now merged.
@ArmageddonKnight please open a tracking issue in the main TVM repo to start tracking the progress. Thanks.

[renfeier]

@ArmageddonKnight
hello, i have a question looking forward to you reply.
Is it workable to use DietCode + vm compiler + vm runtime to implement graphlevel dynamic shape inference ?
do you have any suggestions?

[ArmageddonKnight]

It is workable. The VM can gather the workloads and pack them as a single dynamic-shape workload and forward it to DietCode. DietCode does the auto-scheduling and send the generated code back to the VM.

[renfeier]

It is workable. The VM can gather the workloads and pack them as a single dynamic-shape workload and forward it to DietCode. DietCode does the auto-scheduling and send the generated code back to the VM.
Thank you very much for your prompt reply. :-D, b（￣▽￣）d