Merge pull request #205 from Xilinx/tiagot.support_disposable_dense_i…

…nt_allowzero_canonicalization feat: support disposable dense attribute representations for constants in allowzero reshape canonicalization.
Xilinx · Sep 27, 2024 · cc79422 · cc79422
2 parents cfa4445 + 9726dd7
commit cc79422
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/src/Dialect/ONNX/ONNXOps/Canonicalize.cpp b/src/Dialect/ONNX/ONNXOps/Canonicalize.cpp
@@ -329,7 +329,15 @@ bool isConstantOpWithNoZeroElements(Value constVal) {
     return false;
 
   ONNXConstantOp constOp = constVal.getDefiningOp<ONNXConstantOp>();
-  auto intElemsAttr = cast<mlir::DenseIntElementsAttr>(constOp.getValueAttr());
+  DenseElementsAttr intElemsAttr;
+  if (auto elms =
+          dyn_cast<mlir::DenseIntElementsAttr>(constOp.getValueAttr())) {
+    intElemsAttr = elms;
+  } else if (auto elms = dyn_cast<mlir::DisposableElementsAttr>(
+                 constOp.getValueAttr())) {
+    intElemsAttr = dyn_cast_or_null<mlir::DenseIntElementsAttr>(
+        elms.toDenseElementsAttr());
+  }
   if (!intElemsAttr)
     return false;