Page MenuHomePhabricator

[MLIR] Add loop coalesce utility for affine.for
ClosedPublic

Authored by arnab-oss on Aug 16 2021, 5:47 AM.

Details

Summary

Add loop coalesce utility for affine.for. This expects loops to have
been normalized a-priori. This works for both constant as well non
constant upper bounds having single/multiple result upper bound affine
map.

With contributions from Arnab Dutta and Uday Bondhugula.

Diff Detail

Event Timeline

arnab-oss created this revision.Aug 16 2021, 5:47 AM
arnab-oss requested review of this revision.Aug 16 2021, 5:47 AM
This comment was removed by arnab-oss.
This comment was removed by arnab-oss.
This comment was removed by arnab-oss.
bondhugula edited reviewers, added: ayzhuang; removed: avarmapml, pr4tgpt.Aug 17 2021, 5:01 AM

@ftynse Could you please review this? I can review too but I also contributed a part of the code!

bondhugula added inline comments.Aug 21 2021, 7:00 AM
mlir/include/mlir/Transforms/LoopUtils.h
246–249

It would good to mention in what order these loops appear - "outermost to innermost".

a list of ... loops outermost to innermost ...

mlir/lib/Transforms/LoopCoalescing.cpp
100–101

Can we handle both in one walk/pass?

mlir/lib/Transforms/Utils/LoopUtils.cpp
2116

More natural: orgiUbMap.getNumResults() == 1

2127–2128

1 can be on the first line? clang-format

2166–2185

This needs better code comments.

2167

Wrong tense here. Rephrase comment please.

2193

Avoid name like second.

bondhugula requested changes to this revision.Aug 21 2021, 9:52 PM

This already crashes on an input like this:

