[Mlir-commits] [mlir] d926b33 - [mlir] add complex type to getZeroAttr

[Aart Bik]

Author: Aart Bik
Date: 2022-07-07T16:58:59-07:00
New Revision: d926b3307e2f90d8b7f770e6f40fbbe5dcc1f887

URL: https://github.com/llvm/llvm-project/commit/d926b3307e2f90d8b7f770e6f40fbbe5dcc1f887
DIFF: https://github.com/llvm/llvm-project/commit/d926b3307e2f90d8b7f770e6f40fbbe5dcc1f887.diff

LOG: [mlir] add complex type to getZeroAttr

Fixes issue encountered with <sparse> complex constant
https://github.com/llvm/llvm-project/issues/56428

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D129325

Added: 
    mlir/test/Integration/Dialect/LLVMIR/CPU/test-complex-sparse-constant.mlir

Modified: 
    mlir/lib/IR/BuiltinAttributes.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 35926522e0805..80b91ca5bb050 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -1589,6 +1589,18 @@ Attribute SparseElementsAttr::getZeroAttr() const {
   if (eltType.isa<FloatType>())
     return FloatAttr::get(eltType, 0);
 
+  // Handle complex elements.
+  if (auto complexTy = eltType.dyn_cast<ComplexType>()) {
+    auto eltType = complexTy.getElementType();
+    Attribute zero;
+    if (eltType.isa<FloatType>())
+      zero = FloatAttr::get(eltType, 0);
+    else // must be integer
+      zero = IntegerAttr::get(eltType, 0);
+    return ArrayAttr::get(complexTy.getContext(),
+                          ArrayRef<Attribute>{zero, zero});
+  }
+
   // Handle string type.
   if (getValues().isa<DenseStringElementsAttr>())
     return StringAttr::get("", eltType);

diff  --git a/mlir/test/Integration/Dialect/LLVMIR/CPU/test-complex-sparse-constant.mlir b/mlir/test/Integration/Dialect/LLVMIR/CPU/test-complex-sparse-constant.mlir
new file mode 100644
index 0000000000000..a42c64d57a871
--- /dev/null
+++ b/mlir/test/Integration/Dialect/LLVMIR/CPU/test-complex-sparse-constant.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-opt %s --convert-memref-to-llvm | \
+// RUN:   mlir-cpu-runner -e entry -entry-point-result=void
+
+//
+// Code should not crash on the complex32 sparse constant.
+//
+module attributes {llvm.data_layout = ""} {
+  memref.global "private" constant @"__constant_32xcomplex<f32>_0" : memref<32xcomplex<f32>> =
+     sparse<[[1], [28], [31]],
+            [(1.000000e+00,0.000000e+00), (2.000000e+00,0.000000e+00), (3.000000e+00,0.000000e+00)]
+	    > {alignment = 128 : i64}
+  llvm.func @entry() {
+     %0 = memref.get_global @"__constant_32xcomplex<f32>_0" : memref<32xcomplex<f32>>
+     llvm.return
+  }
+}