[mlir] introduce data layout entry for index type

[ftynse]

Summary

Index type is an integer type of target-specific bitwidth present in many MLIR
operations (loops, memory accesses). Converting values of this type to
fixed-size integers has always been problematic. Introduce a data layout entry
to specify the bitwidth of index in a given layout scope, defaulting to 64
bits, which is a commonly used assumption, e.g., in constants.

Port builtin-to-LLVM type conversion to use this data layout entry when
converting index type and untie it from pointer size. This is particularly
relevant for GPU targets. Keep a possibility to forcibly override the index
type in lowerings.

Depends On D98525

[ftynse]

cmake

[herhut]

Generally looks good. Just one question. Thanks for adding support for this!

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

56

Now that this is data layout based, should it not query for the local context rather than assuming a global value? E.g., we could lower a module that contains gpu.module all at once and use different index sizes for them. Currently we have to lower in two steps.

[ftynse]

Rebase

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

56

This would be a much bigger change potentially affecting core infra: the type converter will have to know in which IR scope it is called, which could further affect things like materialization that currently assumes type conversion is scope-agnostic. It will also need a sort of analysis that maintains the DataLayout cache. I will eventually look into that, but this change enables the data layout-driven flow without breaking the current two-stage approach.

[herhut]

Thanks for cleaning this up.

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

56

Agreed. This makes sense longer term but the current solution does not regress what we currently have. As long as one only applies the conversion to one data layout scope at a time (currently likely a module) this is fine.

[mehdi_amini]

mlir/docs/DataLayout.md

273

I think this is leaving some confusion in what is the difference between an index and an integer of the provided size?
I.e. if you have a layout that is providing a size for the "index" type, then what is the difference with just using this integer everywhere?
If there are no semantics difference between the index and an integer in such case, it looks like it makes index just a "frontend" type, but I'm not sure why it just shouldn't be immediately "lowered" to an integer when the size is known instead of keeping metadata (otherwise you have two representation for the same thing (a sized integer) which complexify the IR and the compiler in general).

mlir/docs/DataLayout.md

 
 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.
 
+#### `index` type
+
+Index type is an integer type used for target-specific size information in,
+e.g., `memref` operations. Its data layout is parameterized by a single integer
+data layout entry that specifies its bitwidth. For example,
+
+```
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<index, 32>
+>} {}
+```
+
+specifies that `index` has 32 bits. All other layout properties of `index` match
+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.
+
 ### 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/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

public:
   /// Returns the data layout to use during and after conversion.
   const llvm::DataLayout &getDataLayout() { return options.dataLayout; }
 
   /// Gets the LLVM representation of the index type. The returned type is an
   /// integer type with the size configured for this type converter.
   Type getIndexType();
 
   /// Gets the bitwidth of the index type when converted to LLVM.
