[TIR] Handle DeclBuffer in CacheReadWrite schedule primitive

[Lunderberg]

Part of changes being split out from #14778 into independent portions. This commit allows TIR cache_read and cache_write schedule primitives to preserve DeclBuffer nodes.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[tqchen]

@Lunderberg there seems to be a regression introduced by this PR

[Lunderberg]

@tqchen Is there a specific test case for the regression? I pulled the latest main (31be726) and re-ran the tests introduced by this PR, and they are currently passing.

[tqchen]

see https://ci.tlcpack.ai/blue/organizations/jenkins/tvm-minimal-cross-isa/detail/nightly/140/pipeline/84 test_meta_schedule_schedule_rule_mlt_tc

[tqchen]

Not sure if it is a failed or flaky test, it is something that appears in this PR, and some followup commits. Perhaps some recent changes introduces some non-determinsm?

[Lunderberg]

Thank you, and I can reproduce that failure. It does seem to be a flaky test, with 28 failures out of 100 repetitions on main.

Though, on bisecting, it looks like the flaky failures were introduced in 02dc191 (#15081). The previous commit, 68ac909 (#15073), is able to pass the test 100/100 times, while #15081 shows about the same 30% failure rate.

[Lunderberg]

I have a quick bugfix PR at #15128, which should resolve the flaky test. I didn't tracked down why that specific overwrite results in failure to recognize a stage pipeline, but that does resolve the specific failure mode.