-
Notifications
You must be signed in to change notification settings - Fork 0
Lower AtenEmptyMemoryFormat with zero dimensions through tosa-linalg pipeline #6
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: mw_main
Are you sure you want to change the base?
Lower AtenEmptyMemoryFormat with zero dimensions through tosa-linalg pipeline #6
Conversation
for (auto s : shape) | ||
size *= s; | ||
|
||
if (size == 0) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for fixing this, Hari. Do you think it makes sense to create an issue for the upstream if it is legal for tensor to have zero dimensions?
|
||
if (size == 0) { | ||
return rewriter.notifyMatchFailure(op, | ||
"Shape must not have zero dimensions"); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: The error message can say "Shape must not have a dimension of size zero"
// CHECK: %[[EMPTY_TENSOR:.*]] = tensor.empty() : tensor<1x0x256xf32> | ||
// CHECK: return %[[EMPTY_TENSOR]] : tensor<1x0x256xf32> | ||
// CHECK: } | ||
func.func @torch.aten.empty.memory_format() -> !torch.vtensor<[1,0,256],f32>{ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the empty tensor is expected to be lowered later to linalg, an e2e test can help lock down the expected behavior
for (auto s : shape) | ||
size *= s; | ||
|
||
if (size == 0) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This change should be made in the upstream torch-mlir
repo directly instead of our forked repo. Generally, the goal is to keep the divergence with the upstream repo as minimal as possible.
Given that the current tosa pipeline produces invalid tosa.const op we can fix the issue and lock it down using the same torch-to-tosa-backend
test you've added here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we have time to have it accepted upstream and merged with torch-mlir
? Also how do we lock down the torch-to-tosa-linalg
test upstream, I feel like these test tasks will just be lost.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we have time to have it accepted upstream and merged with torch-mlir?
Maybe, it's a simple enough change. In any case, we probably will update torch-mlir during LCM.
we lock down the torch-to-tosa-linalg test upstream
We don't for this particular case. There is already a test in this repo that locks down that bailing out of tosa lowering for a particular unsupported op will trigger the linalg lowering in our pipeline. That is enough to lock down the our pipeline. We don't need to add test for all such unsupported ops.
…eration during folding (llvm#4274) This fixes a SEGFAULT in the GreedyPatternRewriteDriver and adds a missing size check to the `torch.aten._assert_tensor_metadata` operation. Erasing an operation during folding is not allowed. Folding the operation may eithermodify it in place or return a set of replacements, but may not erase the operation. (see https://github.com/llvm/llvm-project/blob/e56384ff540e68f9d0500fa27a95354c0730e37b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp#L492-L508) Doing this causes a SEGFAULT (witnessed on macOS Sequoia 15.5, Apple M4): ``` Stack dump: 0. Program arguments: build/bin/torch-mlir-opt -canonicalize --split-input-file -verify-diagnostics test/Dialect/Torch/invalid_canonicalize.mlir #0 0x0000000104091524 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (build/bin/torch-mlir-opt+0x10140d524) #1 0x000000010408fa5c llvm::sys::RunSignalHandlers() (build/bin/torch-mlir-opt+0x10140ba5c) #2 0x0000000104091bc8 SignalHandler(int, __siginfo*, void*) (build/bin/torch-mlir-opt+0x10140dbc8) #3 0x0000000181e10624 (/usr/lib/system/libsystem_platform.dylib+0x1804ac624) #4 0x0000000103c1f7a8 (anonymous namespace)::GreedyPatternRewriteDriver::processWorklist() (build/bin/torch-mlir-opt+0x100f9b7a8) #5 0x0000000103c1cf4c mlir::applyPatternsGreedily(mlir::Region&, mlir::FrozenRewritePatternSet const&, mlir::GreedyRewriteConfig, bool*) (build/bin/torch-mlir-opt+0x100f98f4c) #6 0x0000000102c8f62c (anonymous namespace)::Canonicalizer::runOnOperation() (build/bin/torch-mlir-opt+0x10000b62c) llvm#7 0x0000000103c72fa4 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (build/bin/torch-mlir-opt+0x100feefa4) llvm#8 0x0000000103c750d4 mlir::PassManager::run(mlir::Operation*) (build/bin/torch-mlir-opt+0x100ff10d4) llvm#9 0x0000000102c8d774 performActions(llvm::raw_ostream&, std::__1::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) (build/bin/torch-mlir-opt+0x100009774) llvm#10 0x0000000102c8d35c llvm::LogicalResult llvm::function_ref<llvm::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>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&)::$_0>(long, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) (build/bin/torch-mlir-opt+0x10000935c) llvm#11 0x000000010403194c mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<llvm::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, llvm::StringRef, llvm::StringRef)::$_0::operator()(llvm::StringRef) const (build/bin/torch-mlir-opt+0x1013ad94c) llvm#12 0x00000001040316a4 mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<llvm::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, llvm::StringRef, llvm::StringRef) (build/bin/torch-mlir-opt+0x1013ad6a4) llvm#13 0x0000000102c87078 mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) (build/bin/torch-mlir-opt+0x100003078) llvm#14 0x0000000102c8731c mlir::MlirOptMain(int, char**, llvm::StringRef, llvm::StringRef, mlir::DialectRegistry&) (build/bin/torch-mlir-opt+0x10000331c) llvm#15 0x0000000102c87538 mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&) (build/bin/torch-mlir-opt+0x100003538) llvm#16 0x0000000102c85cd0 main (build/bin/torch-mlir-opt+0x100001cd0) llvm#17 0x0000000181a36b98 build/tools/torch-mlir/test/Dialect/Torch/Output/invalid_canonicalize.mlir.script: line 1: 72586 Segmentation fault: 11 build/bin/torch-mlir-opt -canonicalize --split-input-file -verify-diagnostics test/Dialect/Torch/invalid_canonicalize.mlir ``` Since the `torch.aten._assert_tensor_metadata` operation is only used for static assertion during compile time the folding can be replaced by a canonicalization that checks the assert and then uses a rewriter to erase the operation. The second commit deals with a missing size check in the assert operation before using a zip operation. Without the explicit checkout of the size, the would assert not fail in case the size of the dimensions were the same, but there are either less or more dimensions in the input than specified in the assert. --------- Signed-off-by: Florian Walbroel <walbroel@roofline.ai> Co-authored-by: Florian Walbroel <walbroel@roofline.ai>
This change handles the case where
torch.aten.empty.memory_format
with 0 sized dimensions fails to lower totosa.const
This was seen in the SAM2 imagePredictor model.We skip the lowering in tosa and use the existing lowering in the linalg pipeline which loweres to
tensor.empty
.