lib/Dialect/Torch/IR/TorchOps.cpp: fix: use-after-free: erasing an operation during folding (#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)
 #7 0x0000000103c72fa4 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (build/bin/torch-mlir-opt+0x100feefa4)
 #8 0x0000000103c750d4 mlir::PassManager::run(mlir::Operation*) (build/bin/torch-mlir-opt+0x100ff10d4)
 #9 0x0000000102c8d774 performActions(llvm::raw_ostream&, std::__1::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) (build/bin/torch-mlir-opt+0x100009774)
#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)
#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)
#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)
#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)
#14 0x0000000102c8731c mlir::MlirOptMain(int, char**, llvm::StringRef, llvm::StringRef, mlir::DialectRegistry&) (build/bin/torch-mlir-opt+0x10000331c)
#15 0x0000000102c87538 mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&) (build/bin/torch-mlir-opt+0x100003538)
#16 0x0000000102c85cd0 main (build/bin/torch-mlir-opt+0x100001cd0)
#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>

Author: Manewing

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -14374,7 +14374,7 @@ def Torch_Aten_AssertTensorMetadataOp : Torch_Op<"aten._assert_tensor_metadata",
       printDefaultTorchOp(printer, *this, 6, 0);
     }
   }];
-  let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Torch_AtenDiagonalOp : Torch_Op<"aten.diagonal", [

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -5490,48 +5490,64 @@ OpFoldResult PrimsConvertElementTypeOp::fold(FoldAdaptor adaptor) {
 // Aten_AssertTensorMetadataOp
 //===----------------------------------------------------------------------===//
 
-LogicalResult Aten_AssertTensorMetadataOp::fold(
-    FoldAdaptor adaptor, SmallVectorImpl<::mlir::OpFoldResult> &results) {
-  Value input = getA();
-  auto inputType = cast<BaseTensorType>(input.getType());
-  if (!inputType.hasDtype() || !inputType.hasSizes())
-    return failure();
+namespace {
+class EraseAssertMetadataPattern
+    : public OpRewritePattern<Aten_AssertTensorMetadataOp> {
+public:
+  using OpRewritePattern<Aten_AssertTensorMetadataOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(Aten_AssertTensorMetadataOp op,
+                                PatternRewriter &rewriter) const override {
+    Value input = op.getA();
+    auto inputType = cast<BaseTensorType>(input.getType());
+    if (!inputType.hasDtype() || !inputType.hasSizes())
+      return failure();
 
-  // TODO: Add checks for stride, device, and layout when we can extract that
-  // information from the torch tensor. For now, we can only get the shape and
-  // dtype info from the tensor hence adding checks for them.
+    // TODO: Add checks for stride, device, and layout when we can extract that
+    // information from the torch tensor. For now, we can only get the shape and
+    // dtype info from the tensor hence adding checks for them.
 
-  // convert size to a list of integers.
-  SmallVector<int64_t> size;
-  if (!isa<Torch::NoneType>(getSize().getType())) {
-    if (!matchPattern(getSize(), m_TorchListOfConstantInts(size))) {
-      return emitOpError("expected dtype to be a constant int");
+    // convert size to a list of integers.
+    SmallVector<int64_t> size;
+    if (!isa<Torch::NoneType>(op.getSize().getType())) {
+      if (!matchPattern(op.getSize(), m_TorchListOfConstantInts(size))) {
+        return op.emitOpError("expected dtype to be a constant int");
+      }
+      if (inputType.getSizes().size() != size.size() ||
+          !llvm::all_of(llvm::zip(inputType.getSizes(), size),
+                        [](const auto &pair) {
+                          return std::get<0>(pair) == std::get<1>(pair);
+                        }))
+        return op.emitOpError(
+            "Failed to canonicalize the _assert_tensor_metadata op since "
+            "the sizes do not match");
     }
-    if (!llvm::all_of(llvm::zip(inputType.getSizes(), size),
-                      [](const auto &pair) {
-                        return std::get<0>(pair) == std::get<1>(pair);
-                      }))
-      return emitOpError("Failed to fold the _assert_tensor_metadata op since "
-                         "the sizes do not match");
-  }
 
-  // convert dtype to an integer.
-  int64_t dtype;
-  if (!isa<Torch::NoneType>(getDtype().getType())) {
-    if (!matchPattern(getDtype(), m_TorchConstantInt(&dtype))) {
-      return emitOpError("expected dtype to be a constant int");
+    // convert dtype to an integer.
+    int64_t dtype;
+    if (!isa<Torch::NoneType>(op.getDtype().getType())) {
+      if (!matchPattern(op.getDtype(), m_TorchConstantInt(&dtype))) {
+        return op.emitOpError("expected dtype to be a constant int");
+      }
+      FailureOr<Type> inputDtype =
+          getTypeForScalarType(getContext(), (torch_upstream::ScalarType)dtype);
+      if (failed(inputDtype))
+        return failure();
+      if (inputType.getDtype() != inputDtype)
+        return op.emitOpError(
+            "Failed to canonicalize the _assert_tensor_metadata op since "
+            "the dtype does not match");
     }
-    FailureOr<Type> inputDtype =
-        getTypeForScalarType(getContext(), (torch_upstream::ScalarType)dtype);
-    if (failed(inputDtype))
-      return failure();
-    if (inputType.getDtype() != inputDtype)
-      return emitOpError("Failed to fold the _assert_tensor_metadata op since "
-                         "the dtype does not match");
+
+    rewriter.eraseOp(op);
+    return success();
   }
+};
+} // namespace
 
-  getOperation()->erase();
-  return success();
+void Aten_AssertTensorMetadataOp::getCanonicalizationPatterns(
+    RewritePatternSet &patterns, MLIRContext *context) {
+  patterns.add<EraseAssertMetadataPattern>(context);
 }
 
 //===----------------------------------------------------------------------===//

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -1044,7 +1044,7 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::as_strided : (Tensor, int[], int[], int?) -> (Tensor)")
     emit(
         "aten::_assert_tensor_metadata : (Tensor, int[]?, int[]?, int?, Device?, int?) -> ()",
-        has_folder=True,
+        has_canonicalizer=True,
     )
     emit("aten::diagonal : (Tensor, int, int, int) -> (Tensor)")
     emit("aten::diagonal_copy : (Tensor, int, int, int) -> (Tensor)")

test/Dialect/Torch/canonicalize.mlir

@@ -29,6 +29,17 @@ func.func @torch.runtime.assert() {
   return
 }
 
+// CHECK-LABEL:   func.func @torch.aten.assert_tensor_metadata
+// CHECK-NEXT:      return
+func.func @torch.aten.assert_tensor_metadata() {
+  %int4 = torch.constant.int 4
+  %none = torch.constant.none
+  %1 = tensor.empty() : tensor<1x1x128x128xi64>
+  %2 = torch_c.from_builtin_tensor %1 : tensor<1x1x128x128xi64> -> !torch.vtensor<[1,1,128,128],si64>
+  torch.aten._assert_tensor_metadata %2, %none, %none, %int4, %none, %none : !torch.vtensor<[1,1,128,128],si64>, !torch.none, !torch.none, !torch.int, !torch.none, !torch.none
+  return
+}
+
 // CHECK-LABEL:   func.func @torch.aten.ones_item
 // CHECK:           %[[CONST:.*]] = torch.constant.int 1
 // CHECK:           return %[[CONST]] : !torch.int

test/Dialect/Torch/invalid_canonicalize.mlir

@@ -0,0 +1,40 @@
+// RUN: torch-mlir-opt -canonicalize --split-input-file -verify-diagnostics %s
+
+func.func @torch.aten.assert_tensor_metadata_invalid_dtype() {
+  %int8 = torch.constant.int 8
+  %none = torch.constant.none
+  %1 = tensor.empty() : tensor<1x1x128x128xi64>
+  %2 = torch_c.from_builtin_tensor %1 : tensor<1x1x128x128xi64> -> !torch.vtensor<[1,1,128,128],si64>
+  // expected-error @+1 {{torch.aten._assert_tensor_metadata' op Failed to canonicalize the _assert_tensor_metadata op since the dtype does not match}}
+  torch.aten._assert_tensor_metadata %2, %none, %none, %int8, %none, %none : !torch.vtensor<[1,1,128,128],si64>, !torch.none, !torch.none, !torch.int, !torch.none, !torch.none
+  return
+}
+
+func.func @torch.aten.assert_tensor_metadata_invalid_size() {
+  %int0 = torch.constant.int 0
+  %int2 = torch.constant.int 2
+  %int3 = torch.constant.int 3
+  %sizes = torch.prim.ListConstruct %int0, %int2, %int3
+        : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  %int4 = torch.constant.int 4
+  %none = torch.constant.none
+  %1 = tensor.empty() : tensor<1x1x128x128xi64>
+  %2 = torch_c.from_builtin_tensor %1 : tensor<1x1x128x128xi64> -> !torch.vtensor<[1,1,128,128],si64>
+  // expected-error @+1 {{'torch.aten._assert_tensor_metadata' op Failed to canonicalize the _assert_tensor_metadata op since the sizes do not match}}
+  torch.aten._assert_tensor_metadata %2, %sizes, %none, %int4, %none, %none : !torch.vtensor<[1,1,128,128],si64>, !torch.list<int>, !torch.none, !torch.int, !torch.none, !torch.none
+  return
+}
+
+func.func @torch.aten.assert_tensor_metadata_invalid_size_extra_dim() {
+  %int1 = torch.constant.int 1
+  %int4 = torch.constant.int 4
+  %int128 = torch.constant.int 128
+  %sizes = torch.prim.ListConstruct %int1, %int1, %int128, %int128, %int4
+        : (!torch.int, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  %none = torch.constant.none
+  %1 = tensor.empty() : tensor<1x1x128x128xi64>
+  %2 = torch_c.from_builtin_tensor %1 : tensor<1x1x128x128xi64> -> !torch.vtensor<[1,1,128,128],si64>
+  // expected-error @+1 {{'torch.aten._assert_tensor_metadata' op Failed to canonicalize the _assert_tensor_metadata op since the sizes do not match}}
+  torch.aten._assert_tensor_metadata %2, %sizes, %none, %int4, %none, %none : !torch.vtensor<[1,1,128,128],si64>, !torch.list<int>, !torch.none, !torch.int, !torch.none, !torch.none
+  return
+}