[Mlir-commits] [mlir] [MLIR][linalg] Handle invalid elementwise kind safely (PR #221835)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

The generic elementwise builder accepts an attribute list. Avoid passing an uninitialized kind to region construction if a caller omits the required kind, and reject attributes of the wrong type rather than relying on an assertion.

Found by Coverity.

Assisted-by: Codex

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


3 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+13-3) 
- (modified) mlir/test/Dialect/Linalg/elementwise/roundtrip.mlir (+3) 
- (modified) mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp (+18) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 9563fcef16d69..09d0ba4d2a31c 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5237,17 +5237,27 @@ void ElementwiseOp::print(OpAsmPrinter &p) {
 void ElementwiseOp::regionBuilder(
     ImplicitLocOpBuilder &b, Block &block, ArrayRef<NamedAttribute> attrs,
     function_ref<InFlightDiagnostic()> emitError) {
-  ElementwiseKind elemwiseKind;
+  std::optional<ElementwiseKind> elemwiseKind;
   for (auto attr : attrs) {
     if (attr.getName() == b.getStringAttr("kind")) {
       auto kindAttr = dyn_cast<ElementwiseKindAttr>(attr.getValue());
-      assert(kindAttr && "op kind attribute incorrectly set");
+      if (!kindAttr) {
+        if (emitError)
+          emitError() << "'kind' must be an ElementwiseKindAttr";
+        return;
+      }
       elemwiseKind = kindAttr.getValue();
       break;
     }
   }
 
-  ArityGroupAndKind groupAndKind = getArityGroupAndKind(elemwiseKind);
+  if (!elemwiseKind) {
+    if (emitError)
+      emitError() << "missing required 'kind' attribute";
+    return;
+  }
+
+  ArityGroupAndKind groupAndKind = getArityGroupAndKind(*elemwiseKind);
   auto arityGroup = groupAndKind.arityGroup;
   auto kind = groupAndKind.kind;
   if (emitError && block.getNumArguments() !=
diff --git a/mlir/test/Dialect/Linalg/elementwise/roundtrip.mlir b/mlir/test/Dialect/Linalg/elementwise/roundtrip.mlir
index 20ebdd992b5a1..29cf1b351deeb 100644
--- a/mlir/test/Dialect/Linalg/elementwise/roundtrip.mlir
+++ b/mlir/test/Dialect/Linalg/elementwise/roundtrip.mlir
@@ -1,4 +1,7 @@
 // RUN: mlir-opt %s -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -split-input-file \
+// RUN:   -test-linalg-elementwise-fusion-patterns=test-invalid-elementwise-kind-builder \
+// RUN:   | FileCheck %s
 //
 // Note - the functions are named @{unary|binary}_{identity|transpose|broadcast|transpose_a|...}_{exp|mul|div|..}
 
diff --git a/mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp b/mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp
index ee0b19006fc62..f6fcb546906d7 100644
--- a/mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp
+++ b/mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp
@@ -143,6 +143,11 @@ struct TestLinalgElementwiseFusion
       llvm::cl::desc("Test fusion of producer ops with multiple uses"),
       llvm::cl::init(false)};
 
+  Option<bool> testInvalidElementwiseKindBuilder{
+      *this, "test-invalid-elementwise-kind-builder",
+      llvm::cl::desc("Test building an elementwise op with an invalid kind"),
+      llvm::cl::init(false)};
+
   ListOption<int64_t> collapseDimensions{
       *this, "collapse-dimensions-control",
       llvm::cl::desc("Test controlling dimension collapse pattern")};
@@ -151,6 +156,19 @@ struct TestLinalgElementwiseFusion
     MLIRContext *context = &this->getContext();
     func::FuncOp funcOp = this->getOperation();
 
+    if (testInvalidElementwiseKindBuilder) {
+      OpBuilder builder = OpBuilder::atBlockBegin(&funcOp.front());
+      linalg::ElementwiseOp missingKindOp = linalg::ElementwiseOp::create(
+          builder, funcOp.getLoc(), ValueRange{}, ValueRange{});
+      missingKindOp.erase();
+
+      NamedAttribute kind = builder.getNamedAttr("kind", builder.getUnitAttr());
+      linalg::ElementwiseOp wrongKindOp = linalg::ElementwiseOp::create(
+          builder, funcOp.getLoc(), ValueRange{}, ValueRange{}, kind);
+      wrongKindOp.erase();
+      return;
+    }
+
     if (fuseGenericOps) {
       RewritePatternSet fusionPatterns(context);
       auto controlFn = [](OpOperand *operand) { return true; };

``````````

</details>


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