[mlir] provide a version of data layout size hooks in bits

[ftynse]

Summary

This is useful for bit-packing types such as vectors and tuples as well as for
exotic architectures that have non-8-bit bytes.

Depends On D98500

[ftynse]

rebase

mlir/docs/DataLayout.md

 `DataLayout` class is as follows.
 
 ```c++
 class DataLayout {
 public:
   explicit DataLayout(DataLayoutOpInterface scope);
 
   unsigned getTypeSize(Type type) const;
+  unsigned getTypeSizeInBits(Type type) const;
   unsigned getTypeABIAlignment(Type type) const;
   unsigned getTypePreferredAlignment(Type type) const;
 };
 ```
 
 The user can construct the `DataLayout` object for the scope of interest. Since
 the data layout properties are fixed in the scope, they will be computed only
 once upon first request and cached for further use. Therefore,

 ### Dialects with Data Layout Identifiers
 
 For data layout entries that are not related to a particular type class, the key
 of the entry is an Identifier that belongs to some dialect. In this case, the
 dialect is expected to implement the `DataLayoutDialectInterface`. This dialect
 provides hooks for verifying the validity of the entry value attributes and for
 and the compatibility of nested entries.
 
+### Bits and Bytes
+
+Two versions of hooks are provided for sizes: in bits and in bytes. The version
+in bytes has a default implementation that derives the size in bytes by rounding
+up the result of division of the size in bits by 8. Types exclusively targeting
+architectures with different assumptions can override this. Operations can
+redefine this for all types, providing scoped versions for cases of byte sizes
+other than eight without having to modify types, including built-in types.
+
 ### Query Dispatch
 
 The overall flow of a data layout property query is as follows.
 
 1.  The user constructs a `DataLayout` at the given scope. The constructor
     fetches the data layout specification and combines it with those of
     enclosing scopes (layouts are expected to be compatible).
 2.  The user calls `DataLayout::query(Type ty)`.

 
 Note: these values are selected for consistency with the
 [default data layout in LLVM](https://llvm.org/docs/LangRef.html#data-layout),
 which MLIR assumed until the introduction of proper data layout modeling, and
 with the
 [modeling of n-D vectors](https://mlir.llvm.org/docs/Dialects/Vector/#deeperdive).
 They **may change** in the future.
 
+### Byte Size
+
+The default data layout assumes 8-bit bytes.
+
 ### DLTI Dialect
 
 The [DLTI](Dialects/DLTI.md) dialect provides the attributes implementing
 `DataLayoutSpecInterface` and `DataLayoutEntryInterface`, as well as a dialect
 attribute that can be used to attach the specification to a given operation. The
 verifier of this attribute triggers those of the specification and checks the
 compatiblity of nested specifications.

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

 class DataLayoutOpInterface;
 class DataLayoutSpecInterface;
 class ModuleOp;
 
 namespace detail {
 /// Default handler for the type size request. Computes results for built-in
 /// types and dispatches to the DataLayoutTypeInterface for other types.
 unsigned getDefaultTypeSize(Type type, const DataLayout &dataLayout,
-                            ArrayRef<DataLayoutEntryInterface> params);
+                            DataLayoutEntryListRef params);
+
+/// Default handler for the type size in bits request. Computes results for
+/// built-in types and dispatches to the DataLayoutTypeInterface for other
+/// types.
+unsigned getDefaultTypeSizeInBits(Type type, const DataLayout &dataLayout,
+                                  DataLayoutEntryListRef params);
 
 /// Default handler for the required alignemnt request. Computes results for
 /// built-in types and dispatches to the DataLayoutTypeInterface for other
 /// types.
 unsigned getDefaultABIAlignment(Type type, const DataLayout &dataLayout,
                                 ArrayRef<DataLayoutEntryInterface> params);
 
 /// Default handler for the preferred alignemnt request. Computes results for

 class DataLayout {
 public:
   explicit DataLayout(DataLayoutOpInterface op);
   explicit DataLayout(ModuleOp op);
 
   /// Returns the size of the given type in the current scope.
   unsigned getTypeSize(Type t) const;
 
+  /// Returns the size in bits of the given type in the current scope.
+  unsigned getTypeSizeInBits(Type t) const;
+
   /// Returns the required alignment of the given type in the current scope.
   unsigned getTypeABIAlignment(Type t) const;
 
   /// Returns the preferred of the given type in the current scope.
   unsigned getTypePreferredAlignment(Type t) const;
 
 private:
   /// Combined layout spec at the given scope.
#endif
 
   /// Operation defining the scope of requests.
   // TODO: this is mutable because the generated interface method are not const.
   // Update the generator to support const methods and change this to const.
   mutable Operation *scope;
 
   /// Caches for individual requests.
   mutable DenseMap<Type, unsigned> sizes;
+  mutable DenseMap<Type, unsigned> bitsizes;
   mutable DenseMap<Type, unsigned> abiAlignments;
   mutable DenseMap<Type, unsigned> preferredAlignments;
 };
 
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_DATALAYOUTINTERFACES_H

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

StaticInterfaceMethod<
                       "for recursive queries.",
       /*retTy=*/"unsigned",
       /*methodName=*/"getTypeSize",
       /*args=*/(ins "::mlir::Type":$type,
                     "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return ::mlir::detail::getDefaultTypeSize(type, dataLayout, params);
+        unsigned bits = ConcreteOp::getTypeSizeInBits(type, dataLayout, params);
+        return ::llvm::divideCeil(bits, 8);
+      }]
+    >,
+    StaticInterfaceMethod<
+      /*description=*/"Returns the size of the given type in bits computed "
+                      "using the relevant entries. The data layout object can "
+                      "be used for recursive queries.",
+      /*retTy=*/"unsigned",
+      /*methodName=*/"getTypeSizeInBits",
+      /*args=*/(ins "::mlir::Type":$type,
+                    "const ::mlir::DataLayout &":$dataLayout,
+                    "::mlir::DataLayoutEntryListRef":$params),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultTypeSizeInBits(type, dataLayout,
+                                                        params);
       }]
     >,
     StaticInterfaceMethod<
       /*description=*/"Returns the alignment required by the ABI for the given "
                       "type computed using the relevant entries. The data "
                       "layout object can be used for recursive queries.",
       /*retTy=*/"unsigned",
       /*methodName=*/"getTypeABIAlignment",
