add layout test, address reviewer comments and fix a bug in layout methods

Author: fabianmcg

mlir/include/mlir/Dialect/Ptr/IR/PtrDialect.td

@@ -55,8 +55,8 @@ def Ptr_PtrType : Ptr_Type<"Ptr", "ptr", [
   let parameters = (ins OptionalParameter<"Attribute">:$memorySpace);
   let assemblyFormat = "(`<` $memorySpace^ `>`)?";
   let builders = [
-    TypeBuilder<(ins CArg<"Attribute", "nullptr">:$addressSpace), [{
-      return $_get($_ctxt, addressSpace);
+    TypeBuilder<(ins CArg<"Attribute", "nullptr">:$memorySpace), [{
+      return $_get($_ctxt, memorySpace);
     }]>,
     TypeBuilder<(ins CArg<"unsigned">:$addressSpace), [{
       return $_get($_ctxt, IntegerAttr::get(IntegerType::get($_ctxt, 32),
@@ -69,7 +69,7 @@ def Ptr_PtrType : Ptr_Type<"Ptr", "ptr", [
     Attribute getDefaultMemorySpace() const;
 
     /// Returns the memory space as an unsigned number.
-    int64_t getAddressSpace() const;
+    uint64_t getAddressSpace() const;
   }];
 }
 

mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp

@@ -36,35 +36,32 @@ getPointerDataLayoutEntry(DataLayoutEntryListRef params, PtrType type,
   for (DataLayoutEntryInterface entry : params) {
     if (!entry.isTypeEntry())
       continue;
-    if (cast<PtrType>(entry.getKey().get<Type>()).getAddressSpace() ==
-        type.getAddressSpace()) {
+    if (cast<PtrType>(entry.getKey().get<Type>()).getMemorySpace() ==
+        type.getMemorySpace()) {
       currentEntry = entry.getValue();
       break;
     }
   }
+  bool isSizeOrIndex =
+      pos == PtrDLEntryPos::Size || pos == PtrDLEntryPos::Index;
   if (currentEntry) {
     std::optional<uint64_t> value = extractPointerSpecValue(currentEntry, pos);
     // If the optional `PtrDLEntryPos::Index` entry is not available, use the
     // pointer size as the index bitwidth.
     if (!value && pos == PtrDLEntryPos::Index)
       value = extractPointerSpecValue(currentEntry, PtrDLEntryPos::Size);
-    bool isSizeOrIndex =
-        pos == PtrDLEntryPos::Size || pos == PtrDLEntryPos::Index;
     return *value / (isSizeOrIndex ? 1 : kBitsInByte);
   }
 
   // If not found, and this is the pointer to the default memory space, assume
   // 64-bit pointers.
-  if (type.getAddressSpace() == 0) {
-    bool isSizeOrIndex =
-        pos == PtrDLEntryPos::Size || pos == PtrDLEntryPos::Index;
+  if (type.getMemorySpace() == type.getDefaultMemorySpace())
     return isSizeOrIndex ? kDefaultPointerSizeBits : kDefaultPointerAlignment;
-  }
 
   return std::nullopt;
 }
 
-int64_t PtrType::getAddressSpace() const { return 0; }
+uint64_t PtrType::getAddressSpace() const { return 0; }
 
 Attribute PtrType::getDefaultMemorySpace() const { return nullptr; }
 
@@ -85,9 +82,10 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
           return false;
         });
     if (it == oldLayout.end()) {
-      llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
+      it = llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
         if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey())) {
-          return llvm::cast<PtrType>(type).getAddressSpace() == 0;
+          auto ptrTy = llvm::cast<PtrType>(type);
+          return ptrTy.getMemorySpace() == ptrTy.getDefaultMemorySpace();
         }
         return false;
       });

mlir/test/Dialect/Ptr/layout.mlir

@@ -0,0 +1,87 @@
+// RUN: mlir-opt --test-data-layout-query --split-input-file --verify-diagnostics %s | FileCheck %s
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!ptr.ptr, dense<[32, 32, 64]> : vector<3xi64>>,
+  #dlti.dl_entry<!ptr.ptr<5>, dense<[64, 64, 64]> : vector<3xi64>>,
+  #dlti.dl_entry<!ptr.ptr<4>, dense<[32, 64, 64, 24]> : vector<4xi64>>,
+  #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>,
+  #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>,
+  #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
+  #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>
+>} {
+  // CHECK: @spec
+  func.func @spec() {
+    // CHECK: alignment = 4
+    // CHECK: alloca_memory_space = 5
+    // CHECK: bitsize = 32
+    // CHECK: global_memory_space = 2
+    // CHECK: index = 32
+    // CHECK: preferred = 8
+    // CHECK: program_memory_space = 3
+    // CHECK: size = 4
+    // CHECK: stack_alignment = 128
+    "test.data_layout_query"() : () -> !ptr.ptr
+    // CHECK: alignment = 4
+    // CHECK: alloca_memory_space = 5
+    // CHECK: bitsize = 32
+    // CHECK: global_memory_space = 2
+    // CHECK: index = 32
+    // CHECK: preferred = 8
+    // CHECK: program_memory_space = 3
+    // CHECK: size = 4
+    // CHECK: stack_alignment = 128
+    "test.data_layout_query"() : () -> !ptr.ptr<3>
+    // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 5
+    // CHECK: bitsize = 64
+    // CHECK: global_memory_space = 2
+    // CHECK: index = 64
+    // CHECK: preferred = 8
+    // CHECK: program_memory_space = 3
+    // CHECK: size = 8
+    // CHECK: stack_alignment = 128
+    "test.data_layout_query"() : () -> !ptr.ptr<5>
+    // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 5
+    // CHECK: bitsize = 32
+    // CHECK: global_memory_space = 2
+    // CHECK: index = 24
+    // CHECK: preferred = 8
+    // CHECK: program_memory_space = 3
+    // CHECK: size = 4
+    // CHECK: stack_alignment = 128
+    "test.data_layout_query"() : () -> !ptr.ptr<4>
+    return
+  }
+}
+
+// -----
+
+// expected-error@below {{expected layout attribute for '!ptr.ptr' to be a dense integer elements attribute with 3 or 4 elements}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!ptr.ptr, dense<[64.0, 64.0, 64.0]> : vector<3xf32>>
+>} {
+  func.func @pointer() {
+    return
+  }
+}
+
+// -----
+
+// expected-error@below {{preferred alignment is expected to be at least as large as ABI alignment}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!ptr.ptr, dense<[64, 64, 32]> : vector<3xi64>>
+>} {
+  func.func @pointer() {
+    return
+  }
+}
+
+// -----
+
+// expected-error @below {{expected i64 parameters for '!ptr.ptr'}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!ptr.ptr, dense<[32, 32, 64]> : vector<3xi32>>
+>} {
+}
+