[mlir] Add hoisting of transfer ops in affine loops

[jsetoain]

Summary

The only way to do this with the current hoisting strategy is by
lowering Affine to Scf first, but that prevents further passes on
Affine.

[jsetoain]

Make utility function names more specific

[jsetoain]

Fix typo

[jsetoain]

clang-format

[bondhugula]

This looks great to me - thanks! Mostly minor comments.

mlir/include/mlir/Dialect/Affine/Utils.h

336

Nit: doesn't
Similarly, apostrophes missing below.

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

521–522

dyn_cast here to avoid an extra check below.

528

dyn_cast here again.

534

You should actually check earlier and advance the walk if it's something other than affine.for or scf.for at line 428; either that or assert here with an extra else or use cast above for AffineForOp in the conditional.

mlir/lib/Dialect/SCF/Utils/Utils.cpp

41–44

This should've been in the mlir::scf namespace. You can update it in this PR itself if you wish.

151

SCFForLoop -> SCFFor or SCFLoop for conciseness.

[jsetoain]

Address comments

[jsetoain]

Thanks for the review!

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

521–522

I haven't found a way to dyn_cast the result of the dyn_cast. I'm not very familiar with llvm's static polymorphism constructs, so I'm not sure what the problem is, but I get errors like:

error: invalid cast of an rvalue expression of type ‘std::nullptr_t’ to type ‘llvm::CastInfo<mlir::scf::ForOp, mlir::LoopLikeOpInterface, void>::CastReturnType’ {aka ‘mlir::scf::ForOp&’}

Alternatively, I can write: dyn_cast<scf::ForOp>(transferRead->getParentOp()) but I'm not sure it's worth the loss of clarity.

mlir/lib/Dialect/SCF/Utils/Utils.cpp

41–44

Indeed, this has been bugging me since I submitted the patch. I agree, it would be better if instead of using namespaces in the function names, we placed these utility functions in their respective namespaces. It's a straightforward NFC, but it would splash everywhere; I will push a separated patch as soon as I'm done with this review.

[bondhugula]

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

521–522

Please see https://llvm.org/docs/ProgrammersManual.html#the-isa-cast-and-dyn-cast-templates.

Use dyn_cast instead of isa + cast, i.e.,