def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
   }];
 
   let methods = [
     InterfaceMethod<
       /*description=*/"Returns the size of this type in bytes.",
       /*retTy=*/"unsigned",
       /*methodName=*/"getTypeSize",
       /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
+                    "::mlir::DataLayoutEntryListRef":$params),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        unsigned bits = $_type.getTypeSizeInBits(dataLayout, params);
+        return ::llvm::divideCeil(bits, 8);
+      }]
+    >,
+    InterfaceMethod<
+      /*description=*/"Returns the size of this type in bits.",
+      /*retTy=*/"unsigned",
+      /*methodName=*/"getTypeSizeInBits",
+      /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params)
     >,
     InterfaceMethod<
       /*description=*/"Returns the ABI-required alignment for this type, "
                       "in bytes",
       /*retTy=*/"unsigned",
       /*methodName=*/"getABIAlignment",
       /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

os << "neither the scoping op nor the type class provide data layout "
         "information for "
      << type;
   llvm::report_fatal_error(os.str());
 }
 
 unsigned
 mlir::detail::getDefaultTypeSize(Type type, const DataLayout &dataLayout,
                                  ArrayRef<DataLayoutEntryInterface> params) {
+  unsigned bits = getDefaultTypeSizeInBits(type, dataLayout, params);
+  return llvm::divideCeil(bits, 8);
+}
+
+unsigned mlir::detail::getDefaultTypeSizeInBits(Type type,
+                                                const DataLayout &dataLayout,
+                                                DataLayoutEntryListRef params) {
   if (type.isa<IntegerType, FloatType>())
-    return llvm::divideCeil(type.getIntOrFloatBitWidth(), 8);
+    return type.getIntOrFloatBitWidth();
 
   // Sizes of vector types are rounded up to those of types with closest
-  // power-of-two number of elements.
+  // power-of-two number of elements in the innermost dimension. We also assume
+  // there is no bit-packing at the moment element sizes are taken in bytes and
+  // multiplied with 8 bits.
   // TODO: make this extensible.
   if (auto vecType = type.dyn_cast<VectorType>())
-    return llvm::PowerOf2Ceil(vecType.getNumElements()) *
-           dataLayout.getTypeSize(vecType.getElementType());
+    return vecType.getNumElements() / vecType.getShape().back() *
+           llvm::PowerOf2Ceil(vecType.getShape().back()) *
+           dataLayout.getTypeSize(vecType.getElementType()) * 8;
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
-    return typeInterface.getTypeSize(dataLayout, params);
+    return typeInterface.getTypeSizeInBits(dataLayout, params);
 
   reportMissingDataLayout(type);
 }
 
 unsigned mlir::detail::getDefaultABIAlignment(
     Type type, const DataLayout &dataLayout,
     ArrayRef<DataLayoutEntryInterface> params) {
   // Natural alignment is the closest power-of-two number above.
if (originalLayout) {
       if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
         return iface.getTypeSize(ty, *this, list);
       return detail::getDefaultTypeSize(ty, *this, list);
     }
     return detail::getDefaultTypeSize(ty, *this, {});
   });
 }
 
