[Mlir-commits] [mlir] [mlir][tensor] Guard FromElementsOp fold (PR #178996)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tensor

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

FromElementsOp::fold currently attempts to build
a DenseElementsAttr whenever all elem fold results are non-null. This can include ub.poison ops,
which triggers an assertion failure when trying to cast elements to a DenseElementsAttr.
This commit adds a guard to prevent attempting
to fold FromElementsOps that contain ub.poison
elements.

Fixes https://github.com/llvm/llvm-project/issues/178818

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


1 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+18-3) 


``````````diff
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index d885d2c871e3f..6ac4e8dcc531d 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1456,9 +1456,24 @@ void FromElementsOp::build(OpBuilder &builder, OperationState &result,
 }
 
 OpFoldResult FromElementsOp::fold(FoldAdaptor adaptor) {
-  if (!llvm::is_contained(adaptor.getElements(), nullptr))
-    return DenseElementsAttr::get(getType(), adaptor.getElements());
-  return {};
+  auto elements = adaptor.getElements();
+  // If any of the elements is null, we cannot fold.
+  if (llvm::is_contained(elements, nullptr))
+    return {};
+
+  Type elemTy = getType().getElementType();
+
+  // DenseElementsAttr expects element attributes consistent
+  // with the element type. In particular, index/integer tensors
+  // must use IntegerAttr.
+  if (isa<IndexType>(elemTy) || isa<IntegerType>(elemTy)) {
+    for (Attribute a : elements) {
+      auto intAttr = llvm::dyn_cast<IntegerAttr>(a);
+      if (!intAttr || intAttr.getType() != elemTy)
+        return {};
+    }
+  }
+  return DenseElementsAttr::get(getType(), elements);
 }
 
 namespace {

``````````

</details>


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