if (auto scfFor = dyn_cast<scf::ForOp>(loop)) {
  ...

mlir/lib/Dialect/SCF/Utils/Utils.cpp

148

innermost

[bondhugula]

mlir/lib/Dialect/Affine/Utils/Utils.cpp

1877–1941

Can we templatize this function and make it work for both AffineForOp and scf::ForOp? You'll need YieldOpTy (can be scf::YieldOp and AffineYieldOp) and ForOpTy (can be scf::ForOp or AffineForOp). Other than the op creation calls, isn't everything else the same?

[jsetoain]

Try dyn_cast

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

521–522

I believe I didn't explain myself very well there. I understand how LLVM's RTTI is supposed to work, I don't understand why it's not working. I'm getting that error in my previous comment, but I don't understand enough about the implementation to figure out why.

I'm going to push the change anyway. Best case scenario, it's something wrong with my compiler. Worst case scenario, it will leave a nice log of the error I'm seeing.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

521–522

I am missing overall context, but you cant dyn_cast result of the dyn_cast....

IIUC what is being suggested here is

if (auto scfForOp = dyn_cast<scf::ForOp>(loop)) {
  ....
} else if (auto affineForOp = dyn_cast<AffineForOp>(loop)) {

}

or better yet, use TypeSwitch

return TypeSwitch<Operation *, WalkResult>(loop)
  .Case<scf::ForOp>([](scf::ForOp scfForOp) {...})
  .Case<AffineForOp>([](AffineForOp affineForOp) {...})
  .Default([](Operation *) { return WalkResult::interrupt(); }

(and move the loop.erase into the if in first case or callback for the Case in the second example).

[jsetoain]

Address comments (1/n)

[jsetoain]

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

521–522

That's what I suspected, and the problem is that loop is the result of a dyn_cast, so it can't be dyn_cast-ed again. Thankfully, TypeSwitch does the trick.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

521–522

dyn_cast, etc. works only on Operation *. When you define TypeSwitch<Operation *, ....>, the LoopLikeInterface op gets implicitly cast to Operation *, and then TypeSwitch works the same way as dyn_cast. So with the if version you would need if (auto scfForOp = dyn_cast<scf::ForOp>(loop.getOperation()))

[mravishankar]

I dont have further comments on this patch. This only landed on my radar cause of this statement `I'm going to push the change anyway. Best case scenario, it's something wrong with my compiler. Worst case scenario, it will leave a nice log of the error I'm seeing.`
I might have mis-judged the sentiment, but pushes that could cause failures would cause pain downstream (and would probably get reverted anyway). Apologies if I misread this.

[jsetoain]

Templatize function

[jsetoain]

In D137600#3948191, @mravishankar wrote:

I dont have further comments on this patch. This only landed on my radar cause of this statement `I'm going to push the change anyway. Best case scenario, it's something wrong with my compiler. Worst case scenario, it will leave a nice log of the error I'm seeing.`
I might have mis-judged the sentiment, but pushes that could cause failures would cause pain downstream (and would probably get reverted anyway). Apologies if I misread this.

My apologies. English is not my native language, I didn't mean it like that, but I see where the confusion may have come from. What I was trying to say is "I will update this patch in revision, let the automatic checks fail where they fail for me locally, and then I can point to those logs to explain my problem". You can rest assure, I never had the intention to land this patch until all issues are clarified and everyone is happy with my code, and I would never land a patch that doesn't pass all tests in my machine :-) I appreciate the clarification on dyn_cast issue, this was a new one for me, thanks!

mlir/lib/Dialect/Affine/Utils/Utils.cpp

1877–1941

That and the access to the YieldOp operands (which I find strange). I considered the template early on, but was not convinced with the result. I'm uploading my templated version, let me know what you think. My issues are:

1. It only really works for scf::ForOp and AffineForOp
2. I'm getting away with not making SCF dependent on Affine only because that's the only alternative to scf::ForOp. It's creating a non-explicit dependency that I don't like.
3. The whole thing should probably live somewhere else, but at the same time, this is code depends on a lot of for loop-related functions and methods that don't work for all possible loops.
4. It only "coincidentally" works for both Affine and SCF fors because they have methods with the same identifiers. It feels like abusing templates to avoid copying a couple dozen lines of code.

That's why I originally scrapped it and decided in favour of total isolation. Whether this is preferable or not I will leave it up to you, you have the experience to judge :-)

[jsetoain]

Remove unnecessary empty line

[bondhugula]

LGTM - thanks.

[bondhugula]

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

54

This much C++ code in the header file isn't okay. I understand this was added here because it's a template function. Instead, you can just define it out of line and explicitly instantiate for both scf::ForOp and AffineForOp in the C++ file (two extra lines right below the definition).

[jsetoain]

Replace the template with an existing function

[jsetoain]

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

54

I don't think that would have worked. Since I needed to instantiate the template in two different C++ files, the template had to be visible from both. In any case, I found a function that almost does what I need, and I can get by using that one plus some additional code right after that, which I believe is preferable.

[bondhugula]

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

54

It doesn't matter how many different places you need to use the instantiation. What I meant was to add an explicit instantiation in Utils.cpp. Anything linking with the library would have the necessary things to link to. I don't see why it wouldn't work. For eg. see SideEffectInterfaces.cpp:

template bool mlir::hasEffect<MemoryEffects::Allocate>(Operation *, Value);
template bool mlir::hasEffect<MemoryEffects::Free>(Operation *, Value);
template bool mlir::hasEffect<MemoryEffects::Read>(Operation *, Value);
template bool mlir::hasEffect<MemoryEffects::Write>(Operation *, Value);

[jsetoain]

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

54

Ah! I completely misunderstood what you were asking for, my apologies. Nevertheless, am I right to assume you're okay with the alternative? If so, I will leave this here for a couple more days and land it. Thanks for all your help (and patience)!

[hctim]

It looks like the update to Dialect/Linalg/hoisting.mlir breaks ASan, MSan and UBSan: https://lab.llvm.org/buildbot/#/builders/5/builds/29773

You can find instructions on reproducing the sanitizer bots here: https://github.com/google/sanitizers/wiki/SanitizerBotReproduceBuild. The buildbot_fast is the one is probably easiest to use and reflects all the breakages.

******************** TEST 'MLIR :: Dialect/Linalg/hoisting.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/mlir-opt /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/test/Dialect/Linalg/hoisting.mlir -test-linalg-hoisting=test-hoist-redundant-transfers -allow-unregistered-dialect -split-input-file | /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/test/Dialect/Linalg/hoisting.mlir
--
Exit Code: 1
Command Output (stderr):
--
=================================================================
==455613==ERROR: AddressSanitizer: stack-use-after-scope on address 0x7f43c2dbea50 at pc 0x5607ff79ca82 bp 0x7ffd712d54a0 sp 0x7ffd712d5498
READ of size 8 at 0x7f43c2dbea50 thread T0
    #0 0x5607ff79ca81 in mlir::ValueRange::dereference_iterator(llvm::PointerUnion<mlir::Value const*, mlir::OpOperand*, mlir::detail::OpResultImpl*> const&, long) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/OperationSupport.cpp:613:12
    #1 0x5607feb11ea4 in operator* /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:1281:14
    #2 0x5607feb11ea4 in deref<0UL, 1UL> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:768:23
    #3 0x5607feb11ea4 in operator* /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:793:12
    #4 0x5607feb11ea4 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:543:26
    #5 0x5607feb11ea4 in Case<mlir::AffineForOp, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:534:30)> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/TypeSwitch.h:122:16
    #6 0x5607feb11ea4 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:534:12
    #7 0x5607feb11ea4 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Visitors.h:230:14
    #8 0x5607feb11ea4 in mlir::WalkResult llvm::function_ref<mlir::WalkResult (mlir::Operation*)>::callback_fn<std::__1::enable_if<!llvm::is_one_of<mlir::vector::TransferReadOp, mlir::Operation*, mlir::Region*, mlir::Block*>::value && std::is_same<mlir::WalkResult, mlir::WalkResult>::value, mlir::WalkResult>::type mlir::detail::walk<(mlir::WalkOrder)1, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2, mlir::vector::TransferReadOp, mlir::WalkResult>(mlir::Operation*, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2&&)::'lambda'(mlir::Operation*)>(long, mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
    #9 0x5607ff7e8a04 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #10 0x5607ff7e8a04 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #11 0x5607ff7e8a04 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #12 0x5607ff7e8a04 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #13 0x5607feb060c4 in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::vector::TransferReadOp, mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Visitors.h:233:10
    #14 0x5607feb060c4 in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Operation.h:579:12
    #15 0x5607feb060c4 in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/OpDefinition.h:147:19
    #16 0x5607feb060c4 in mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:10
    #17 0x5607feb050f7 in (anonymous namespace)::TestLinalgHoisting::runOnOperation() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/test/lib/Dialect/Linalg/TestLinalgHoisting.cpp:48:5
    #18 0x5607ff3b525a in mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:465:11
    #19 0x5607ff3b6164 in mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:529:16
    #20 0x5607ff3b9655 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:749:36
    #21 0x5607ff3b9655 in failableParallelForEach<std::__1::__wrap_iter<OpPMInfo *>, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:739:20) &> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Threading.h:46:18
    #22 0x5607ff3b9655 in failableParallelForEach<std::__1::vector<OpPMInfo, std::__1::allocator<OpPMInfo> > &, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:739:20) &> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Threading.h:92:10
    #23 0x5607ff3b9655 in mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:759:14
    #24 0x5607ff3b52e3 in runOnOperation /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:650:5
    #25 0x5607ff3b52e3 in mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:463:14
    #26 0x5607ff3bcfb4 in runPipeline /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:529:16
    #27 0x5607ff3bcfb4 in runPasses /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:834:10
    #28 0x5607ff3bcfb4 in mlir::PassManager::run(mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:814:60
    #29 0x5607ff3abb7a in performActions(llvm::raw_ostream&, bool, bool, llvm::SourceMgr&, mlir::MLIRContext*, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:91:17
    #30 0x5607ff3ab260 in processBuffer /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:139:12
    #31 0x5607ff3ab260 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:181:12
    #32 0x5607ff3ab260 in mlir::LogicalResult llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, bool, bool, bool, bool, bool, bool, bool)::$_0>(long, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
    #33 0x5607ff5341a5 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
    #34 0x5607ff5341a5 in mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool)::$_0::operator()(llvm::StringRef) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:88:16
    #35 0x5607ff532f46 in interleave<const llvm::StringRef *, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:81:23), (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2085:43), void> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2067:5
    #36 0x5607ff532f46 in interleave<llvm::SmallVector<llvm::StringRef, 8U>, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:81:23), llvm::raw_ostream, llvm::StringRef> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2085:3
    #37 0x5607ff532f46 in mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:91:3
    #38 0x5607ff3a51f5 in mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, bool, bool, bool, bool, bool, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:186:10
    #39 0x5607ff3a65ef in MlirOptMain /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:209:10
    #40 0x5607ff3a65ef in mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:306:14
    #41 0x5607fac120ff in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/tools/mlir-opt/mlir-opt.cpp:244:7
    #42 0x7f43c481dd8f  (/lib/x86_64-linux-gnu/libc.so.6+0x29d8f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #43 0x7f43c481de3f in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x29e3f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #44 0x5607fab5b5e4 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/mlir-opt+0x60ce5e4)
Address 0x7f43c2dbea50 is located in stack of thread T0 at offset 592 in frame
    #0 0x5607feb0fc9f in mlir::WalkResult llvm::function_ref<mlir::WalkResult (mlir::Operation*)>::callback_fn<std::__1::enable_if<!llvm::is_one_of<mlir::vector::TransferReadOp, mlir::Operation*, mlir::Region*, mlir::Block*>::value && std::is_same<mlir::WalkResult, mlir::WalkResult>::value, mlir::WalkResult>::type mlir::detail::walk<(mlir::WalkOrder)1, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2, mlir::vector::TransferReadOp, mlir::WalkResult>(mlir::Operation*, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2&&)::'lambda'(mlir::Operation*)>(long, mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:44
  This frame has 31 object(s):
    [32, 56) 'ref.tmp.i.i'
    [96, 104) 'newForOp.i.i24.i'
    [128, 136) 'ref.tmp.i.i25.i'
    [160, 216) 'ref.tmp14.i.i.i'
    [256, 258) 'ref.tmp3.i.i.i.i.i.i'
    [272, 280) 'newForOp.i.i.i'
    [304, 312) 'ref.tmp.i.i16.i'
    [336, 352) 'agg.tmp7.i.i.i'
    [368, 432) 'ref.tmp8.i.i.i'
    [464, 472) 'ref.tmp10.i.i.i'
    [496, 512) 'agg.tmp15.i.i.i'
    [528, 536) 'ref.tmp16.i.i.i'
    [560, 576) 'ref.tmp35.i.i.i'
    [592, 600) 'ref.tmp37.i.i.i' <== Memory access at offset 592 is inside this variable
    [624, 648) '__begin3.i.i.i'
    [688, 704) 'it.i.i.i'
    [720, 728) 'ref.tmp53.i.i.i'
    [752, 808) 'ref.tmp56.i.i.i'
    [848, 872) 'ref.tmp.i517.i.i'
    [912, 936) 'ref.tmp.i444.i.i'
    [976, 1000) 'ref.tmp.i365.i.i'
    [1040, 1048) 'ref.tmp.i.i.i.i'
    [1072, 1088) 'ref.tmp.i.i.i'
    [1104, 1112) 'transferRead.i.i'
    [1136, 1152) 'loop.i.i'
    [1168, 1224) 'forwardSlice.i.i'
    [1264, 1272) 'transferWrite.i.i'
    [1296, 1304) 'candidateWrite.i.i'
    [1328, 1360) 'dom.i.i'
    [1392, 1424) 'b.i.i'
    [1456, 1504) 'yieldFn.i.i'
HINT: this may be a false positive if your program uses some custom stack unwind mechanism, swapcontext or vfork
      (longjmp and C++ exceptions *are* supported)
SUMMARY: AddressSanitizer: stack-use-after-scope /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/OperationSupport.cpp:613:12 in mlir::ValueRange::dereference_iterator(llvm::PointerUnion<mlir::Value const*, mlir::OpOperand*, mlir::detail::OpResultImpl*> const&, long)
Shadow bytes around the buggy address:
  0x0fe8f85afcf0: f5 f5 f5 f5 f5 f5 f5 f5 f5 f5 f5 f5 f5 f5 f5 f5
  0x0fe8f85afd00: f1 f1 f1 f1 f8 f8 f8 f2 f2 f2 f2 f2 f8 f2 f2 f2
  0x0fe8f85afd10: f8 f2 f2 f2 f8 f8 f8 f8 f8 f8 f8 f2 f2 f2 f2 f2
  0x0fe8f85afd20: f8 f2 00 f2 f2 f2 f8 f2 f2 f2 00 00 f2 f2 f8 f8
  0x0fe8f85afd30: f8 f8 f8 f8 f8 f8 f2 f2 f2 f2 f8 f2 f2 f2 00 00
=>0x0fe8f85afd40: f2 f2 f8 f2 f2 f2 f8 f8 f2 f2[f8]f2 f2 f2 00 00
  0x0fe8f85afd50: 00 f2 f2 f2 f2 f2 00 00 f2 f2 f8 f2 f2 f2 f8 f8
  0x0fe8f85afd60: f8 f8 f8 f8 f8 f2 f2 f2 f2 f2 f8 f8 f8 f2 f2 f2
  0x0fe8f85afd70: f2 f2 f8 f8 f8 f2 f2 f2 f2 f2 f8 f8 f8 f2 f2 f2
  0x0fe8f85afd80: f2 f2 f8 f2 f2 f2 f8 f8 f2 f2 00 f2 f2 f2 00 00
  0x0fe8f85afd90: f2 f2 00 00 00 00 00 00 00 f2 f2 f2 f2 f2 00 f2
Shadow byte legend (one shadow byte represents 8 application bytes):
  Addressable:           00
  Partially addressable: 01 02 03 04 05 06 07 
  Heap left redzone:       fa
  Freed heap region:       fd
  Stack left redzone:      f1
  Stack mid redzone:       f2
  Stack right redzone:     f3
  Stack after return:      f5
  Stack use after scope:   f8
  Global redzone:          f9
  Global init order:       f6
  Poisoned by user:        f7
  Container overflow:      fc
  Array cookie:            ac
  Intra object redzone:    bb
  ASan internal:           fe
  Left alloca redzone:     ca
  Right alloca redzone:    cb
==455613==ABORTING

******************** TEST 'MLIR :: Dialect/Linalg/hoisting.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/mlir-opt /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/test/Dialect/Linalg/hoisting.mlir -test-linalg-hoisting=test-hoist-redundant-transfers -allow-unregistered-dialect -split-input-file | /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/test/Dialect/Linalg/hoisting.mlir
--
Exit Code: 1
Command Output (stderr):
--
/b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Operation.h:165:52: runtime error: member call on misaligned address 0x000000000051 for type 'mlir::Block', which requires 8 byte alignment
0x000000000051: note: pointer points here
<memory cannot be printed>
    #0 0x5583a6b717ba in getParentOp /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Operation.h:165:52
    #1 0x5583a6b717ba in mlir::Operation::isProperAncestor(mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Operation.cpp:177:26
    #2 0x5583a6b99283 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
    #3 0x5583a6b99283 in mlir::Value::replaceUsesWithIf(mlir::Value, llvm::function_ref<bool (mlir::OpOperand&)>) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Value.cpp:85:9
    #4 0x5583a6718589 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:545:31
    #5 0x5583a6718589 in llvm::TypeSwitch<mlir::Operation*, mlir::WalkResult>& llvm::TypeSwitch<mlir::Operation*, mlir::WalkResult>::Case<mlir::AffineForOp, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2::operator()(mlir::vector::TransferReadOp) const::'lambda'(mlir::AffineForOp)>(mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2::operator()(mlir::vector::TransferReadOp) const::'lambda'(mlir::AffineForOp)&&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/TypeSwitch.h:122:16
    #6 0x5583a6717ddf in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:534:12
    #7 0x5583a6717ddf in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Visitors.h:230:14
    #8 0x5583a6717ddf in mlir::WalkResult llvm::function_ref<mlir::WalkResult (mlir::Operation*)>::callback_fn<std::__1::enable_if<!llvm::is_one_of<mlir::vector::TransferReadOp, mlir::Operation*, mlir::Region*, mlir::Block*>::value && std::is_same<mlir::WalkResult, mlir::WalkResult>::value, mlir::WalkResult>::type mlir::detail::walk<(mlir::WalkOrder)1, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2, mlir::vector::TransferReadOp, mlir::WalkResult>(mlir::Operation*, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2&&)::'lambda'(mlir::Operation*)>(long, mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
    #9 0x5583a6b9e8d6 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
    #10 0x5583a6b9e8d6 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:181:12
    #11 0x5583a6b9e852 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #12 0x5583a6b9e852 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #13 0x5583a6b9e852 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #14 0x5583a6b9e852 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #15 0x5583a67143be in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::vector::TransferReadOp, mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Visitors.h:233:10
    #16 0x5583a67143be in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Operation.h:579:12
    #17 0x5583a67143be in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/OpDefinition.h:147:19
    #18 0x5583a67143be in mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:10
    #19 0x5583a671402b in (anonymous namespace)::TestLinalgHoisting::runOnOperation() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/test/lib/Dialect/Linalg/TestLinalgHoisting.cpp:48:5
    #20 0x5583a69eda15 in mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:465:11
    #21 0x5583a69ee212 in mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:529:16
    #22 0x5583a69f61bb in mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_14::operator()(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:749:36
    #23 0x5583a69ef9fe in failableParallelForEach<std::__1::__wrap_iter<OpPMInfo *>, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:739:20) &> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Threading.h:46:18
    #24 0x5583a69ef9fe in failableParallelForEach<std::__1::vector<OpPMInfo, std::__1::allocator<OpPMInfo> > &, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:739:20) &> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Threading.h:92:10
    #25 0x5583a69ef9fe in mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:759:14
    #26 0x5583a69ed94a in runOnOperation /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:650:5
    #27 0x5583a69ed94a in mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:463:14
    #28 0x5583a69ee212 in mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:529:16
    #29 0x5583a69f0fa5 in runPasses /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:834:10
    #30 0x5583a69f0fa5 in mlir::PassManager::run(mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:814:60
    #31 0x5583a69e9ff7 in performActions(llvm::raw_ostream&, bool, bool, llvm::SourceMgr&, mlir::MLIRContext*, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:91:17
    #32 0x5583a69e9c47 in processBuffer /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:139:12
    #33 0x5583a69e9c47 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:181:12
    #34 0x5583a69e9c47 in mlir::LogicalResult llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, bool, bool, bool, bool, bool, bool, bool)::$_0>(long, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
    #35 0x5583a6a9253c in llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::operator()(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
    #36 0x5583a6a92703 in mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool)::$_0::operator()(llvm::StringRef) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:88:16
    #37 0x5583a6a923af in interleave<const llvm::StringRef *, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:81:23), (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2085:43), void> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2067:5
    #38 0x5583a6a923af in interleave<llvm::SmallVector<llvm::StringRef, 8U>, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:81:23), llvm::raw_ostream, llvm::StringRef> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2085:3
    #39 0x5583a6a923af in mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:91:3
    #40 0x5583a69e8556 in mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, bool, bool, bool, bool, bool, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:186:10
    #41 0x5583a69e8648 in mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, mlir::PassPipelineCLParser const&, mlir::DialectRegistry&, bool, bool, bool, bool, bool, bool, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:209:10
    #42 0x5583a69e8ca9 in mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:306:14
    #43 0x5583a5124295 in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/tools/mlir-opt/mlir-opt.cpp:244:7
    #44 0x7fe2f378ad8f  (/lib/x86_64-linux-gnu/libc.so.6+0x29d8f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #45 0x7fe2f378ae3f in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x29e3f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #46 0x5583a510b364 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/mlir-opt+0x404e364)
SUMMARY: UndefinedBehaviorSanitizer: undefined-behavior /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Operation.h:165:52 in

==1418730==WARNING: MemorySanitizer: use-of-uninitialized-value
    #0 0x55da63fd5ae3 in use_begin /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/UseDefLists.h:202:56
    #1 0x55da63fd5ae3 in use_begin /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Value.h:189:49
    #2 0x55da63fd5ae3 in getUses /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Value.h:193:39
    #3 0x55da63fd5ae3 in mlir::Value::replaceUsesWithIf(mlir::Value, llvm::function_ref<bool (mlir::OpOperand&)>) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Value.cpp:84:52
    #4 0x55da63431716 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:545:31
    #5 0x55da63431716 in Case<mlir::AffineForOp, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:534:30)> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/TypeSwitch.h:122:16
    #6 0x55da63431716 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:534:12
    #7 0x55da63431716 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Visitors.h:230:14
    #8 0x55da63431716 in mlir::WalkResult llvm::function_ref<mlir::WalkResult (mlir::Operation*)>::callback_fn<std::__1::enable_if<!llvm::is_one_of<mlir::vector::TransferReadOp, mlir::Operation*, mlir::Region*, mlir::Block*>::value && std::is_same<mlir::WalkResult, mlir::WalkResult>::value, mlir::WalkResult>::type mlir::detail::walk<(mlir::WalkOrder)1, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2, mlir::vector::TransferReadOp, mlir::WalkResult>(mlir::Operation*, mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp)::$_2&&)::'lambda'(mlir::Operation*)>(long, mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
    #9 0x55da63fdfb44 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
    #10 0x55da63fdfb44 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:181:12
    #11 0x55da63fdfa78 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #12 0x55da63fdfa78 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #13 0x55da63fdfa78 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #14 0x55da63fdfa78 in mlir::detail::walk(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/IR/Visitors.cpp:174:13
    #15 0x55da6342641a in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::vector::TransferReadOp, mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Visitors.h:233:10
    #16 0x55da6342641a in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Operation.h:579:12
    #17 0x55da6342641a in walk<(mlir::WalkOrder)1, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:15), mlir::WalkResult> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/OpDefinition.h:147:19
    #18 0x55da6342641a in mlir::linalg::hoistRedundantVectorTransfers(mlir::func::FuncOp) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp:425:10
    #19 0x55da63425946 in (anonymous namespace)::TestLinalgHoisting::runOnOperation() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/test/lib/Dialect/Linalg/TestLinalgHoisting.cpp:48:5
    #20 0x55da63bf1614 in mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:465:11
    #21 0x55da63bf258b in mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:529:16
    #22 0x55da63bf62b9 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:749:36
    #23 0x55da63bf62b9 in failableParallelForEach<std::__1::__wrap_iter<OpPMInfo *>, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:739:20) &> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Threading.h:46:18
    #24 0x55da63bf62b9 in failableParallelForEach<std::__1::vector<OpPMInfo, std::__1::allocator<OpPMInfo> > &, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:739:20) &> /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/Threading.h:92:10
    #25 0x55da63bf62b9 in mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:759:14
    #26 0x55da63bf16a2 in runOnOperation /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:650:5
    #27 0x55da63bf16a2 in mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:463:14
    #28 0x55da63bf99d4 in runPipeline /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:529:16
    #29 0x55da63bf99d4 in runPasses /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:834:10
    #30 0x55da63bf99d4 in mlir::PassManager::run(mlir::Operation*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Pass/Pass.cpp:814:60
    #31 0x55da63be8316 in performActions(llvm::raw_ostream&, bool, bool, llvm::SourceMgr&, mlir::MLIRContext*, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:91:17
    #32 0x55da63be79a8 in processBuffer /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:139:12
    #33 0x55da63be79a8 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:181:12
    #34 0x55da63be79a8 in mlir::LogicalResult llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, bool, bool, bool, bool, bool, bool, bool)::$_0>(long, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
    #35 0x55da63d5aec0 in operator() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
    #36 0x55da63d5aec0 in mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool)::$_0::operator()(llvm::StringRef) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:88:16
    #37 0x55da63d5a712 in interleave<const llvm::StringRef *, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:81:23), (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2085:43), void> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2067:5
    #38 0x55da63d5a712 in interleave<llvm::SmallVector<llvm::StringRef, 8U>, (lambda at /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:81:23), llvm::raw_ostream, llvm::StringRef> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/STLExtras.h:2085:3
    #39 0x55da63d5a712 in mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Support/ToolUtilities.cpp:91:3
    #40 0x55da63be2c46 in mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, bool, bool, bool, bool, bool, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:186:10
    #41 0x55da63be3de7 in MlirOptMain /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:209:10
    #42 0x55da63be3de7 in mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:306:14
    #43 0x55da5fe0ae09 in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/tools/mlir-opt/mlir-opt.cpp:244:7
    #44 0x7ff04c884d8f  (/lib/x86_64-linux-gnu/libc.so.6+0x29d8f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #45 0x7ff04c884e3f in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x29e3f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #46 0x55da5fd8d324 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_msan/bin/mlir-opt+0x22bb324)
SUMMARY: MemorySanitizer: use-of-uninitialized-value /b/sanitizer-x86_64-linux-fast/build/llvm-project/mlir/include/mlir/IR/UseDefLists.h:202:56 in use_begin

[jsetoain]

In D137600#3974885, @hctim wrote:

It looks like the update to Dialect/Linalg/hoisting.mlir breaks ASan, MSan and UBSan: https://lab.llvm.org/buildbot/#/builders/5/builds/29773

Yes, I saw that, I'm working on it right now. Thanks for the suggestions!

[jsetoain]

Fix sanitizer issues

[jsetoain]

clang-format

mlir/include/mlir/Dialect/Affine/Utils.h

 /// (noninclusive) and prior to `memOp` (e.g. on a control flow/op path
 /// between the operations) do not have the potential memory effect
 /// `EffectType` on `memOp`. `memOp`  is an operation that reads or writes to
 /// a memref. For example, if `EffectType` is MemoryEffects::Write, this method
 /// will check if there is no write to the memory between `start` and `memOp`
 /// that would change the read within `memOp`.
 template <typename EffectType, typename T>
 bool hasNoInterveningEffect(Operation *start, T memOp);
+
+/// Replace the `Affine for loop` with `newIterOperands` added as new
+/// initialization values. `newYieldValuesFn` is a callback that can be used to
+/// specify the additional values to be yielded by the loop. The number of
+/// values returned by the callback should match the number of new
+/// initialization values. This function
+/// - Moves (i.e. doesn't clone) operations from the `loop` to the newly created
+///   loop
+/// - Replaces the uses of `loop` with the new loop.
+/// - `loop` isn't erased, but is left in a "no-op" state where the body of the
+///   loop just yields the basic block arguments that correspond to the
+///   initialization values of a loop. The loop is dead after this method.
+/// - If `replaceIterOperandsUsesInLoop` is true, all uses of the
+///   `newIterOperands` within the generated new loop are replaced
+///   with the corresponding `BlockArgument` in the loop body.
+using NewYieldValueFn = std::function<SmallVector<Value>(
+    OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs)>;
+AffineForOp
+replaceAffineForWithNewYields(OpBuilder &builder, AffineForOp loop,
+                              ValueRange newIterOperands,
+                              const NewYieldValueFn &newYieldValuesFn,
+                              bool replaceIterOperandsUsesInLoop = true);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_UTILS_H

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

 class ValueRange;
 class Value;
 
 namespace func {
 class CallOp;
 class FuncOp;
 } // namespace func
 
-/// Replace the `loop` with `newIterOperands` added as new initialization
-/// values. `newYieldValuesFn` is a callback that can be used to specify
-/// the additional values to be yielded by the loop. The number of
+/// Replace the `SCF for loop` with `newIterOperands` added as new
+/// initialization values. `newYieldValuesFn` is a callback that can be used to
+/// specify the additional values to be yielded by the loop. The number of
 /// values returned by the callback should match the number of new
 /// initialization values. This function
-/// - Moves (i.e. doesnt clone) operations from the `loop` to the newly created
+/// - Moves (i.e. doesn't clone) operations from the `loop` to the newly created
 ///   loop
 /// - Replaces the uses of `loop` with the new loop.
-/// - `loop` isnt erased, but is left in a "no-op" state where the body of the
+/// - `loop` isn't erased, but is left in a "no-op" state where the body of the
 ///   loop just yields the basic block arguments that correspond to the
 ///   initialization values of a loop. The loop is dead after this method.
 /// - If `replaceIterOperandsUsesInLoop` is true, all uses of the
 ///   `newIterOperands` within the generated new loop are replaced
 ///   with the corresponding `BlockArgument` in the loop body.
 using NewYieldValueFn = std::function<SmallVector<Value>(
     OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs)>;
-scf::ForOp replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
-                                    ValueRange newIterOperands,
-                                    const NewYieldValueFn &newYieldValuesFn,
-                                    bool replaceIterOperandsUsesInLoop = true);
+scf::ForOp
+replaceSCFForWithNewYields(OpBuilder &builder, scf::ForOp loop,
 
 /// Update a perfectly nested loop nest to yield new values from the innermost
 /// loop and propagating it up through the loop nest. This function
 /// - Expects `loopNest` to be a perfectly nested loop with outer most loop
 ///   first and innermost loop last.
 /// - `newIterOperands` are the initialization values to be used for the
 ///    outermost loop
 /// - `newYielValueFn` is the callback that generates the new values to be

mlir/lib/Dialect/Affine/Utils/Utils.cpp

FailureOr<SmallVector<Value>> mlir::delinearizeIndex(OpBuilder &b, Location loc,
   for (Value divisor : divisors) {
     DivModValue divMod = getDivMod(b, loc, residual, divisor);
     results.push_back(divMod.quotient);
     residual = divMod.remainder;
   }
   results.push_back(residual);
   return results;
 }
+
+AffineForOp
+mlir::replaceAffineForWithNewYields(OpBuilder &builder, AffineForOp loop,
+                                    ValueRange newIterOperands,
+                                    const NewYieldValueFn &newYieldValuesFn,
+                                    bool replaceIterOperandsUsesInLoop) {
+  // Create a new loop before the existing one, with the extra operands.
+  OpBuilder::InsertionGuard g(builder);
+  builder.setInsertionPoint(loop);
+  auto operands = llvm::to_vector(loop.getIterOperands());
+  operands.append(newIterOperands.begin(), newIterOperands.end());
+  AffineForOp newLoop = builder.create<AffineForOp>(
+      loop.getLoc(), loop.getLowerBound().getOperands(),
+      loop.getLowerBound().getMap(), loop.getUpperBound().getOperands(),
+      loop.getUpperBound().getMap(), loop.getStep(), operands,
+      [](OpBuilder &, Location, Value, ValueRange) {});
+
+  Block *loopBody = loop.getBody();
+  Block *newLoopBody = newLoop.getBody();
+
+  // Move the body of the original loop to the new loop.
+  newLoopBody->getOperations().splice(newLoopBody->end(),
+                                      loopBody->getOperations());
+
+  // Generate the new yield values to use by using the callback and append the
+  // yield values to the affine.yield operation.
+  auto yield = cast<AffineYieldOp>(newLoopBody->getTerminator());
+  ArrayRef<BlockArgument> newBBArgs =
+      newLoopBody->getArguments().take_back(newIterOperands.size());
+  {
+    OpBuilder::InsertionGuard g(builder);
+    builder.setInsertionPoint(yield);
+    SmallVector<Value> newYieldedValues =
+        newYieldValuesFn(builder, loop.getLoc(), newBBArgs);
+    assert(newIterOperands.size() == newYieldedValues.size() &&
+           "expected as many new yield values as new iter operands");
+    yield.operandsMutable().append(newYieldedValues);
+  }
+
+  // Remap the BlockArguments from the original loop to the new loop
+  // BlockArguments.
+  ArrayRef<BlockArgument> bbArgs = loopBody->getArguments();
+  for (auto it :
+       llvm::zip(bbArgs, newLoopBody->getArguments().take_front(bbArgs.size())))
+    std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
+
+  if (replaceIterOperandsUsesInLoop) {
+    // Replace all uses of `newIterOperands` with the corresponding basic block
+    // arguments.
+    for (auto it : llvm::zip(newIterOperands, newBBArgs)) {
+      std::get<0>(it).replaceUsesWithIf(std::get<1>(it), [&](OpOperand &use) {
+        Operation *user = use.getOwner();
+        return newLoop->isProperAncestor(user);
+      });
+    }
+  }
+
+  // Replace all uses of the original loop with corresponding values from the
+  // new loop.
+  loop.replaceAllUsesWith(
+      newLoop.getResults().take_front(loop.getNumResults()));
+
+  // Add a fake yield to the original loop body that just returns the
+  // BlockArguments corresponding to the iter_args. This makes it a no-op loop.
+  // The loop is dead. The caller is expected to erase it.
+  builder.setInsertionPointToEnd(loopBody);
+  builder.create<AffineYieldOp>(loop->getLoc(), loop.getRegionIterArgs());
+
+  return newLoop;
+}

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

 //===- Hoisting.cpp - Linalg hoisting transformations ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements functions concerned with hoisting invariant operations
 // in the context of Linalg transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
 using llvm::dbgs;
 
 #define DEBUG_TYPE "linalg-hoisting"
 
 #define DBGS() (dbgs() << '[' << DEBUG_TYPE << "] ")
 
else
     yieldOp->setOperand(initArgNumber, write.transferWriteOp.getSource());
 
   // Rewrite `loop` with additional new yields.
   OpBuilder b(read.transferReadOp);
   NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
                                 ArrayRef<BlockArgument> newBBArgs) {
     return SmallVector<Value>{write.transferWriteOp.getVector()};
   };
