[Mlir-commits] [mlir] [mlir][linalg] Add Check for Reduction Operation in Contraction Body (PR #123134)

[Ayokunle Amodu]

https://github.com/ayokunle321 created https://github.com/llvm/llvm-project/pull/123134

Fixes issue #122258.

Addresses a crash in the linalg-specialize-generic-ops pass caused by a null pointer dereference when attempting to specialize a linalg generic operation with a reduction iterator. The issue occurs because the pass assumes the presence of a valid reduction operation in the contraction body but does not verify its existence.

Cause: In the isContractionBody function, the getDefiningOp() method was returning nullptr for the yielded value when no operation explicitly defined the reduction result in the block. This led to an invalid dereference and a crash.

A check was added to ensure the presence of a valid reduction operation in the contraction body. The fix verifies that the function returns a false if getDefiningOp() returns a nullptr.

>From 80200160d4c4288f604046440e6dd36a924f0fdc Mon Sep 17 00:00:00 2001
From: Ayokunle Amodu <121697771+ayokunle321 at users.noreply.github.com>
Date: Wed, 15 Jan 2025 15:32:35 -0700
Subject: [PATCH] added check for reduction op in contraction body

---
 mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index caf9cdb3a3eb4f..14f8f9e8fdd3b4 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -281,6 +281,12 @@ bool mlir::linalg::detail::isContractionBody(
 
   Value yielded = getSourceSkipUnary(terminator->getOperand(0));
   Operation *reductionOp = yielded.getDefiningOp();
+
+  if (!reductionOp){
+    errs << "expected reduction op in body";
+    return false;
+  }
+
   if (reductionOp->getNumResults() != 1 || reductionOp->getNumOperands() != 2) {
     errs << "expected reduction op to be binary";
     return false;