-  unsigned getIndexTypeBitwidth() { return options.indexBitwidth; }
+  unsigned getIndexTypeBitwidth() { return options.getIndexBitwidth(); }
 
   /// Gets the pointer bitwidth.
   unsigned getPointerBitwidth(unsigned addressSpace = 0);
 
 protected:
   /// Pointer to the LLVM dialect.
   LLVM::LLVMDialect *llvmDialect;
 

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

 //===- ConvertStandardToLLVMPass.h - Pass entrypoint ------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_
 #define MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_
 
 #include "llvm/IR/DataLayout.h"
 
 #include <memory>
 
 namespace mlir {
+class DataLayout;
 class LLVMTypeConverter;
+class MLIRContext;
 class ModuleOp;
 template <typename T>
 class OperationPass;
 class RewritePatternSet;
 using OwningRewritePatternList = RewritePatternSet;
 
 /// Value to pass as bitwidth for the index type when the converter is expected
 /// to derive the bitwidth from the LLVM data layout.
 static constexpr unsigned kDeriveIndexBitwidthFromDataLayout = 0;
 
 /// Options to control the Standard dialect to LLVM lowering. The struct is used
 /// to share lowering options between passes, patterns, and type converter.
-struct LowerToLLVMOptions {
+class LowerToLLVMOptions {
+public:
+  explicit LowerToLLVMOptions(MLIRContext *ctx);
+  explicit LowerToLLVMOptions(MLIRContext *ctx, DataLayout dl);
+
   bool useBarePtrCallConv = false;
   bool emitCWrappers = false;
-  unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout;
+
   /// Use aligned_alloc for heap allocations.
   bool useAlignedAlloc = false;
 
   /// The data layout of the module to produce. This must be consistent with the
   /// data layout used in the upper levels of the lowering pipeline.
   // TODO: this should be replaced by MLIR data layout when one exists.
   llvm::DataLayout dataLayout = llvm::DataLayout("");
 
-  /// Get a statically allocated copy of the default LowerToLLVMOptions.
-  static const LowerToLLVMOptions &getDefaultOptions() {
-    static LowerToLLVMOptions options;
-    return options;
+  /// Set the index bitwidth to the given value.
+  void overrideIndexBitwidth(unsigned bitwidth) {
+    assert(bitwidth != kDeriveIndexBitwidthFromDataLayout &&
+           "can only override to a concrete bitwidth");
+    indexBitwidth = bitwidth;
   }
+
+  /// Get the index bitwidth.
+  unsigned getIndexBitwidth() const { return indexBitwidth; }
+
+private:
+  unsigned indexBitwidth;
 };
 
 /// Collect a set of patterns to convert memory-related operations from the
 /// Standard dialect to the LLVM dialect, excluding non-memory-related
 /// operations and FuncOp.
 void populateStdToLLVMMemoryConversionPatterns(LLVMTypeConverter &converter,
                                                RewritePatternSet &patterns);
 

 /// pattern lifetime.
 void populateStdToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                          RewritePatternSet &patterns);
 
 /// Creates a pass to convert the Standard dialect into the LLVMIR dialect.
 /// stdlib malloc/free is used by default for allocating memrefs allocated with
 /// memref.alloc, while LLVM's alloca is used for those allocated with
 /// memref.alloca.
+std::unique_ptr<OperationPass<ModuleOp>> createLowerToLLVMPass();
 std::unique_ptr<OperationPass<ModuleOp>>
-createLowerToLLVMPass(const LowerToLLVMOptions &options =
-                          LowerToLLVMOptions::getDefaultOptions());
+createLowerToLLVMPass(const LowerToLLVMOptions &options);
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_

mlir/include/mlir/Dialect/GPU/GPUDialect.h

 //===- GPUDialect.h - MLIR Dialect for GPU Kernels --------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the GPU kernel-related operations and puts them in the
 // corresponding dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_GPU_GPUDIALECT_H
 #define MLIR_DIALECT_GPU_GPUDIALECT_H
 
+#include "mlir/Dialect/DLTI/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"

mlir/include/mlir/Dialect/GPU/GPUOps.td

 //===-- GPUOps.td - GPU dialect operation definitions ------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Defines some operations of the GPU dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef GPU_OPS
 #define GPU_OPS
 
+include "mlir/Dialect/DLTI/DLTIBase.td"
 include "mlir/Dialect/GPU/GPUBase.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/SymbolInterfaces.td"
+include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // GPU Dialect operations.
 //===----------------------------------------------------------------------===//
 
 class GPU_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<GPU_Dialect, mnemonic, traits>;
let description = [{
     Either none or all work items of a workgroup need to execute this op
     in convergence.
   }];
   let parser = [{ return success(); }];
   let printer = [{ p << getOperationName(); }];
 }
 
 def GPU_GPUModuleOp : GPU_Op<"module", [
-  IsolatedFromAbove, SymbolTable, Symbol,
+  DataLayoutOpInterface, HasDefaultDLTIDataLayout, IsolatedFromAbove,
+  SymbolTable, Symbol,
   SingleBlockImplicitTerminator<"ModuleEndOp">
 ]> {
   let summary = "A top level compilation unit containing code to be run on a GPU.";
   let description = [{
     GPU module contains code that is intended to be run on a GPU. A host device
     can launch this code through a gpu.launc_func that creates a fully
     qualified symbol through the gpu.module's symbol and a gpu.func symbol
     contained in the gpu.module.

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

 /// all its ancestors will be computed and used to handle further requests. For
 /// efficiency, results to all requests will be cached in this object.
 /// Therefore, if the data layout spec for the scoping operation, or any of the
 /// enclosing operations, changes, the cache is no longer valid. The user is
 /// responsible creating a new DataLayout object after any spec change. In debug
 /// mode, the cache validity is being checked in every request.
 class DataLayout {
 public:
+  explicit DataLayout();
   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;

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

struct LowerGpuOpsToNVVMOpsPass
   LowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) {
     this->indexBitwidth = indexBitwidth;
   }
 
   void runOnOperation() override {
     gpu::GPUModuleOp m = getOperation();
 
     /// Customize the bitwidth used for the device side index computations.
-    LowerToLLVMOptions options = {/*useBarePtrCallConv =*/false,
-                                  /*emitCWrappers =*/true,
-                                  /*indexBitwidth =*/indexBitwidth,
-                                  /*useAlignedAlloc =*/false};
+    LowerToLLVMOptions options(
+        m.getContext(),
+        DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));
+    options.emitCWrappers = true;
+    if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
+      options.overrideIndexBitwidth(indexBitwidth);
 
     /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory
     /// space 5 for private memory attributions, but NVVM represents private
     /// memory allocations as local `alloca`s in the default address space. This
     /// converter drops the private memory space to support the use case above.
     LLVMTypeConverter converter(m.getContext(), options);
     converter.addConversion([&](MemRefType type) -> Optional<Type> {
       if (type.getMemorySpaceAsInt() !=

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

struct LowerGpuOpsToROCDLOpsPass
   LowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth) {
     this->indexBitwidth = indexBitwidth;
   }
 
   void runOnOperation() override {
     gpu::GPUModuleOp m = getOperation();
 
     /// Customize the bitwidth used for the device side index computations.
-    LowerToLLVMOptions options = {/*useBarePtrCallConv =*/false,
-                                  /*emitCWrappers =*/true,
-                                  /*indexBitwidth =*/indexBitwidth,
-                                  /*useAlignedAlloc =*/false};
+    LowerToLLVMOptions options(
+        m.getContext(),
+        DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));
+    options.emitCWrappers = true;
+    if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
+      options.overrideIndexBitwidth(indexBitwidth);
     LLVMTypeConverter converter(m.getContext(), options);
 
     RewritePatternSet patterns(m.getContext());
     RewritePatternSet llvmPatterns(m.getContext());
 
     populateGpuRewritePatterns(patterns);
     (void)applyPatternsAndFoldGreedily(m, std::move(patterns));
 

mlir/lib/Conversion/SPIRVToLLVM/ConvertLaunchFuncToLLVMCalls.cpp

void runOnOperation() override {
 
     // Erase the GPU module.
     for (auto gpuModule :
          llvm::make_early_inc_range(module.getOps<gpu::GPUModuleOp>()))
       gpuModule.erase();
 
     // Specify options to lower Standard to LLVM and pull in the conversion
     // patterns.
-    LowerToLLVMOptions options = {
-        /*useBarePtrCallConv=*/false,
-        /*emitCWrappers=*/true,
-        /*indexBitwidth=*/kDeriveIndexBitwidthFromDataLayout};
+    LowerToLLVMOptions options(module.getContext());
+    options.emitCWrappers = true;
     auto *context = module.getContext();
     RewritePatternSet patterns(context);
     LLVMTypeConverter typeConverter(context, options);
     populateStdToLLVMConversionPatterns(typeConverter, patterns);
     patterns.add<GPULaunchLowering>(typeConverter);
 
     // Pull in SPIR-V type conversion patterns to convert SPIR-V global
     // variable's type to LLVM dialect type.

mlir/lib/Conversion/StandardToLLVM/CMakeLists.txt

 add_mlir_conversion_library(MLIRStandardToLLVM
   StandardToLLVM.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/StandardToLLVM
 
   DEPENDS
   MLIRConversionPassIncGen
   intrinsics_gen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
+  MLIRDataLayoutInterfaces
   MLIRLLVMIR
   MLIRMath
   MLIRMemRef
   MLIRTransforms
   )

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
if (!llvmTy)
     return failure();
 
   result.push_back(llvmTy);
   return success();
 }
 
 /// Create an LLVMTypeConverter using default LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx)
-    : LLVMTypeConverter(ctx, LowerToLLVMOptions::getDefaultOptions()) {}
+    : LLVMTypeConverter(ctx, LowerToLLVMOptions(ctx)) {}
 
 /// Create an LLVMTypeConverter using custom LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
                                      const LowerToLLVMOptions &options)
     : llvmDialect(ctx->getOrLoadDialect<LLVM::LLVMDialect>()),
       options(options) {
   assert(llvmDialect && "LLVM IR dialect is not registered");
-  if (options.indexBitwidth == kDeriveIndexBitwidthFromDataLayout)
-    this->options.indexBitwidth = options.dataLayout.getPointerSizeInBits();
 
   // Register conversions for the builtin types.
   addConversion([&](ComplexType type) { return convertComplexType(type); });
   addConversion([&](FloatType type) { return convertFloatType(type); });
   addConversion([&](FunctionType type) { return convertFunctionType(type); });
   addConversion([&](IndexType type) { return convertIndexType(type); });
   addConversion([&](IntegerType type) { return convertIntegerType(type); });
   addConversion([&](MemRefType type) { return convertMemRefType(type); });
if (failed(LLVM::LLVMDialect::verifyDataLayoutString(
               getOperation().emitError() << message.str();
             }))) {
       signalPassFailure();
       return;
     }
 
     ModuleOp m = getOperation();
 
-    LowerToLLVMOptions options = {useBarePtrCallConv, emitCWrappers,
-                                  indexBitwidth, useAlignedAlloc,
-                                  llvm::DataLayout(this->dataLayout)};
+    LowerToLLVMOptions options(&getContext(), DataLayout(m));
+    options.useBarePtrCallConv = useBarePtrCallConv;
+    options.emitCWrappers = emitCWrappers;
+    if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
+      options.overrideIndexBitwidth(indexBitwidth);
+    options.useAlignedAlloc = useAlignedAlloc;
+    options.dataLayout = llvm::DataLayout(this->dataLayout);
     LLVMTypeConverter typeConverter(&getContext(), options);
 
     RewritePatternSet patterns(&getContext());
     populateStdToLLVMConversionPatterns(typeConverter, patterns);
 
     LLVMConversionTarget target(getContext());
     if (failed(applyPartialConversion(m, target, std::move(patterns))))
       signalPassFailure();
     m->setAttr(LLVM::LLVMDialect::getDataLayoutAttrName(),
                StringAttr::get(m.getContext(), this->dataLayout));
   }
 };
 } // end namespace
 
 mlir::LLVMConversionTarget::LLVMConversionTarget(MLIRContext &ctx)
     : ConversionTarget(ctx) {
   this->addLegalDialect<LLVM::LLVMDialect>();
   this->addIllegalOp<LLVM::DialectCastOp>();
   this->addIllegalOp<math::TanhOp>();
 }
 
+std::unique_ptr<OperationPass<ModuleOp>> mlir::createLowerToLLVMPass() {
+  return std::make_unique<LLVMLoweringPass>();
+}
+
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createLowerToLLVMPass(const LowerToLLVMOptions &options) {
   return std::make_unique<LLVMLoweringPass>(
-      options.useBarePtrCallConv, options.emitCWrappers, options.indexBitwidth,
-      options.useAlignedAlloc, options.dataLayout);
+      options.useBarePtrCallConv, options.emitCWrappers,
+      options.getIndexBitwidth(), options.useAlignedAlloc, options.dataLayout);
+}
+
+mlir::LowerToLLVMOptions::LowerToLLVMOptions(MLIRContext *ctx)
+    : LowerToLLVMOptions(ctx, DataLayout()) {}
+
+mlir::LowerToLLVMOptions::LowerToLLVMOptions(MLIRContext *ctx,
+                                             mlir::DataLayout dl) {
+  indexBitwidth = dl.getTypeSizeInBits(IndexType::get(ctx));
 }

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPU
   MLIRGPUOpsIncGen
   MLIRGPUOpInterfacesIncGen
   MLIRGPUPassIncGen
   MLIRParallelLoopMapperAttrGen
   MLIRParallelLoopMapperEnumsGen
 
   LINK_LIBS PUBLIC
   MLIRAsync
+  MLIRDataLayoutInterfaces
+  MLIRDLTI
   MLIREDSC
   MLIRIR
   MLIRMemRef
   MLIRLLVMIR
   MLIRLLVMToLLVMIRTranslation
   MLIRSCF
   MLIRPass
   MLIRSideEffectInterfaces

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

static LLVM_ATTRIBUTE_NORETURN void reportMissingDataLayout(Type type) {
   std::string message;
   llvm::raw_string_ostream os(message);
   os << "neither the scoping op nor the type class provide data layout "
         "information for "
      << type;
   llvm::report_fatal_error(os.str());
 }
 
+/// Returns the bitwidth of the index type if specified in the param list.
+/// Assumes 64-bit index otherwise.
+static unsigned getIndexBitwidth(DataLayoutEntryListRef params) {
+  if (params.empty())
+    return 64;
+  auto attr = params.front().getValue().cast<IntegerAttr>();
+  return attr.getValue().getZExtValue();
+}
+
 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 type.getIntOrFloatBitWidth();
 
+  // Index is an integer of some bitwidth.
+  if (type.isa<IndexType>())
+    return dataLayout.getTypeSizeInBits(
+        IntegerType::get(type.getContext(), getIndexBitwidth(params)));
+
   // Sizes of vector types are rounded up to those of types with closest
   // 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 vecType.getNumElements() / vecType.getShape().back() *
            llvm::PowerOf2Ceil(vecType.getShape().back()) *
            dataLayout.getTypeSize(vecType.getElementType()) * 8;
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     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 (type.isa<FloatType, VectorType>())
     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
 
+  // Index is an integer of some bitwidth.
+  if (type.isa<IndexType>())
+    return dataLayout.getTypeABIAlignment(
+        IntegerType::get(type.getContext(), getIndexBitwidth(params)));
+
   if (auto intType = type.dyn_cast<IntegerType>()) {
     return intType.getWidth() < 64
                ? llvm::PowerOf2Ceil(llvm::divideCeil(intType.getWidth(), 8))
                : 4;
   }
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getABIAlignment(dataLayout, params);
 
   reportMissingDataLayout(type);
 }
 
 unsigned mlir::detail::getDefaultPreferredAlignment(
     Type type, const DataLayout &dataLayout,
     ArrayRef<DataLayoutEntryInterface> params) {
   // Preferred alignment is same as natural for floats and vectors.
   if (type.isa<FloatType, VectorType>())
     return dataLayout.getTypeABIAlignment(type);
 
   // Preferred alignment is the cloest power-of-two number above for integers
   // (ABI alignment may be smaller).
-  if (auto intType = type.dyn_cast<IntegerType>())
+  if (type.isa<IntegerType, IndexType>())
     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getPreferredAlignment(dataLayout, params);
 
   reportMissingDataLayout(type);
 }
 

 void checkMissingLayout(DataLayoutSpecInterface originalLayout, OpTy op) {
   if (!originalLayout) {
     assert((!op || !op.getDataLayoutSpec()) &&
            "could not compute layout information for an op (failed to "
            "combine attributes?)");
   }
 }
 
+mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
+
 mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
     : originalLayout(getCombinedDataLayout(op)), scope(op) {
 #ifndef NDEBUG
   checkMissingLayout(originalLayout, op);
   collectParentLayouts(op, layoutStack);
 #endif
 }
 
LogicalResult mlir::detail::verifyDataLayoutSpec(DataLayoutSpecInterface spec,
   // Second, dispatch verifications of entry groups to types or dialects they
   // are are associated with.
   DenseMap<TypeID, DataLayoutEntryList> types;
   DenseMap<Identifier, DataLayoutEntryInterface> ids;
   spec.bucketEntriesByType(types, ids);
 
   for (const auto &kvp : types) {
     auto sampleType = kvp.second.front().getKey().get<Type>();
+    if (sampleType.isa<IndexType>()) {
+      assert(kvp.second.size() == 1 &&
+             "expected one data layout entry for non-parametric 'index' type");
+      if (!kvp.second.front().getValue().isa<IntegerAttr>())
+        return emitError(loc)
+               << "expected integer attribute in the data layout entry for "
+               << sampleType;
+      continue;
+    }
+
     if (isa<BuiltinDialect>(&sampleType.getDialect()))
       return emitError(loc) << "unexpected data layout for a built-in type";
 
     auto dlType = sampleType.dyn_cast<DataLayoutTypeInterface>();
     if (!dlType)
       return emitError(loc)
              << "data layout specified for a type that does not support it";
     if (failed(dlType.verifyEntries(kvp.second, loc)))

mlir/test/Interfaces/DataLayoutInterfaces/types.mlir

+// RUN: mlir-opt --split-input-file --verify-diagnostics --test-data-layout-query %s | FileCheck %s
+
+// expected-error@below {{expected integer attribute in the data layout entry for 'index'}}
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<index, [32]>>} {
+}
+
+// -----
+
+// CHECK-LABEL: @index
+module @index attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<index, 32>>} {
+  func @query() {
+    // CHECK: bitsize = 32
+    "test.data_layout_query"() : () -> index
+    return
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @index_default
+module @index_default {
+  func @query() {
+    // CHECK: bitsize = 64
+    "test.data_layout_query"() : () -> index
+    return
+  }
+}

mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp

static LogicalResult runMLIRPasses(ModuleOp module) {
 
   passManager.addPass(createGpuKernelOutliningPass());
   passManager.addPass(createLegalizeStdOpsForSPIRVLoweringPass());
   passManager.addPass(createConvertGPUToSPIRVPass());
   OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
   modulePM.addPass(spirv::createLowerABIAttributesPass());
   modulePM.addPass(spirv::createUpdateVersionCapabilityExtensionPass());
   passManager.addPass(createConvertGpuLaunchFuncToVulkanLaunchFuncPass());
-  LowerToLLVMOptions llvmOptions = {
-      /*useBarePtrCallConv =*/false,
-      /*emitCWrappers = */ true,
-      /*indexBitwidth =*/kDeriveIndexBitwidthFromDataLayout};
+  LowerToLLVMOptions llvmOptions(module.getContext(), DataLayout(module));
+  llvmOptions.emitCWrappers = true;
   passManager.addPass(createLowerToLLVMPass(llvmOptions));
   passManager.addPass(createConvertVulkanLaunchFuncToVulkanCallsPass());
   return passManager.run(module);
 }
 
 int main(int argc, char **argv) {
   llvm::llvm_shutdown_obj x;
   registerPassManagerCLOptions();