#map0 = affine_map<(d0) -> (d0 * 110)>
#map1 = affine_map<(d0) -> (696, d0 * 110 + 110)>
func @test(%arg0: memref<2088x2048xf64>, %arg1: memref<2048x2048xf64>, %arg2: memref<2088x2048xf64>) { 
  affine.for %arg3 = 0 to 4 { 
    affine.for %arg4 = 0 to 7 { 
      affine.for %arg5 = 0 to 128 { 
        affine.for %arg6 = #map0(%arg4) to min #map1(%arg4) {
        }
      }
    }
  } 
  return
}
... mlir/lib/IR/AffineMap.cpp:297: int64_t mlir::AffineMap::getSingleConstantResult() const: Assertion `isSingleConstant() && "map must have a single constant result"' failed.
PLEASE submit a bug report to https://bugs.llvm.org/ and include the crash backtrace.
Stack dump:
0.	Program arguments: /home/uday/llvm-project-upstream/build/bin/mlir-opt -loop-coalescing coalesce_crash.mlir
 #0 0x00007f7fa4887c54 PrintStackTraceSignalHandler(void*) Signals.cpp:0:0
 #1 0x00007f7fa488536e SignalHandler(int) Signals.cpp:0:0
 #2 0x00007f7fa42a6210 (/lib/x86_64-linux-gnu/libc.so.6+0x46210)
 #3 0x00007f7fa42a618b raise /build/glibc-eX1tMB/glibc-2.31/signal/../sysdeps/unix/sysv/linux/raise.c:51:1
 #4 0x00007f7fa4285859 abort /build/glibc-eX1tMB/glibc-2.31/stdlib/abort.c:81:7
 #5 0x00007f7fa4285729 get_sysdep_segment_value /build/glibc-eX1tMB/glibc-2.31/intl/loadmsgcat.c:509:8
 #6 0x00007f7fa4285729 _nl_load_domain /build/glibc-eX1tMB/glibc-2.31/intl/loadmsgcat.c:970:34
 #7 0x00007f7fa4296f36 (/lib/x86_64-linux-gnu/libc.so.6+0x36f36)
 #8 0x00007f7fa495e0fb mlir::AffineMap::getSingleConstantResult() const (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRIR.so.14git+0xba0fb)
 #9 0x00007f7fa4d39248 mlir::AffineForOp::getConstantLowerBound() (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRAffine.so.14git+0x4d248)
#10 0x00007f7fa4230048 mlir::coalesceLoops(llvm::MutableArrayRef<mlir::AffineForOp>) (/home/uday/llvm-project-upstream/build/bin/../lib/../lib/libMLIRTransformUtils.so.14git+0x97048)
#11 0x00007f7fa5435045 _ZN4llvm12function_refIFvPN4mlir9OperationEEE11callback_fnIZNS1_6detail4walkILNS1_9WalkOrderE1EZN12_GLOBAL__N_118LoopCoalescingPass13runOnFunctionEvEUlNS1_11AffineForOpEE0_SC_vEENSt9enable_ifIXaantsrNS_11disjunctionIJSt7is_sameIT1_S3_ESG_ISH_PNS1_6RegionEESG_ISH_PNS1_5BlockEEEEE5valuesrSG_IT2_vE5valueESQ_E4typeES3_OT0_EUlS3_E_EEvlS3_ LoopCoalescing.cpp:0:0
#12 0x00007f7fa4a3b2bb mlir::detail::walk(mlir::Operation*, llvm::function_ref<void (mlir::Operation*)>, mlir::WalkOrder) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRIR.so.14git+0x1972bb)
#13 0x00007f7fa5435405 (anonymous namespace)::LoopCoalescingPass::runOnFunction() LoopCoalescing.cpp:0:0
#14 0x00007f7fa54009c5 mlir::FunctionPass::runOnOperation() BufferDeallocation.cpp:0:0
#15 0x00007f7fa512fcdd mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x29cdd)
#16 0x00007f7fa5130020 mlir::detail::OpToOpPassAdaptor::runPipeline(llvm::iterator_range<llvm::pointee_iterator<std::unique_ptr<mlir::Pass, std::default_delete<mlir::Pass> >*, mlir::Pass> >, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x2a020)
#17 0x00007f7fa512e971 mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x28971)
#18 0x00007f7fa512fac7 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x29ac7)
#19 0x00007f7fa5130020 mlir::detail::OpToOpPassAdaptor::runPipeline(llvm::iterator_range<llvm::pointee_iterator<std::unique_ptr<mlir::Pass, std::default_delete<mlir::Pass> >*, mlir::Pass> >, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x2a020)
#20 0x00007f7fa51313c1 mlir::PassManager::run(mlir::Operation*) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x2b3c1)
#21 0x00007f7fa68291eb performActions(llvm::raw_ostream&, bool, bool, llvm::SourceMgr&, mlir::MLIRContext*, mlir::PassPipelineCLParser const&) (.constprop.0) MlirOptMain.cpp:0:0
#22 0x00007f7fa68296c1 processBuffer(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer> >, bool, bool, bool, bool, mlir::PassPipelineCLParser const&, mlir::DialectRegistry&) MlirOptMain.cpp:0:0
#23 0x00007f7fa6829b4e mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer> >, mlir::PassPipelineCLParser const&, mlir::DialectRegistry&, bool, bool, bool, bool, bool) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIROptLib.so.14git+0x7b4e)
#24 0x00007f7fa682b62b mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&, bool) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIROptLib.so.14git+0x962b)
#25 0x0000556e63a101c8 main (/home/uday/llvm-project-upstream/build/bin/mlir-opt+0x191c8)
#26 0x00007f7fa42870b3 __libc_start_main /build/glibc-eX1tMB/glibc-2.31/csu/../csu/libc-start.c:342:3
#27 0x0000556e63a0834e _start (/home/uday/llvm-project-upstream/build/bin/mlir-opt+0x1134e)
Aborted (core dumped)

The test coverage has to be improved as well.

mlir/test/Transforms/loop-coalescing.mlir
229

Instead of MAP, use IDENTITY. Easier to read.

301

Add a comment here as to which scenario it's testing or rename the function to make it descriptive.

337

Also add a negative test case - one that is not supported for coalescing to make sure it doesn't do the wrong thing.

This revision now requires changes to proceed.Aug 21 2021, 9:52 PM
ayzhuang requested changes to this revision.Aug 23 2021, 11:15 AM
ayzhuang added inline comments.
mlir/lib/Transforms/Utils/LoopUtils.cpp
2092

Return failure if !loop.hasConstantLowerBound()

2115

is this check needed? The upper bound map should have at least 1 result.

arnab-oss updated this revision to Diff 370234.EditedSep 2 2021, 5:11 AM
arnab-oss marked 11 inline comments as done.

Updating D108126: [MLIR] Add loop coalesce utility for affine.for

Addressed reviews

*Add extra checks so that code does not crash on running the loop
coalescing pass on unnormalized loop.

*Refactored the code and improved documentation.

*Check for and coalesce affine.for and scf.for loops in a single pass.

This comment was removed by arnab-oss.
Herald added a project: Restricted Project. · View Herald TranscriptSep 2 2021, 5:16 AM
This comment was removed by arnab-oss.
bondhugula accepted this revision.Sep 2 2021, 9:52 AM

LGTM - thanks.

@ayzhuang Do you have any more comments here?

ayzhuang accepted this revision.Sep 3 2021, 8:43 AM

LGTM, thanks.

This revision is now accepted and ready to land.Sep 3 2021, 8:43 AM
bondhugula added a comment.EditedSep 8 2021, 5:22 AM

@arnab-oss Your revision adds five whitespace errors. You can check these with a $ git diff --check HEAD~. I'll fix these before committing this time.

$ git diff --check HEAD~
mlir/test/Transforms/loop-coalescing.mlir:230: trailing whitespace.
+// Check coalescing of affine.for loops when all the loops have non constant upper bounds. 
mlir/test/Transforms/loop-coalescing.mlir:349: trailing whitespace.
+func @test_loops_do_not_get_coalesced() { 
mlir/test/Transforms/loop-coalescing.mlir:350: trailing whitespace.
+  affine.for %i = 0 to 7 { 
mlir/test/Transforms/loop-coalescing.mlir:351: trailing whitespace.
+    affine.for %j = #map0(%i) to min #map1(%i) { 
mlir/test/Transforms/loop-coalescing.mlir:353: trailing whitespace.
+    } 
This revision was landed with ongoing or failed builds.Sep 8 2021, 5:33 AM
This revision was automatically updated to reflect the committed changes.