[Mlir-commits] [mlir] [mlir][linalg] Add pattern to bubble-up pack through expand shape op (PR #93529)

Scott Todd llvmlistbot at llvm.org
Tue Jun 25 11:11:10 PDT 2024 

================

----------------
ScottTodd wrote:

Something in this new code is triggering a crash (compiler assert) in the downstream IREE project: https://github.com/iree-org/iree/issues/17734. If I revert this PR locally, the crash goes away.

I don't have a reduced test case yet and the input program is large (12MB) + specific to our downstream project.

* Assert + stack trace:
    ```
    Assertion failed: input.size() == permutation.size() && "expected input rank to equal permutation rank", file D:\dev\projects\iree\third_party\llvm-project\mlir\include\mlir/Dialect/Utils/IndexingUtils.h, line 204
    Please report issues to https://github.com/iree-org/iree/issues and include the crash backtrace.
    Stack dump:
    0.	Program arguments: D:\\dev\\projects\\iree-build\\tools\\iree-compile.exe D:/tmp/open_llama_3b_v2/open-llama-3b-v2-f16.mlir --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-cpu-features=host -o D:/tmp/open_llama_3b_v2/open-llama-3b-v2-f16_cpu.vmfb --mlir-print-ir-before-all --mlir-elide-elementsattrs-if-larger=8 --mlir-elide-resource-strings-if-larger=8 --mlir-disable-threading
    Exception Code: 0x80000003
     #0 0x00007ff64dcf8e95 HandleAbort D:\dev\projects\iree\third_party\llvm-project\llvm\lib\Support\Windows\Signals.inc:425:0
     #1 0x00007ffe7d561881 (C:\WINDOWS\System32\ucrtbase.dll+0x71881)
     #2 0x00007ffe7d562851 (C:\WINDOWS\System32\ucrtbase.dll+0x72851)
     #3 0x00007ffe7d5641b5 (C:\WINDOWS\System32\ucrtbase.dll+0x741b5)
     #4 0x00007ffe7d5644f1 (C:\WINDOWS\System32\ucrtbase.dll+0x744f1)
     #5 0x00007ff651c4f70f mlir::applyPermutation<class llvm::SmallVector<__int64, 2>>(class llvm::ArrayRef<class llvm::SmallVector<__int64, 2>>, class llvm::ArrayRef<__int64>) D:\dev\projects\iree\third_party\llvm-project\mlir\include\mlir\Dialect\Utils\IndexingUtils.h:205:0
     #6 0x00007ff651c49c11 mlir::applyPermutation<class llvm::SmallVector<__int64, 2>>(class llvm::SmallVectorImpl<class llvm::SmallVector<__int64, 2>> const &, class llvm::ArrayRef<__int64>) D:\dev\projects\iree\third_party\llvm-project\mlir\include\mlir\Dialect\Utils\IndexingUtils.h:214:0
     #7 0x00007ff651c41c2b mlir::applyPermutationToVector<class llvm::SmallVector<__int64, 2>, 1>(class llvm::SmallVector<class llvm::SmallVector<__int64, 2>, 1> &, class llvm::ArrayRef<__int64>) D:\dev\projects\iree\third_party\llvm-project\mlir\include\mlir\Dialect\Utils\IndexingUtils.h:225:0
     #8 0x00007ff6531d1eab `anonymous namespace'::applyPermutationAndReindexReassoc D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Dialect\Linalg\Transforms\DataLayoutPropagation.cpp:610:0
     #9 0x00007ff6531d268d `anonymous namespace'::bubbleUpPackOpThroughCollapseShape D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Dialect\Linalg\Transforms\DataLayoutPropagation.cpp:687:0
    #10 0x00007ff6531d3cd6 `anonymous namespace'::BubbleUpPackOpThroughReshapeOp::matchAndRewrite D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Dialect\Linalg\Transforms\DataLayoutPropagation.cpp:849:0
    #11 0x00007ff650b2bbe4 mlir::detail::OpOrInterfaceRewritePatternBase<class mlir::tensor::PackOp>::matchAndRewrite(class mlir::Operation *, class mlir::PatternRewriter &) const D:\dev\projects\iree\third_party\llvm-project\mlir\include\mlir\IR\PatternMatch.h:332:0
    #12 0x00007ff65209e8eb <lambda_033eed04a8a10a7b33015298d48d216a>::operator() D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Rewrite\PatternApplicator.cpp:212:0
    #13 0x00007ff65209c275 mlir::PatternApplicator::matchAndRewrite(class mlir::Operation *, class mlir::PatternRewriter &, class llvm::function_ref<(class mlir::Pattern const &)>, class llvm::function_ref<(class mlir::Pattern const &)>, class llvm::function_ref<(class mlir::Pattern const &)>) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Rewrite\PatternApplicator.cpp:233:0
    #14 0x00007ff650f1f91e `anonymous namespace'::GreedyPatternRewriteDriver::processWorklist D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\GreedyPatternRewriteDriver.cpp:617:0
    #15 0x00007ff650f220e2 llvm::function_ref<void __cdecl(void)>::callback_fn<<lambda_56efa1fe2231a48e07ce9bd5369059af> > D:\dev\projects\iree\third_party\llvm-project\llvm\include\llvm\ADT\STLFunctionalExtras.h:45:0
    #16 0x00007ff650f214ae `anonymous namespace'::RegionPatternRewriteDriver::simplify D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\GreedyPatternRewriteDriver.cpp:872:0
    #17 0x00007ff650f1d38e mlir::applyPatternsAndFoldGreedily(class mlir::Region &, class mlir::FrozenRewritePatternSet const &, class mlir::GreedyRewriteConfig, bool *) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Transforms\Utils\GreedyPatternRewriteDriver.cpp:920:0
    #18 0x00007ff651c78d1d mlir::iree_compiler::GlobalOptimization::`anonymous namespace'::DataLayoutPropagationPass::runOnOperation D:\dev\projects\iree\compiler\src\iree\compiler\GlobalOptimization\DataLayoutPropagation.cpp:31:0
    #19 0x00007ff64e0cead0 llvm::function_ref<void __cdecl(void)>::callback_fn<<lambda_e8f8990a45bf3495636c03506b9db479> > D:\dev\projects\iree\third_party\llvm-project\llvm\include\llvm\ADT\STLFunctionalExtras.h:45:0
    #20 0x00007ff64e0c8637 mlir::detail::OpToOpPassAdaptor::run(class mlir::Pass *, class mlir::Operation *, class mlir::AnalysisManager, bool, unsigned int) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:533:0
    #21 0x00007ff64e0c883d mlir::detail::OpToOpPassAdaptor::runPipeline(class mlir::OpPassManager &, class mlir::Operation *, class mlir::AnalysisManager, bool, unsigned int, class mlir::PassInstrumentor *, struct mlir::PassInstrumentation::PipelineParentInfo const *) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:593:0
    #22 0x00007ff64e0c77bb mlir::detail::OpToOpPassAdaptor::runOnOperationImpl(bool) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:734:0
    #23 0x00007ff64e0ceb23 llvm::function_ref<void __cdecl(void)>::callback_fn<<lambda_e8f8990a45bf3495636c03506b9db479> > D:\dev\projects\iree\third_party\llvm-project\llvm\include\llvm\ADT\STLFunctionalExtras.h:45:0
    #24 0x00007ff64e0c8637 mlir::detail::OpToOpPassAdaptor::run(class mlir::Pass *, class mlir::Operation *, class mlir::AnalysisManager, bool, unsigned int) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:533:0
    #25 0x00007ff64e0c883d mlir::detail::OpToOpPassAdaptor::runPipeline(class mlir::OpPassManager &, class mlir::Operation *, class mlir::AnalysisManager, bool, unsigned int, class mlir::PassInstrumentor *, struct mlir::PassInstrumentation::PipelineParentInfo const *) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:593:0
    #26 0x00007ff64e0c6d7b mlir::PassManager::runPasses(class mlir::Operation *, class mlir::AnalysisManager) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:904:0
    #27 0x00007ff64e0c6b3e mlir::PassManager::run(class mlir::Operation *) D:\dev\projects\iree\third_party\llvm-project\mlir\lib\Pass\Pass.cpp:883:0
    #28 0x00007ff64dca71c7 mlir::iree_compiler::embed::`anonymous namespace'::Invocation::runPipeline D:\dev\projects\iree\compiler\src\iree\compiler\API\Internal\CompilerDriver.cpp:995:0
    #29 0x00007ff64dc657ac <lambda_139d4d9eb9ed714e768e1c22e93f7b10>::operator() D:\dev\projects\iree\compiler\src\iree\compiler\Tools\iree_compile_lib.cc:254:0
    #30 0x00007ff64dc5ba18 mlir::iree_compiler::runIreecMain(int, char **) D:\dev\projects\iree\compiler\src\iree\compiler\Tools\iree_compile_lib.cc:355:0
    #31 0x00007ff658023d34 __scrt_common_main_seh d:\a01\_work\43\s\src\vctools\crt\vcstartup\src\startup\exe_common.inl:288:0
    #32 0x00007ffe7e987344 (C:\WINDOWS\System32\KERNEL32.DLL+0x17344)
    #33 0x00007ffe7f9bcc91 (C:\WINDOWS\SYSTEM32\ntdll.dll+0x4cc91)
    ```

* Here's a bit of printf debugging:

    ```mlir
    // 100s of these, which are fine
    // Calling bubbleUpPackOpThroughCollapseShape with tensor::CollapseShapeOp:
    %collapsed_116 = tensor.collapse_shape %114 [[0, 1], [2], [3]] : tensor<4x32x100x?xf16> into tensor<128x100x?xf16>
    // ...and tensor::PackOp:
    %pack_119 = tensor.pack %collapsed_116 padding_value(%cst_27 : f16) outer_dims_perm = [0, 2, 1] inner_dims_pos = [2, 1] inner_tiles = [8, 1] into %117 : tensor<128x100x?xf16> -> tensor<128x?x100x8x1xf16>
    
    // crash right after this
    // Calling bubbleUpPackOpThroughCollapseShape with tensor::CollapseShapeOp:
    %collapsed_2681 = tensor.collapse_shape %expanded_2680 [[0], [1, 2], [3]] : tensor<4x32x1x100xf16> into tensor<4x32x100xf16>
    // ...and tensor::PackOp:
    %pack_2682 = tensor.pack %collapsed_2681 inner_dims_pos = [0, 2] inner_tiles = [1, 1] into %2209 : tensor<4x32x100xf16> -> tensor<4x32x100x1x1xf16>
    ```

* Here's the IR before we call into this code and crash (13000 lines, can try reducing): https://gist.github.com/ScottTodd/d5f9721307e78cada067a81e60a471c0

https://github.com/llvm/llvm-project/pull/93529

More information about the Mlir-commits mailing list