-  auto newForOp = replaceLoopWithNewYields(
+  auto newForOp = replaceSCFForWithNewYields(
       b, forOp, read.transferReadOp.getVector(), yieldFn);
 
   // Transfer write has been hoisted, need to update the vector and tensor
   // source. Replace the result of the loop to use the new tensor created
   // outside the loop.
   // Depending on whether a insert_slice is present or not, it carries the
   // update on the tensor operands.
   if (write.insertSliceOp) {
func.walk(
         [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
 
     func.walk([&](vector::TransferReadOp transferRead) {
       if (!transferRead.getShapedType().isa<MemRefType>())
         return WalkResult::advance();
 
       LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
                         << *transferRead.getOperation() << "\n");
-      auto loop = dyn_cast<scf::ForOp>(transferRead->getParentOp());
+      auto loop = dyn_cast<LoopLikeOpInterface>(transferRead->getParentOp());
       LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()
                         << "\n");
-      if (!loop)
+      if (!isa_and_nonnull<scf::ForOp, AffineForOp>(loop))
         return WalkResult::advance();
 
       LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()
                         << "\n");
 
       SetVector<Operation *> forwardSlice;
       getForwardSlice(transferRead.getOperation(), &forwardSlice);
 
func.walk([&](vector::TransferReadOp transferRead) {
       transferWrite->moveAfter(loop);
 
       // Rewrite `loop` with new yields by cloning and erase the original loop.
       OpBuilder b(transferRead);
       NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
                                     ArrayRef<BlockArgument> newBBArgs) {
         return SmallVector<Value>{transferWrite.getVector()};
       };
-      auto newForOp =
-          replaceLoopWithNewYields(b, loop, transferRead.getVector(), yieldFn);
 
       // Transfer write has been hoisted, need to update the written vector by
       // the value yielded by the newForOp.
+      return TypeSwitch<Operation *, WalkResult>(loop)
+        .Case<scf::ForOp>([&](scf::ForOp scfForOp) {
+          auto newForOp = replaceSCFForWithNewYields(
+            b, scfForOp, transferRead.getVector(), yieldFn);
           transferWrite.getVectorMutable().assign(newForOp.getResults().back());
-
           changed = true;
           loop.erase();
-      // Need to interrupt and restart because erasing the loop messes up the
-      // walk.
+          // Need to interrupt and restart because erasing the loop messes up
+          // the walk.
           return WalkResult::interrupt();
+        })
+        .Case<AffineForOp>([&](AffineForOp affineForOp) {
+          auto newForOp = replaceAffineForWithNewYields(
+            b, cast<AffineForOp>(loop), transferRead.getVector(), yieldFn);
+          transferWrite.getVectorMutable().assign(newForOp.getResults().back());
+          changed = true;
+          loop.erase();
+          // Need to interrupt and restart because erasing the loop messes up
+          // the walk.
+          return WalkResult::interrupt();
+        })
+        .Default([](Operation *) {
+          return WalkResult::interrupt();
+        });
     });
   }
 }