+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, {});
+  });
+}
+
 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);

mlir/test/Interfaces/DataLayoutInterfaces/module.mlir

 // RUN: mlir-opt --test-data-layout-query %s | FileCheck %s
 
 module attributes { dlti.dl_spec = #dlti.dl_spec<
       #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 12]>,
       #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 32]>>} {
   // CHECK-LABEL: @module_level_layout
   func @module_level_layout() {
      // CHECK: alignment = 32
+     // CHECK: bitsize = 12
      // CHECK: preferred = 1
-     // CHECK: size = 12
+     // CHECK: size = 2
     "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
     return
   }
 }

mlir/test/Interfaces/DataLayoutInterfaces/query.mlir

 // RUN: mlir-opt --test-data-layout-query %s | FileCheck %s
 
 // CHECK-LABEL: @no_layout_builtin
 func @no_layout_builtin() {
   // CHECK: alignment = 4
+  // CHECK: bitsize = 32
   // CHECK: preferred = 4
   // CHECK: size = 4
   "test.data_layout_query"() : () -> i32
   // CHECK: alignment = 8
+  // CHECK: bitsize = 64
   // CHECK: preferred = 8
   // CHECK: size = 8
   "test.data_layout_query"() : () -> f64
   return
 }
 
 // CHECK-LABEL: @no_layout_custom
 func @no_layout_custom() {
   // CHECK: alignment = 1
+  // CHECK: bitsize = 1
   // CHECK: preferred = 1
   // CHECK: size = 1
   "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
   return
 }
 
 // CHECK-LABEL: @layout_op_no_layout
 func @layout_op_no_layout() {
   "test.op_with_data_layout"() ({
     // CHECK: alignment = 1
+    // CHECK: bitsize = 1
     // CHECK: preferred = 1
     // CHECK: size = 1
     "test.data_layout_query"() : () -> !test.test_type_with_layout<1000>
     "test.maybe_terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
 // CHECK-LABEL: @layout_op
 func @layout_op() {
   "test.op_with_data_layout"() ({
     // CHECK: alignment = 20
+    // CHECK: bitsize = 10
     // CHECK: preferred = 1
-    // CHECK: size = 10
+    // CHECK: size = 2
     "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
     "test.maybe_terminator"() : () -> ()
   }) { dlti.dl_spec = #dlti.dl_spec<
       #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 10]>,
       #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 20]>
   >} : () -> ()
   return
 }
 
 // Make sure the outer op with layout may be missing the spec.
 // CHECK-LABEL: @nested_inner_only
 func @nested_inner_only() {
   "test.op_with_data_layout"() ({
     "test.op_with_data_layout"() ({
       // CHECK: alignment = 20
+      // CHECK: bitsize = 10
       // CHECK: preferred = 1
-      // CHECK: size = 10
+      // CHECK: size = 2
       "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
       "test.maybe_terminator"() : () -> ()
     }) { dlti.dl_spec = #dlti.dl_spec<
         #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 10]>,
         #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 20]>
     >} : () -> ()
     "test.maybe_terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
 // Make sure the inner op with layout may be missing the spec.
 // CHECK-LABEL: @nested_outer_only
 func @nested_outer_only() {
   "test.op_with_data_layout"() ({
     "test.op_with_data_layout"() ({
       // CHECK: alignment = 20
+      // CHECK: bitsize = 10
       // CHECK: preferred = 1
-      // CHECK: size = 10
+      // CHECK: size = 2
       "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
       "test.maybe_terminator"() : () -> ()
     }) : () -> ()
     "test.maybe_terminator"() : () -> ()
   }) { dlti.dl_spec = #dlti.dl_spec<
       #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 10]>,
       #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 20]>
     >} : () -> ()
   return
 }
 
 // CHECK-LABEL: @nested_middle_only
 func @nested_middle_only() {
   "test.op_with_data_layout"() ({
     "test.op_with_data_layout"() ({
       "test.op_with_data_layout"() ({
         // CHECK: alignment = 20
+        // CHECK: bitsize = 10
         // CHECK: preferred = 1
-        // CHECK: size = 10
+        // CHECK: size = 2
         "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
         "test.maybe_terminator"() : () -> ()
     }) : () -> ()
     "test.maybe_terminator"() : () -> ()
     }) { dlti.dl_spec = #dlti.dl_spec<
         #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 10]>,
         #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 20]>
       >} : () -> ()
     "test.maybe_terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
 // CHECK-LABEL: @nested_combine_with_missing
 func @nested_combine_with_missing() {
   "test.op_with_data_layout"() ({
     "test.op_with_data_layout"() ({
       "test.op_with_data_layout"() ({
         // CHECK: alignment = 20
+        // CHECK: bitsize = 10
         // CHECK: preferred = 30
-        // CHECK: size = 10
+        // CHECK: size = 2
         "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
         "test.maybe_terminator"() : () -> ()
       }) : () -> ()
     "test.maybe_terminator"() : () -> ()
     }) { dlti.dl_spec = #dlti.dl_spec<
         #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 10]>,
         #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 20]>
       >} : () -> ()
     // CHECK: alignment = 1
+    // CHECK: bitsize = 42
     // CHECK: preferred = 30
-    // CHECK: size = 42
+    // CHECK: size = 6
     "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
     "test.maybe_terminator"() : () -> ()
   }) { dlti.dl_spec = #dlti.dl_spec<
       #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 42]>,
       #dlti.dl_entry<!test.test_type_with_layout<30>, ["preferred", 30]>
   >}: () -> ()
   return
 }
 
 // CHECK-LABEL: @nested_combine_all
 func @nested_combine_all() {
   "test.op_with_data_layout"() ({
     "test.op_with_data_layout"() ({
       "test.op_with_data_layout"() ({
         // CHECK: alignment = 20
+        // CHECK: bitsize = 3
         // CHECK: preferred = 30
-        // CHECK: size = 3
+        // CHECK: size = 1
         "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
         "test.maybe_terminator"() : () -> ()
       }) { dlti.dl_spec = #dlti.dl_spec<
           #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 3]>,
           #dlti.dl_entry<!test.test_type_with_layout<30>, ["preferred", 30]>
         >} : () -> ()
       // CHECK: alignment = 20
+      // CHECK: bitsize = 10
       // CHECK: preferred = 30
-      // CHECK: size = 10
+      // CHECK: size = 2
       "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
       "test.maybe_terminator"() : () -> ()
     }) { dlti.dl_spec = #dlti.dl_spec<
         #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 10]>,
         #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 20]>
       >} : () -> ()
     // CHECK: alignment = 1
+    // CHECK: bitsize = 42
     // CHECK: preferred = 30
-    // CHECK: size = 42
+    // CHECK: size = 6
     "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
     "test.maybe_terminator"() : () -> ()
   }) { dlti.dl_spec = #dlti.dl_spec<
       #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 42]>,
       #dlti.dl_entry<!test.test_type_with_layout<30>, ["preferred", 30]>
   >}: () -> ()
   return
 }

mlir/test/lib/Dialect/Test/TestTypes.h

public:
   using Base::Base;
 
   static TestTypeWithLayout get(MLIRContext *ctx, unsigned key) {
     return Base::get(ctx, key);
   }
 
   unsigned getKey() { return getImpl()->key; }
 
-  unsigned getTypeSize(const DataLayout &dataLayout,
+  unsigned getTypeSizeInBits(const DataLayout &dataLayout,
                              DataLayoutEntryListRef params) const {
     return extractKind(params, "size");
   }
 
   unsigned getABIAlignment(const DataLayout &dataLayout,
                            DataLayoutEntryListRef params) const {
     return extractKind(params, "alignment");
   }
 

mlir/test/lib/Transforms/TestDataLayoutQuery.cpp

func.walk([&](test::DataLayoutQueryOp op) {
         layouts.try_emplace(module, module);
 
       Operation *closest = (scope && module && module->isProperAncestor(scope))
                                ? scope.getOperation()
                                : module.getOperation();
 
       const DataLayout &layout = layouts.find(closest)->getSecond();
       unsigned size = layout.getTypeSize(op.getType());
+      unsigned bitsize = layout.getTypeSizeInBits(op.getType());
       unsigned alignment = layout.getTypeABIAlignment(op.getType());
       unsigned preferred = layout.getTypePreferredAlignment(op.getType());
       op->setAttrs(
           {builder.getNamedAttr("size", builder.getIndexAttr(size)),
+           builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)),
            builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
            builder.getNamedAttr("preferred", builder.getIndexAttr(preferred))});
     });
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestDataLayoutQuery() {
   PassRegistration<TestDataLayoutQuery>("test-data-layout-query",
                                         "Test data layout queries");
 }
 } // namespace test
 } // namespace mlir

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

 /// than once. Handy to check if the cache works.
 struct SingleQueryType
     : public Type::TypeBase<SingleQueryType, Type, TypeStorage,
                             DataLayoutTypeInterface::Trait> {
   using Base::Base;
 
   static SingleQueryType get(MLIRContext *ctx) { return Base::get(ctx); }
 
-  unsigned getTypeSize(const DataLayout &layout,
-                       DataLayoutEntryListRef params) {
+  unsigned getTypeSizeInBits(const DataLayout &layout,
+                             DataLayoutEntryListRef params) const {
     static bool executed = false;
     if (executed)
       llvm::report_fatal_error("repeated call");
 
     executed = true;
     return 1;
   }
 
struct OpWithLayout : public Op<OpWithLayout, DataLayoutOpInterface::Trait> {
   using Op::Op;
 
   static StringRef getOperationName() { return "dltest.op_with_layout"; }
 
   DataLayoutSpecInterface getDataLayoutSpec() {
     return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName);
   }
 
-  static unsigned getTypeSize(Type type, const DataLayout &dataLayout,
+  static unsigned getTypeSizeInBits(Type type, const DataLayout &dataLayout,
                                     DataLayoutEntryListRef params) {
     // Make a recursive query.
     if (type.isa<FloatType>())
-      return dataLayout.getTypeSize(
+      return dataLayout.getTypeSizeInBits(
           IntegerType::get(type.getContext(), type.getIntOrFloatBitWidth()));
 
     // Handle built-in types that are not handled by the default process.
     if (auto iType = type.dyn_cast<IntegerType>()) {
       for (DataLayoutEntryInterface entry : params)
         if (entry.getKey().dyn_cast<Type>() == type)
-          return entry.getValue().cast<IntegerAttr>().getValue().getZExtValue();
-      return iType.getIntOrFloatBitWidth();
+          return 8 *
+                 entry.getValue().cast<IntegerAttr>().getValue().getZExtValue();
+      return 8 * iType.getIntOrFloatBitWidth();
     }
 
     // Use the default process for everything else.
     return detail::getDefaultTypeSize(type, dataLayout, params);
   }
 
   static unsigned getTypeABIAlignment(Type type, const DataLayout &dataLayout,
                                       DataLayoutEntryListRef params) {
     return llvm::PowerOf2Ceil(getTypeSize(type, dataLayout, params));
   }
 
   static unsigned getTypePreferredAlignment(Type type,
                                             const DataLayout &dataLayout,
                                             DataLayoutEntryListRef params) {
     return 2 * getTypeABIAlignment(type, dataLayout, params);
   }
 };
 
+struct OpWith7BitByte
+    : public Op<OpWith7BitByte, DataLayoutOpInterface::Trait> {
+  using Op::Op;
+
+  static StringRef getOperationName() { return "dltest.op_with_7bit_byte"; }
+
+  DataLayoutSpecInterface getDataLayoutSpec() {
+    return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName);
+  }
+
+  // Bytes are assumed to be 7-bit here.
+  static unsigned getTypeSize(Type type, const DataLayout &dataLayout,
+                              DataLayoutEntryListRef params) {
+    return llvm::divideCeil(dataLayout.getTypeSizeInBits(type), 7);
+  }
+};
+
 /// A dialect putting all the above together.
 struct DLTestDialect : Dialect {
   explicit DLTestDialect(MLIRContext *ctx)
       : Dialect(getDialectNamespace(), ctx, TypeID::get<DLTestDialect>()) {
     ctx->getOrLoadDialect<DLTIDialect>();
     addAttributes<CustomDataLayoutSpec>();
-    addOperations<OpWithLayout>();
+    addOperations<OpWithLayout, OpWith7BitByte>();
     addTypes<SingleQueryType, TypeNoLayout>();
   }
   static StringRef getDialectNamespace() { return "dltest"; }
 
   void printAttribute(Attribute attr,
                       DialectAsmPrinter &printer) const override {
     printer << "spec<";
     llvm::interleaveComma(attr.cast<CustomDataLayoutSpec>().getEntries(),
"dltest.op_with_layout"() : () -> ()
   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)), 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, EmptySpec) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> ()
   )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, SpecWithEntries) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<
   #dlti.dl_entry<i42, 5>,
   #dlti.dl_entry<i16, 6>
 > } : () -> ()
   )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)), 5u);
   EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u);
+  EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 40u);
+  EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 48u);
   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
   EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 8u);
   EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 16u);
   EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 16u);
 
   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 32u);
   EXPECT_EQ(layout.getTypeSize(Float32Type::get(&ctx)), 32u);
+  EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 32)), 256u);
+  EXPECT_EQ(layout.getTypeSizeInBits(Float32Type::get(&ctx)), 256u);
   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 32)), 32u);
   EXPECT_EQ(layout.getTypeABIAlignment(Float32Type::get(&ctx)), 32u);
   EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 32)), 64u);
   EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u);
 }
 
 TEST(DataLayout, Caching) {
   const char *ir = R"MLIR(
"dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> ()
   auto op =
       cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
   DataLayout layout(op);
 
   ASSERT_DEATH(layout.getTypeSize(TypeNoLayout::get(&ctx)),
                "neither the scoping op nor the type class provide data layout "
                "information");
 }
+
+TEST(DataLayout, SevenBitByte) {
+  const char *ir = R"MLIR(
+"dltest.op_with_7bit_byte"() { dltest.layout = #dltest.spec<> } : () -> ()
+  )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.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 42u);
+  EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 32)), 32u);
+  EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u);
+  EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 5u);
+}