Add default DataLayout support for complex numbers

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

diff --git a/mlir/docs/DataLayout.md b/mlir/docs/DataLayout.md
index 6d2aab8..caa90cf 100644 (file)
--- a/mlir/docs/DataLayout.md
+++ b/mlir/docs/DataLayout.md
@@ -271,6 +271,14 @@ those of the integer type with the same bitwidth defined above.
 In absence of the corresponding entry, `index` is assumed to be a 64-bit
 integer.
 
+#### `complex` type
+
+By default complex type is treated like a 2 element structure of its given
+element type. This is to say that each of its elements are aligned to their
+preferred alignment, the entire complex type is also aligned to this preference,
+and the complex type size includes the possible padding between elements to enforce
+alignment.
+
 ### Byte Size
 
 The default data layout assumes 8-bit bytes.

diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index 3369d61..1f219c1 100644 (file)
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -13,6 +13,7 @@
 #include "mlir/IR/Operation.h"
 
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/MathExtras.h"
 
 using namespace mlir;
 
@@ -53,6 +54,17 @@ unsigned mlir::detail::getDefaultTypeSizeInBits(Type type,
   if (type.isa<IntegerType, FloatType>())
     return type.getIntOrFloatBitWidth();
 
+  if (auto ctype = type.dyn_cast<ComplexType>()) {
+    auto et = ctype.getElementType();
+    auto innerAlignment =
+        getDefaultPreferredAlignment(et, dataLayout, params) * 8;
+    auto innerSize = getDefaultTypeSizeInBits(et, dataLayout, params);
+
+    // Include padding required to align the imaginary value in the complex
+    // type.
+    return llvm::alignTo(innerSize, innerAlignment) + innerSize;
+  }
+
   // Index is an integer of some bitwidth.
   if (type.isa<IndexType>())
     return dataLayout.getTypeSizeInBits(
@@ -92,6 +104,9 @@ unsigned mlir::detail::getDefaultABIAlignment(
                : 4;
   }
 
+  if (auto ctype = type.dyn_cast<ComplexType>())
+    return getDefaultABIAlignment(ctype.getElementType(), dataLayout, params);
+
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getABIAlignment(dataLayout, params);
 
@@ -110,6 +125,10 @@ unsigned mlir::detail::getDefaultPreferredAlignment(
   if (type.isa<IntegerType, IndexType>())
     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
 
+  if (auto ctype = type.dyn_cast<ComplexType>())
+    return getDefaultPreferredAlignment(ctype.getElementType(), dataLayout,
+                                        params);
+
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getPreferredAlignment(dataLayout, params);
 

diff --git a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
index 41f21d9..c7f2582 100644 (file)
--- a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
+++ b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
@@ -12,7 +12,18 @@ func @no_layout_builtin() {
   // CHECK: preferred = 8
   // CHECK: size = 8
   "test.data_layout_query"() : () -> f64
+  // CHECK: alignment = 4
+  // CHECK: bitsize = 64
+  // CHECK: preferred = 4
+  // CHECK: size = 8
+  "test.data_layout_query"() : () -> complex<f32>
+  // CHECK: alignment = 1
+  // CHECK: bitsize = 14
+  // CHECK: preferred = 1
+  // CHECK: size = 2
+  "test.data_layout_query"() : () -> complex<i6>
   return
+
 }
 
 // CHECK-LABEL: @no_layout_custom