mlir/lib/Dialect/SCF/Utils/Utils.cpp

 struct LoopParams {
   Value lowerBound;
   Value upperBound;
   Value step;
 };
 } // namespace
 
 scf::ForOp
-mlir::replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
+mlir::replaceSCFForWithNewYields(OpBuilder &builder, scf::ForOp loop,
                                  ValueRange newIterOperands,
                                  const NewYieldValueFn &newYieldValuesFn,
                                  bool replaceIterOperandsUsesInLoop) {
   // Create a new loop before the existing one, with the extra operands.
   OpBuilder::InsertionGuard g(builder);
   builder.setInsertionPoint(loop);
   auto operands = llvm::to_vector(loop.getIterOperands());
   operands.append(newIterOperands.begin(), newIterOperands.end());
   scf::ForOp newLoop = builder.create<scf::ForOp>(
       loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
       operands, [](OpBuilder &, Location, Value, ValueRange) {});
SmallVector<scf::ForOp> mlir::replaceLoopNestWithNewYields(
   //       scf.yield %value
   //     }
   //     scf.yield %2
   //   }
   //   scf.yield %1
   // }
   // ```
   //
-  // The inner most loop is handled using the `replaceLoopWithNewYields`
+  // The innermost loop is handled using the `replaceSCFForWithNewYields`
   // that works on a single loop.
   if (loopNest.size() == 1) {
-    auto innerMostLoop = replaceLoopWithNewYields(
+    auto innerMostLoop = replaceSCFForWithNewYields(
         builder, loopNest.back(), newIterOperands, newYieldValueFn,
         replaceIterOperandsUsesInLoop);
     return {innerMostLoop};
   }
   // The outer loops are modified by calling this method recursively
   // - The return value of the inner loop is the value yielded by this loop.
   // - The region iter args of this loop are the init_args for the inner loop.
   SmallVector<scf::ForOp> newLoopNest;
   NewYieldValueFn fn =
       [&](OpBuilder &innerBuilder, Location loc,
           ArrayRef<BlockArgument> innerNewBBArgs) -> SmallVector<Value> {
     newLoopNest = replaceLoopNestWithNewYields(builder, loopNest.drop_front(),
                                                innerNewBBArgs, newYieldValueFn,
                                                replaceIterOperandsUsesInLoop);
     return llvm::to_vector(llvm::map_range(
         newLoopNest.front().getResults().take_back(innerNewBBArgs.size()),
         [](OpResult r) -> Value { return r; }));
   };
   scf::ForOp outerMostLoop =
-      replaceLoopWithNewYields(builder, loopNest.front(), newIterOperands, fn,
-                               replaceIterOperandsUsesInLoop);
+      replaceSCFForWithNewYields(builder, loopNest.front(), newIterOperands,
+                                 fn, replaceIterOperandsUsesInLoop);
   newLoopNest.insert(newLoopNest.begin(), outerMostLoop);
   return newLoopNest;
 }
 
 /// Outline a region with a single block into a new FuncOp.
 /// Assumes the FuncOp result types is the type of the yielded operands of the
 /// single block. This constraint makes it easy to determine the result.
 /// This method also clones the `arith::ConstantIndexOp` at the start of

mlir/test/Dialect/Linalg/hoisting.mlir

%1 = scf.for %j = %lb to %ub step %step iter_args(%arg5 = %tensor)
     %r01 = vector.transfer_read %w10[%c0, %c3], %cst: tensor<?x?xf32>, vector<2xf32>
     %u01 = "some_use"(%r01) : (vector<2xf32>) -> vector<2xf32>
     %w11 = vector.transfer_write %u01, %w10[%c0, %c3] : vector<2xf32>, tensor<?x?xf32>
     scf.yield %w11 : tensor<?x?xf32>
   }
   return %1 : tensor<?x?xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @hoist_vector_transfer_pairs_in_affine_loops(
+//  CHECK-SAME:   %[[MEMREF0:[a-zA-Z0-9]+]]: memref<64x64xi32>,
+//  CHECK-SAME:   %[[MEMREF1:[a-zA-Z0-9]+]]: memref<64x64xi32>,
+//  CHECK-SAME:   %[[MEMREF2:[a-zA-Z0-9]+]]: memref<64x64xi32>) {
+//       CHECK:   %[[C0:.*]] = arith.constant 0 : i32
+//       CHECK:   affine.for %[[I:.*]] = 0 to 64 {
+//       CHECK:     affine.for %[[J:.*]] = 0 to 64 step 16 {
+//       CHECK:       %[[R0:.*]] = vector.transfer_read %[[MEMREF2]][%[[I]], %[[J]]], %[[C0]] : memref<64x64xi32>, vector<16xi32>
+//       CHECK:       %[[R:.*]] = affine.for %[[K:.*]] = 0 to 64 iter_args(%[[ACC:.*]] = %[[R0]]) -> (vector<16xi32>) {
+//       CHECK:         %[[AV:.*]] = vector.transfer_read %[[MEMREF0]][%[[I]], %[[K]]], %[[C0]] {{.*}}: memref<64x64xi32>, vector<16xi32>
+//       CHECK:         %[[BV:.*]] = vector.transfer_read %[[MEMREF1]][%[[K]], %[[J]]], %[[C0]] {{.*}}: memref<64x64xi32>, vector<16xi32>
+//       CHECK:         %[[T0:.*]] = arith.muli %[[AV]], %[[BV]] : vector<16xi32>
+//       CHECK:         %[[T1:.*]] = arith.addi %[[ACC]], %[[T0]] : vector<16xi32>
+//       CHECK:         affine.yield %[[T1]] : vector<16xi32>
+//       CHECK:       }
+//       CHECK:       vector.transfer_write %[[R]], %[[MEMREF2]][%[[I]], %[[J]]] : vector<16xi32>, memref<64x64xi32>
+//       CHECK:     }
+//       CHECK:   }
+func.func @hoist_vector_transfer_pairs_in_affine_loops(%memref0: memref<64x64xi32>, %memref1: memref<64x64xi32>, %memref2: memref<64x64xi32>) {
+  %c0_i32 = arith.constant 0 : i32
+  affine.for %arg3 = 0 to 64 {
+    affine.for %arg4 = 0 to 64 step 16 {
+      affine.for %arg5 = 0 to 64 {
+        %0 = vector.transfer_read %memref0[%arg3, %arg5], %c0_i32 {permutation_map = affine_map<(d0, d1) -> (0)>} : memref<64x64xi32>, vector<16xi32>
+        %1 = vector.transfer_read %memref1[%arg5, %arg4], %c0_i32 : memref<64x64xi32>, vector<16xi32>
+        %2 = vector.transfer_read %memref2[%arg3, %arg4], %c0_i32 : memref<64x64xi32>, vector<16xi32>
+        %3 = arith.muli %0, %1 : vector<16xi32>
+        %4 = arith.addi %2, %3 : vector<16xi32>
+        vector.transfer_write %4, %memref2[%arg3, %arg4] : vector<16xi32>, memref<64x64xi32>
+      }
+    }
+  }
+  return
+}

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

if (testReplaceWithNewYields) {
           for (auto yieldVal :
                cast<scf::YieldOp>(block->getTerminator()).getResults()) {
             newYieldValues.push_back(
                 b.create<arith::AddFOp>(loc, yieldVal, yieldVal));
           }
           return newYieldValues;
         };
         OpBuilder b(forOp);
-        replaceLoopWithNewYields(b, forOp, newInitValues, fn);
+        replaceSCFForWithNewYields(b, forOp, newInitValues, fn);
       });
     }
   }
 };
 
 struct TestSCFIfUtilsPass
     : public PassWrapper<TestSCFIfUtilsPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFIfUtilsPass)