Check linalg.generic arguments to prevent crashing when they are deleted

[parsifal-47]

During remove-dead-values pass, values are getting deleted and if they appear as arguments to dead operations values disappear before operations are deleted. In this case, arguments may be null. This is a temporary state while the IR is being processed.
When remove-dead-values is trying to delete linalg.generic, it calls isMemoryEffectFree and it iterates over the arguments, some of which are already deleted. This is causing a crash: #118450 which is added as a test.
Thank you!

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Renat Idrisov (parsifal-47)

Changes

During remove-dead-values pass, values are getting deleted and if they appear as arguments to dead operations values disappear before operations are deleted. In this case, arguments may be null. This is a temporary state while the IR is being processed.
When remove-dead-values is trying to delete linalg.generic, it calls isMemoryEffectFree and it iterates over the arguments, some of which are already deleted. This is causing a crash: #118450 which is added as a test.
Thank you!

---

Full diff: https://github.com/llvm/llvm-project/pull/119110.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+2-2)
• (modified) mlir/test/Transforms/remove-dead-values.mlir (+26)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index d9840e3923c4f7..c5b1a3b55126c8 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1199,7 +1199,7 @@ static void getGenericEffectsImpl(
         &effects,
     LinalgOp linalgOp) {
   for (auto [index, operand] : llvm::enumerate(linalgOp.getDpsInputs())) {
-    if (!llvm::isa<MemRefType>(operand.getType()))
+    if (!operand || !llvm::isa<MemRefType>(operand.getType()))
       continue;
     effects.emplace_back(
         MemoryEffects::Read::get(), &linalgOp->getOpOperand(index), /*stage=*/0,
@@ -1207,7 +1207,7 @@ static void getGenericEffectsImpl(
   }
 
   for (OpOperand &operand : linalgOp.getDpsInitsMutable()) {
-    if (!llvm::isa<MemRefType>(operand.get().getType()))
+    if (!operand.get() || !llvm::isa<MemRefType>(operand.get().getType()))
       continue;
     if (linalgOp.payloadUsesValueFromOperand(&operand)) {
       effects.emplace_back(MemoryEffects::Read::get(), &operand, /*stage=*/0,
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 9273ac01e7ccec..fe7bcbc7c490b6 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -73,6 +73,32 @@ func.func @cleanable_loop_iter_args_value(%arg0: index) -> index {
 
 // -----
 
+// Checking that the arguments of linalg.generic are properly handled
+// All code below is removed as a result of the pass
+//
+#map = affine_map<(d0, d1, d2) -> (0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+module {
+  func.func @main() {
+    %cst_3 = arith.constant dense<54> : tensor<1x25x13xi32>
+    %cst_7 = arith.constant dense<11> : tensor<1x25x13xi32>
+    // CHECK-NOT: arith.constant
+    %0 = tensor.empty() : tensor<1x25x13xi32>
+    // CHECK-NOT: tensor
+    %1 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%cst_3, %cst_7 : tensor<1x25x13xi32>, tensor<1x25x13xi32>) outs(%0 : tensor<1x25x13xi32>) {
+    // CHECK-NOT: linalg.generic
+    ^bb0(%in: i32, %in_15: i32, %out: i32):
+      %29 = arith.xori %in, %in_15 : i32
+      // CHECK-NOT: arith.xori
+      linalg.yield %29 : i32
+      // CHECK-NOT: linalg.yield
+    } -> tensor<1x25x13xi32>
+    return
+  }
+}
+
+// -----
+
 // Note that this cleanup cannot be done by the `canonicalize` pass.
 //
 // CHECK-LABEL: func.func private @clean_func_op_remove_argument_and_return_value() {

[parsifal-47]

@CoTinker please take a look, thank you!

[github-actions]

⚠️ We detected that you are using a GitHub private e-mail address to contribute to the repo.
Please turn off Keep my email addresses private setting in your account.
See LLVM Discourse for more information.

[parsifal-47]

@dcaballe @nicolasvasilache @rengolin could you please take a look? Thank you!

[joker-eph]

I'm not sure this fix is correct: seems like working around an issue with the pass: seems like the pass should batch the changes and apply them later instead of eagerly applying the change, breaking the IR, but trying to call functions like isMemoryEffectFree() on the broken IR. I'm afraid that this is getting in a rabbit hole of trying to make every IR helper robust to broken IR, which isn't tractable (nor is the way we usually try to tackle this).

[parsifal-47]

I'm not sure this fix is correct: seems like working around an issue with the pass: seems like the pass should batch the changes and apply them later instead of eagerly applying the change, breaking the IR, but trying to call functions like isMemoryEffectFree() on the broken IR. I'm afraid that this is getting in a rabbit hole of trying to make every IR helper robust to broken IR, which isn't tractable (nor is the way we usually try to tackle this).

I understand, this may be invalid argument, but other operations are unlike linalg.generic they provide "safe" isMemoryEffectFree function which does not break when args are null.

What would you recommend, precompute this function for all nodes in IR and then use it when walking?

Thank you!

[parsifal-47]

I'm not sure this fix is correct: seems like working around an issue with the pass: seems like the pass should batch the changes and apply them later instead of eagerly applying the change, breaking the IR, but trying to call functions like isMemoryEffectFree() on the broken IR. I'm afraid that this is getting in a rabbit hole of trying to make every IR helper robust to broken IR, which isn't tractable (nor is the way we usually try to tackle this).

Another approach to propose is to remove simple operation if at least one of its arguments is already removed without checking for memory safety because this check has been complete on Liveness analysis stage and since the argument is dead, the simple operation is dead and memory safe.

[banach-space]

Hi @parsifal-47 ,

I'm not that familiar with this area, but will try to help.

other operations are unlike linalg.generic they provide "safe" isMemoryEffectFree function which does not break when args are null.

Do you have links? I'm curious where the discrepancy is coming from.

What would you recommend, precompute this function for all nodes in IR and then use it when walking?

IIUC, the suggestion is to pre-compute things to delete, then remove the operation and then the arguments. Perhaps the latter won't be needed if the only consumer is deleted.

Another approach to propose is to remove simple operation if at least one of its arguments is already removed

Is this valid though: "if one argument of Op has been deleted --> Op can be deleted"?

[parsifal-47]

Hi @banach-space,
thank you for helping

other operations are unlike linalg.generic they provide "safe" isMemoryEffectFree function which does not break when args are null.

Do you have links? I'm curious where the discrepancy is coming from.

I browsed the source code and I haven't found any operation which checks its arguments upon isMemoryEffectFree. They are safe because of semantics is different and complex operations like structured control flow is handled separately.

What would you recommend, precompute this function for all nodes in IR and then use it when walking?

IIUC, the suggestion is to pre-compute things to delete, then remove the operation and then the arguments. Perhaps the latter won't be needed if the only consumer is deleted.

Two walks within the pass, one to compute and one to delete? Does not look natural, maybe there is an analysis like Liveness which could be relied on.

Another approach to propose is to remove simple operation if at least one of its arguments is already removed

Is this valid though: "if one argument of Op has been deleted --> Op can be deleted"?

That is the assumption, simple operation cannot survive the removal of positional arguments. If one argument is removed and operation is not, it is a crash.

I am fine with any of these paths and I want to do it right, but it is a bit unclear if I need an attention of code owners and how to get it.
Thank you!

[banach-space]

it is a bit unclear if I need an attention of code owners and how to get it.

Sadly, there are no code-owners in MLIR. Things are bound to improve, see:

• https://discourse.llvm.org/t/rfc-mlir-project-charter-and-restructuring/

However, it will take same time and I suspect you'd like your fix to land sooner than that ;-) I'd normally just ping whoever touched relevant code most recently, but I haven't noticed any familiar (i.e. "active") names. I can try my best to help instead. Mehdi is one of our core maintainers, so we should definitely make sure we act on his feedback.

Now, I see that things go wrong in this hook:

llvm-project/mlir/lib/Transforms/RemoveDeadValues.cpp

Lines 616 to 637 in f85579f

	void RemoveDeadValues::runOnOperation() {
	auto &la = getAnalysis<RunLivenessAnalysis>();
	Operation *module = getOperation();
	module->walk([&](Operation *op) {
	if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
	cleanFuncOp(funcOp, module, la);
	} else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
	cleanRegionBranchOp(regionBranchOp, la);
	} else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
	cleanBranchOp(branchOp, la);
	} else if (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {
	// Nothing to do here because this is a terminator op and it should be
	// honored with respect to its parent
	} else if (isa<CallOpInterface>(op)) {
	// Nothing to do because this op is associated with a function op and gets
	// cleaned when the latter is cleaned.
	} else {
	cleanSimpleOp(op, la);
	}
	});
	}

Clearly, what's currently happening is incorrect:

1. Compute liveness analysis.
2. Walk overall all Ops in a module and remove them if they are not alive and have no memory effects.

Your test case demonstrates that this logic brakes if an operand of an Op is removed before the corresponding Op is processed. The logic needs to be updated.

Two walks within the pass, one to compute and one to delete? Does not look natural, maybe there is an analysis like Liveness which could be relied on.

It won't be as compact, but it will be correct. Correctness is much higher priority. I'd do it in two steps (this is probably what you had in mind):

1. Walk over all Ops within a module and "collect" the ones that are not live and have no memory effects. Lets call the corresponding container (probably a vector): opsToErase.
2. Go over all Ops in opsToErase and remove them (without checking for memory effects).

Perhaps you will have to process "simple" Ops and "everything else" (e.g. functions) separately. I am almost definitely missing some fine details or nuance here ;-)

You may need to experiment a bit. Ultimately, you want to make sure to address Mehdi's feedback and to preserve current functionality.

Once you have an updated version, I am happy to review and we can try pinging somebody more active in this area (though no guarantee they will reply).

Thanks for working on this!

[parsifal-47]

@banach-space thanks a lot for detailed review, I appreciate it, let me work in the better proposal which would follow the correct scheme. Since this is a completely different diff, I am going to close this PR and open another one once I am ready. I will keep you posted, thank you!

[banach-space]

SG, Renat.

One additional data/design point to keep in mind is that right now, mlir::isMemoryEffectFree, receives this invalid linalg.generic:

 "linalg.generic"(<<NULL VALUE>>, <<NULL VALUE>>, <<NULL VALUE>>) <{indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, d2)>, affine_map<(d0, d1, d2) -> (0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>], operandSegmentSizes = array<i32: 2, 1>}> ({
^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
  "linalg.yield"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> ()
}) : (<<NULL TYPE>>, <<NULL TYPE>>, <<NULL TYPE>>) -> tensor<1x25x13xi32>

It shouldn't ;-)