[mlir] forward data layout query to scoping op in absence of specification

[ftynse]

Summary

Even if the layout specification is missing from an op that supports it, the op
is still expected to provide meaningful responses to data layout queries.
Forward them to the op instead of directly calling the default implementation.

Depends On D98524

[ftynse]

rebase

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

static unsigned cachedLookup(Type t, DenseMap<Type, unsigned> &cache,
 
   auto result = cache.try_emplace(t, compute(t));
   return result.first->second;
 }
 
 unsigned mlir::DataLayout::getTypeSize(Type t) const {
   checkValid();
   return cachedLookup(t, sizes, [&](Type ty) {
-    if (originalLayout) {
-      DataLayoutEntryList list = originalLayout.getSpecForType(ty.getTypeID());
-      if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
-        return iface.getTypeSize(ty, *this, list);
-      return detail::getDefaultTypeSize(ty, *this, list);
-    }
-    return detail::getDefaultTypeSize(ty, *this, {});
+    DataLayoutEntryList list;
+    if (originalLayout)
+      list = originalLayout.getSpecForType(ty.getTypeID());
+    if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
   });
 }
 
 unsigned mlir::DataLayout::getTypeSizeInBits(Type t) const {
   checkValid();
   return cachedLookup(t, bitsizes, [&](Type ty) {
-    if (originalLayout) {
-      DataLayoutEntryList list = originalLayout.getSpecForType(ty.getTypeID());
-      if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
-        return iface.getTypeSizeInBits(ty, *this, list);
-      return detail::getDefaultTypeSizeInBits(ty, *this, list);
-    }
-    return detail::getDefaultTypeSizeInBits(ty, *this, {});
+    DataLayoutEntryList list;
+    if (originalLayout)
+      list = originalLayout.getSpecForType(ty.getTypeID());
+    if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
   });
 }
 
 unsigned mlir::DataLayout::getTypeABIAlignment(Type t) const {
   checkValid();
   return cachedLookup(t, abiAlignments, [&](Type ty) {
-    if (originalLayout) {
-      DataLayoutEntryList list = originalLayout.getSpecForType(ty.getTypeID());
-      if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
-        return iface.getTypeABIAlignment(ty, *this, list);
-      return detail::getDefaultABIAlignment(ty, *this, list);
-    }
-    return detail::getDefaultABIAlignment(ty, *this, {});
+    DataLayoutEntryList list;
+    if (originalLayout)
+      list = originalLayout.getSpecForType(ty.getTypeID());
+    if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
   });
 }
 
 unsigned mlir::DataLayout::getTypePreferredAlignment(Type t) const {
   checkValid();
   return cachedLookup(t, preferredAlignments, [&](Type ty) {
-    if (originalLayout) {
-      DataLayoutEntryList list = originalLayout.getSpecForType(ty.getTypeID());
-      if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
-        return iface.getTypePreferredAlignment(ty, *this, list);
-      return detail::getDefaultPreferredAlignment(ty, *this, list);
-    }
-    return detail::getDefaultPreferredAlignment(ty, *this, {});
+    DataLayoutEntryList list;
+    if (originalLayout)
+      list = originalLayout.getSpecForType(ty.getTypeID());
+    if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
   });
 }
 
 //===----------------------------------------------------------------------===//
 // DataLayoutSpecInterface
 //===----------------------------------------------------------------------===//
 
 void DataLayoutSpecInterface::bucketEntriesByType(

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

Type parseType(DialectAsmParser &parser) const override {
     return SingleQueryType::get(parser.getBuilder().getContext());
   }
 };
 
 } // end namespace
 
 TEST(DataLayout, FallbackDefault) {
   const char *ir = R"MLIR(
-"dltest.op_with_layout"() : () -> ()
+module {}
   )MLIR";
 
   DialectRegistry registry;
   registry.insert<DLTIDialect, DLTestDialect>();
   MLIRContext ctx(registry);
 
   OwningModuleRef module = parseSourceString(ir, &ctx);
-  auto op =
-      cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
-  DataLayout layout(op);
+  DataLayout layout(module.get());
   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u);
   EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 2u);
   EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 42u);
   EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 16u);
   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
   EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 2u);
   EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 8u);
   EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u);
 }
 
+TEST(DataLayout, NullSpec) {
+  const char *ir = R"MLIR(
+"dltest.op_with_layout"() : () -> ()
+  )MLIR";
+
+  DialectRegistry registry;
+  registry.insert<DLTIDialect, DLTestDialect>();
+  MLIRContext ctx(registry);
+
+  OwningModuleRef module = parseSourceString(ir, &ctx);
+  auto op =
+      cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
+  DataLayout layout(op);
+  EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u);
+  EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u);
+  EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u);
+  EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 8u * 16u);
+  EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u);
+  EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u);
+  EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u);
+  EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
+}
+
 TEST(DataLayout, EmptySpec) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> ()
   )MLIR";
 
   DialectRegistry registry;
   registry.insert<DLTIDialect, DLTestDialect>();
   MLIRContext ctx(registry);