[TIR, Schedule] Add schedule primitive PadEinsum

[vinx13]

Co-authored-by: Bohan Hou 32121147+spectrometerHBH@users.noreply.github.com

This PR adds a schedule primitive PadEinsum. It is used for computation in Einsum pattern specifically, which cover most cases for tensorization. Different from general cases for padding in https://github.com/apache/tvm-rfcs/blob/main/rfcs/0077-layout-transform-padding.md, this primitive pads the output blocks and the input blocks at once, which eliminates the need to extra arithmetic analysis to provide the guarantee of program correctness.

cc @Hzfengsy @wrongtest-intellif @spectrometerHBH @Lunderberg

[wrongtest-intellif]

Compare to #12720 cc @Lunderberg
Could I understand that it equals with a bundle of operations in certain workload pattern? Like

for buffer in [A_shared, B_shared, C_shared]:
     s.transpose_layout(buffer, (127, 127) -> (128, 128), pad_value=0)
for block in [A, B, C_shared]:
     for axis in s.get_loops(block)
         s.fuse(*s.split(axis, [1, 128]))
s.annotate(C_shared, "en_some_predicate_versus_overcomputation_selection", 1)
        

[vinx13]

Yes. It pads the producers with init value (zero) and over-computes the reduction block

[Hzfengsy]

LGTM

[Lunderberg]

Overall, looks good, but just a few usability questions and potential improvements. I like seeing which assumptions are made here, which lead to a much simpler analysis than the more general case from the padding RFC.

I think the biggest question is the padding specified, and whether it can be specified as both a left/right padding, rather than only padding on the right.

[Lunderberg]

Nitpick: "it's" should be "its", without an apostrphe

[Lunderberg]

It looks like the padding can only be applied to the end of an axis/iterator, and cannot be applied to the beginning. Could we specify two arrays of padding, one for the lower end each block iter and one for the upper end?

[vinx13]

@Lunderberg The current assumption is to over compute the reduction block, and infer the padding of the producer. Since the padding is inferred from buffer access pattern, I think we can't specify the padding as tuple

[Lunderberg]

@vinx13 Thank you, and that makes sense. So, one of the simplifying assumptions that is all padding will only be on one side, and if the padding is allowed on both sides, that wouldn't just add a free parameter for the final output, but also for each producer.

[Lunderberg]

LGTM!

[junrushao]

@vinx13 let's fix the following warnings:

/root/Projects/tvm-dev/src/tir/schedule/primitive/pad_einsum.cc:231:8: warning: 'tvm::tir::PadEinsumRewriter::VisitStmt_' hides overloaded virtual function [-Woverloaded-virtual]
  Stmt VisitStmt_(const ForNode* op) final {
       ^
/root/Projects/tvm-dev/src/tir/schedule/primitive/.././transform.h:134:8: note: hidden overloaded virtual function 'tvm::tir::ReplaceBufferMutator::VisitStmt_' declared here: type mismatch at 1st parameter ('const tvm::tir::BufferStoreNode *' vs 'const tvm::tir::ForNode *')
  Stmt VisitStmt_(const BufferStoreNode* op) final;
       ^
/root/Projects/tvm-dev/src/tir/schedule/primitive/pad_einsum.cc:374:47: warning: lambda capture 'buffer_remap' is not used [-Wunused-lambda-capture]
  auto f_pad_buffer = [&padded_iter_extents, &buffer_remap](Buffer buffer,
                                           ~~~^~~~~~~~~~~~