[mlir][spirv] Initial support for 64 bit index type and builtins

[Hardcode84]

fix gcc build

[Hardcode84]

fix gcc build

[tschuett]

mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h

59

[herhut]

Could this be driven by a data-layout on the gpu module?

[antiagainst]

In D108516#2959848, @herhut wrote:

Could this be driven by a data-layout on the gpu module?

Indeed. I'd like to adopt DataLayout (https://mlir.llvm.org/docs/DataLayout/) but haven't been able to find time for it. If @Hardcode84 you are interested, that'd be a very welcome contribution! But don't want to require refactoring the whole world to introduce a new functionality. :) So I'm fine of landing this to enable your use case for now too. Just a couple of nits.

mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h

59

Nit: move these two booleans together? It helps packing a bit.

59

/// Use 64-bit integers to convert index types.

84

Is there a need to expose this as a public API?

140

Can we still leave the builder as the last parameter for all these functions?

Also here it's a utility function that has nothing to do with index. Just name the parameter as integerType?

142–143

32-bit needs to be changed here.

148–149

Same here: integerType.

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

788

We can get the index type from the typeConverter passed in?

[Hardcode84]

review comments

[herhut]

Basing this on the DataLayout is relatively lightweight. Take a look at how this is done in the LLVM lowering. Unless you had something more in mind.

[antiagainst]

In D108516#2960758, @herhut wrote:

Basing this on the DataLayout is relatively lightweight. Take a look at how this is done in the LLVM lowering. Unless you had something more in mind.

Yup. I'd prefer to do the switch for all, including the existing boolNumBits/etc. and how bitwidths for SPIR-V types are queried. Don't really like to be in a hybrid state where some datalayout configurations are coming from the IR while some are not. That can be quite confusing and error prone.

mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h

struct Options {
     /// runtime to also feed in data with a matched bitwidth and layout for
     /// interface types. The runtime can do that by inspecting the SPIR-V
     /// module.
     ///
     /// If the original scalar type has less than 32-bit, a multiple of its
     /// values will be packed into one 32-bit value to be memory efficient.
     bool emulateNon32BitScalarTypes;
 
+    /// Use 64-bit integers to convert index types.
+    bool use64bitIndex;
+
     /// The number of bits to store a boolean value. It is eight bits by
     /// default.
     unsigned boolNumBits;
 
-    // Note: we need this instead of inline initializers becuase of
+    // Note: we need this instead of inline initializers because of
     // https://bugs.llvm.org/show_bug.cgi?id=36684
-    Options() : emulateNon32BitScalarTypes(true), boolNumBits(8) {}
+    Options()
+        : emulateNon32BitScalarTypes(true), use64bitIndex(false),
+          boolNumBits(8) {}
   };
 
   explicit SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
                               Options options = {});
 
   /// Gets the SPIR-V correspondence for the standard index type.
-  static Type getIndexType(MLIRContext *context);
+  Type getIndexType() const;
 
   /// Returns the corresponding memory space for memref given a SPIR-V storage
   /// class.
   static unsigned getMemorySpaceForStorageClass(spirv::StorageClass);
 
   /// Returns the SPIR-V storage class given a memory space for memref. Return
   /// llvm::None if the memory space does not map to any SPIR-V storage class.
   static Optional<spirv::StorageClass>
   getStorageClassForMemorySpace(unsigned space);
 
   /// Returns the options controlling the SPIR-V type converter.
   const Options &getOptions() const;
 
 private:
   spirv::TargetEnv targetEnv;
   Options options;
+
+  MLIRContext *getContext() const;
 };
 
 //===----------------------------------------------------------------------===//
 // Conversion Target
 //===----------------------------------------------------------------------===//
 
 // The default SPIR-V conversion target.
 //
void populateBuiltinFuncToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                         RewritePatternSet &patterns);
 
 namespace spirv {
 class AccessChainOp;
 
 /// Returns the value for the given `builtin` variable. This function gets or
 /// inserts the global variable associated for the builtin within the nearest
 /// symbol table enclosing `op`. Returns null Value on error.
-Value getBuiltinVariableValue(Operation *op, BuiltIn builtin,
+Value getBuiltinVariableValue(Operation *op, BuiltIn builtin, Type integerType,
                               OpBuilder &builder);
 
 /// Gets the value at the given `offset` of the push constant storage with a
-/// total of `elementCount` 32-bit integers. A global variable will be created
-/// in the nearest symbol table enclosing `op` for the push constant storage if
-/// not existing. Load ops will be created via the given `builder` to load
-/// values from the push constant. Returns null Value on error.
+/// total of `elementCount` `integerType` integers. A global variable will be
+/// created in the nearest symbol table enclosing `op` for the push constant
+/// storage if not existing. Load ops will be created via the given `builder` to
+/// load values from the push constant. Returns null Value on error.
 Value getPushConstantValue(Operation *op, unsigned elementCount,
-                           unsigned offset, OpBuilder &builder);
+                           unsigned offset, Type integerType,
+                           OpBuilder &builder);
 
 /// Generates IR to perform index linearization with the given `indices` and
 /// their corresponding `strides`, adding an initial `offset`.
 Value linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
-                     int64_t offset, Location loc, OpBuilder &builder);
+                     int64_t offset, Type integerType, Location loc,
+                     OpBuilder &builder);
 
 /// Performs the index computation to get to the element at `indices` of the
 /// memory pointed to by `basePtr`, using the layout map of `baseType`.
 /// Returns null if index computation cannot be performed.
 
 // TODO: This method assumes that the `baseType` is a MemRefType with AffineMap
 // that has static strides. Extend to handle dynamic strides.
 spirv::AccessChainOp getElementPtr(SPIRVTypeConverter &typeConverter,
                                    MemRefType baseType, Value basePtr,
                                    ValueRange indices, Location loc,
                                    OpBuilder &builder);
 
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SPIRVCONVERSION_H

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

 template <typename SourceOp, spirv::BuiltIn builtin>
 LogicalResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
     SourceOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   auto index = getLaunchConfigIndex(op);
   if (!index)
     return failure();
 
+  auto *typeConverter = this->template getTypeConverter<SPIRVTypeConverter>();
+  auto indexType = typeConverter->getIndexType();
+
   // SPIR-V invocation builtin variables are a vector of type <3xi32>
-  auto spirvBuiltin = spirv::getBuiltinVariableValue(op, builtin, rewriter);
+  auto spirvBuiltin =
+      spirv::getBuiltinVariableValue(op, builtin, indexType, rewriter);
   rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
-      op, rewriter.getIntegerType(32), spirvBuiltin,
+      op, indexType, spirvBuiltin,
       rewriter.getI32ArrayAttr({index.getValue()}));
   return success();
 }
 
 template <typename SourceOp, spirv::BuiltIn builtin>
 LogicalResult
 SingleDimLaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
     SourceOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
-  auto spirvBuiltin = spirv::getBuiltinVariableValue(op, builtin, rewriter);
+  auto *typeConverter = this->template getTypeConverter<SPIRVTypeConverter>();
+  auto indexType = typeConverter->getIndexType();
+
+  auto spirvBuiltin =
+      spirv::getBuiltinVariableValue(op, builtin, indexType, rewriter);
   rewriter.replaceOp(op, spirvBuiltin);
   return success();
 }
 
 LogicalResult WorkGroupSizeConversion::matchAndRewrite(
     gpu::BlockDimOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   auto index = getLaunchConfigIndex(op);

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

 
 //===----------------------------------------------------------------------===//
 // Utilities
 //===----------------------------------------------------------------------===//
 
 /// Returns a `Value` containing the `dim`-th dimension's size of SPIR-V
 /// location invocation ID. This function will create necessary operations with
 /// `builder` at the proper region containing `op`.
-static Value getLocalInvocationDimSize(Operation *op, int dim, Location loc,
-                                       OpBuilder *builder) {
+static Value getLocalInvocationDimSize(Operation *op, int dim, Type integerType,
+                                       Location loc, OpBuilder *builder) {
   assert(dim >= 0 && dim < 3 && "local invocation only has three dimensions");
   Value invocation = spirv::getBuiltinVariableValue(
-      op, spirv::BuiltIn::LocalInvocationId, *builder);
+      op, spirv::BuiltIn::LocalInvocationId, integerType, *builder);
   Type xType = invocation.getType().cast<ShapedType>().getElementType();
   return builder->create<spirv::CompositeExtractOp>(
       loc, xType, invocation, builder->getI32ArrayAttr({dim}));
 }
 
 //===----------------------------------------------------------------------===//
 // Reduction (single workgroup)
 //===----------------------------------------------------------------------===//
if (llvm::any_of(llvm::drop_begin(localSize.getIntValues(), 1),
     return failure();
 
   // TODO: Query the target environment to make sure the current
   // workload fits in a local workgroup.
 
   Value convertedInput = operands[0], convertedOutput = operands[1];
   Location loc = genericOp.getLoc();
 
+  auto *typeConverter = getTypeConverter<SPIRVTypeConverter>();
+  auto indexType = typeConverter->getIndexType();
+
   // Get the invocation ID.
-  Value x = getLocalInvocationDimSize(genericOp, /*dim=*/0, loc, &rewriter);
+  Value x = getLocalInvocationDimSize(genericOp, /*dim=*/0, indexType, loc,
+                                      &rewriter);
 
   // TODO: Load to Workgroup storage class first.
 
-  auto *typeConverter = getTypeConverter<SPIRVTypeConverter>();
 
   // Get the input element accessed by this invocation.
   Value inputElementPtr = spirv::getElementPtr(
       *typeConverter, originalInputType, convertedInput, {x}, loc, rewriter);
   Value inputElement = rewriter.create<spirv::LoadOp>(loc, inputElementPtr);
 
   // Perform the group reduction operation.
   Value groupOperation;
 #define CREATE_GROUP_NON_UNIFORM_BIN_OP(opKind, spvOp)                         \
   case linalg::RegionMatcher::BinaryOpKind::opKind: {                          \
     groupOperation = rewriter.create<spirv::spvOp>(                            \
         loc, originalInputType.getElementType(), spirv::Scope::Subgroup,       \
         spirv::GroupOperation::Reduce, inputElement,                           \
         /*cluster_size=*/nullptr);                                             \
   } break
   switch (*binaryOpKind) {
     CREATE_GROUP_NON_UNIFORM_BIN_OP(IAdd, GroupNonUniformIAddOp);
   }
 #undef CREATE_GROUP_NON_UNIFORM_BIN_OP
 
   // Get the output element accessed by this reduction.
-  Value zero = spirv::ConstantOp::getZero(
-      typeConverter->getIndexType(rewriter.getContext()), loc, rewriter);
+  Value zero = spirv::ConstantOp::getZero(indexType, loc, rewriter);
   SmallVector<Value, 1> zeroIndices(originalOutputType.getRank(), zero);
   Value outputElementPtr =
       spirv::getElementPtr(*typeConverter, originalOutputType, convertedOutput,
                            zeroIndices, loc, rewriter);
 
   // Write out the final reduction result. This should be only conducted by one
   // invocation. We use spv.GroupNonUniformElect to find the invocation with the
   // lowest ID.

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

if (adaptor.tensor().getDefiningOp<spirv::ConstantOp>()) {
           loc, varType, spirv::StorageClass::Function,
           /*initializer=*/adaptor.tensor());
     } else {
       // Need to store the value to the local variable. It's questionable
       // whether we want to support such case though.
       return failure();
     }
 
+    auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
+    auto indexType = typeConverter.getIndexType();
+
     Value index = spirv::linearizeIndex(adaptor.indices(), strides,
-                                        /*offset=*/0, loc, rewriter);
+                                        /*offset=*/0, indexType, loc, rewriter);
     auto acOp = rewriter.create<spirv::AccessChainOp>(loc, varOp, index);
 
     rewriter.replaceOpWithNewOp<spirv::LoadOp>(extractOp, acOp);
 
     return success();
   }
 
 private:

mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp

LogicalResult ProcessInterfaceVarABI::matchAndRewrite(
   if (!funcOp->getAttrOfType<spirv::EntryPointABIAttr>(
           spirv::getEntryPointABIAttrName())) {
     // TODO: Non-entry point functions are not handled.
     return failure();
   }
   TypeConverter::SignatureConversion signatureConverter(
       funcOp.getType().getNumInputs());
 
+  auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
+  auto indexType = typeConverter.getIndexType();
+
   auto attrName = spirv::getInterfaceVarABIAttrName();
   for (auto argType : llvm::enumerate(funcOp.getType().getInputs())) {
     auto abiInfo = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
         argType.index(), attrName);
     if (!abiInfo) {
       // TODO: For non-entry point functions, it should be legal
       // to pass around scalar/vector values and return a scalar/vector. For now
       // non-entry point functions are not handled in this ABI lowering and will
Value replacement =
         rewriter.create<spirv::AddressOfOp>(funcOp.getLoc(), var);
     // Check if the arg is a scalar or vector type. In that case, the value
     // needs to be loaded into registers.
     // TODO: This is loading value of the scalar into registers
     // at the start of the function. It is probably better to do the load just
     // before the use. There might be multiple loads and currently there is no
     // easy way to replace all uses with a sequence of operations.
     if (argType.value().cast<spirv::SPIRVType>().isScalarOrVector()) {
-      auto indexType = SPIRVTypeConverter::getIndexType(funcOp.getContext());
       auto zero =
           spirv::ConstantOp::getZero(indexType, funcOp.getLoc(), rewriter);
       auto loadPtr = rewriter.create<spirv::AccessChainOp>(
           funcOp.getLoc(), replacement, zero.constant());
       replacement = rewriter.create<spirv::LoadOp>(funcOp.getLoc(), loadPtr);
     }
     signatureConverter.remapInput(argType.index(), replacement);
   }

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

auto structType = needsExplicitLayout(storageClass)
                         : spirv::StructType::get(elementType);
   return spirv::PointerType::get(structType, storageClass);
 }
 
 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
-Type SPIRVTypeConverter::getIndexType(MLIRContext *context) {
-  // Convert to 32-bit integers for now. Might need a way to control this in
-  // future.
-  // TODO: It is probably better to make it 64-bit integers. To
-  // this some support is needed in SPIR-V dialect for Conversion
-  // instructions. The Vulkan spec requires the builtins like
-  // GlobalInvocationID, etc. to be 32-bit (unsigned) integers which should be
-  // SExtended to 64-bit for index computations.
-  return IntegerType::get(context, 32);
+Type SPIRVTypeConverter::getIndexType() const {
+  return IntegerType::get(getContext(), options.use64bitIndex ? 64 : 32);
 }
 
 /// Mapping between SPIR-V storage classes to memref memory spaces.
 ///
 /// Note: memref does not have a defined semantics for each memory space; it
 /// depends on the context where it is used. There are no particular reasons
 /// behind the number assignments; we try to follow NVVM conventions and largely
 /// give common storage classes a smaller number. The hope is use symbolic
#define STORAGE_SPACE_MAP_FN(storage, space)                                   \
   }
 #undef STORAGE_SPACE_MAP_FN
 }
 
 const SPIRVTypeConverter::Options &SPIRVTypeConverter::getOptions() const {
   return options;
 }
 
+MLIRContext *SPIRVTypeConverter::getContext() const {
+  return targetEnv.getAttr().getContext();
+}
+
 #undef STORAGE_SPACE_MAP_LIST
 
 // TODO: This is a utility function that should probably be exposed by the
 // SPIR-V dialect. Keeping it local till the use case arises.
 static Optional<int64_t>
 getTypeNumBytes(const SPIRVTypeConverter::Options &options, Type type) {
   if (type.isa<spirv::ScalarType>()) {
     auto bitWidth = type.getIntOrFloatBitWidth();
SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
   // were tried before.
   //
   // TODO: this assumes that the SPIR-V types are valid to use in
   // the given target environment, which should be the case if the whole
   // pipeline is driven by the same target environment. Still, we probably still
   // want to validate and convert to be safe.
   addConversion([](spirv::SPIRVType type) { return type; });
 
-  addConversion([](IndexType indexType) {
-    return SPIRVTypeConverter::getIndexType(indexType.getContext());
-  });
+  addConversion([this](IndexType /*indexType*/) { return getIndexType(); });
 
   addConversion([this](IntegerType intType) -> Optional<Type> {
     if (auto scalarType = intType.dyn_cast<spirv::ScalarType>())
       return convertScalarType(this->targetEnv, this->options, scalarType);
     return Type();
   });
 
   addConversion([this](FloatType floatType) -> Optional<Type> {

 /// Gets name of global variable for a builtin.
 static std::string getBuiltinVarName(spirv::BuiltIn builtin) {
   return std::string("__builtin_var_") + stringifyBuiltIn(builtin).str() + "__";
 }
 
 /// Gets or inserts a global variable for a builtin within `body` block.
 static spirv::GlobalVariableOp
 getOrInsertBuiltinVariable(Block &body, Location loc, spirv::BuiltIn builtin,
-                           OpBuilder &builder) {
+                           Type integerType, OpBuilder &builder) {
   if (auto varOp = getBuiltinVariable(body, builtin))
     return varOp;
 
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointToStart(&body);
 
   spirv::GlobalVariableOp newVarOp;
   switch (builtin) {
   case spirv::BuiltIn::NumWorkgroups:
   case spirv::BuiltIn::WorkgroupSize:
   case spirv::BuiltIn::WorkgroupId:
   case spirv::BuiltIn::LocalInvocationId:
   case spirv::BuiltIn::GlobalInvocationId: {
-    auto ptrType = spirv::PointerType::get(
-        VectorType::get({3}, builder.getIntegerType(32)),
+    auto ptrType = spirv::PointerType::get(VectorType::get({3}, integerType),
                                            spirv::StorageClass::Input);
     std::string name = getBuiltinVarName(builtin);
     newVarOp =
         builder.create<spirv::GlobalVariableOp>(loc, ptrType, name, builtin);
     break;
   }
   case spirv::BuiltIn::SubgroupId:
   case spirv::BuiltIn::NumSubgroups:
   case spirv::BuiltIn::SubgroupSize: {
-    auto ptrType = spirv::PointerType::get(builder.getIntegerType(32),
-                                           spirv::StorageClass::Input);
+    auto ptrType =
+        spirv::PointerType::get(integerType, spirv::StorageClass::Input);
     std::string name = getBuiltinVarName(builtin);
     newVarOp =
         builder.create<spirv::GlobalVariableOp>(loc, ptrType, name, builtin);
     break;
   }
   default:
     emitError(loc, "unimplemented builtin variable generation for ")
         << stringifyBuiltIn(builtin);
   }
   return newVarOp;
 }
 
 Value mlir::spirv::getBuiltinVariableValue(Operation *op,
                                            spirv::BuiltIn builtin,
+                                           Type integerType,
                                            OpBuilder &builder) {
   Operation *parent = SymbolTable::getNearestSymbolTable(op->getParentOp());
   if (!parent) {
     op->emitError("expected operation to be within a module-like op");
     return nullptr;
   }
 
-  spirv::GlobalVariableOp varOp = getOrInsertBuiltinVariable(
-      *parent->getRegion(0).begin(), op->getLoc(), builtin, builder);
+  spirv::GlobalVariableOp varOp =
+      getOrInsertBuiltinVariable(*parent->getRegion(0).begin(), op->getLoc(),
+                                 builtin, integerType, builder);
   Value ptr = builder.create<spirv::AddressOfOp>(op->getLoc(), varOp);
   return builder.create<spirv::LoadOp>(op->getLoc(), ptr);
 }
 
 //===----------------------------------------------------------------------===//
 // Push constant storage
 //===----------------------------------------------------------------------===//
 
 /// Returns the pointer type for the push constant storage containing
 /// `elementCount` 32-bit integer values.
 static spirv::PointerType getPushConstantStorageType(unsigned elementCount,
-                                                     Builder &builder) {
-  auto arrayType = spirv::ArrayType::get(
-      SPIRVTypeConverter::getIndexType(builder.getContext()), elementCount,
+                                                     Builder &builder,
+                                                     Type indexType) {
+  auto arrayType = spirv::ArrayType::get(indexType, elementCount,
                                          /*stride=*/4);
   auto structType = spirv::StructType::get({arrayType}, /*offsetInfo=*/0);
   return spirv::PointerType::get(structType, spirv::StorageClass::PushConstant);
 }
 
 /// Returns the push constant varible containing `elementCount` 32-bit integer
 /// values in `body`. Returns null op if such an op does not exit.
 static spirv::GlobalVariableOp getPushConstantVariable(Block &body,
                                                        unsigned elementCount) {
static spirv::GlobalVariableOp getPushConstantVariable(Block &body,
   }
   return nullptr;
 }
 
 /// Gets or inserts a global variable for push constant storage containing
 /// `elementCount` 32-bit integer values in `block`.
 static spirv::GlobalVariableOp
 getOrInsertPushConstantVariable(Location loc, Block &block,
-                                unsigned elementCount, OpBuilder &b) {
+                                unsigned elementCount, OpBuilder &b,
+                                Type indexType) {
   if (auto varOp = getPushConstantVariable(block, elementCount))
     return varOp;
 
   auto builder = OpBuilder::atBlockBegin(&block, b.getListener());
-  auto type = getPushConstantStorageType(elementCount, builder);
+  auto type = getPushConstantStorageType(elementCount, builder, indexType);
   const char *name = "__push_constant_var__";
   return builder.create<spirv::GlobalVariableOp>(loc, type, name,
                                                  /*initializer=*/nullptr);
 }
 
 Value spirv::getPushConstantValue(Operation *op, unsigned elementCount,
-                                  unsigned offset, OpBuilder &builder) {
+                                  unsigned offset, Type integerType,
+                                  OpBuilder &builder) {
   Location loc = op->getLoc();
   Operation *parent = SymbolTable::getNearestSymbolTable(op->getParentOp());
   if (!parent) {
     op->emitError("expected operation to be within a module-like op");
     return nullptr;
   }
 
   spirv::GlobalVariableOp varOp = getOrInsertPushConstantVariable(
-      loc, parent->getRegion(0).front(), elementCount, builder);
+      loc, parent->getRegion(0).front(), elementCount, builder, integerType);
 
-  auto i32Type = SPIRVTypeConverter::getIndexType(builder.getContext());
-  Value zeroOp = spirv::ConstantOp::getZero(i32Type, loc, builder);
+  Value zeroOp = spirv::ConstantOp::getZero(integerType, loc, builder);
   Value offsetOp = builder.create<spirv::ConstantOp>(
-      loc, i32Type, builder.getI32IntegerAttr(offset));
+      loc, integerType, builder.getI32IntegerAttr(offset));
   auto addrOp = builder.create<spirv::AddressOfOp>(loc, varOp);
   auto acOp = builder.create<spirv::AccessChainOp>(
       loc, addrOp, llvm::makeArrayRef({zeroOp, offsetOp}));
   return builder.create<spirv::LoadOp>(loc, acOp);
 }
 
 //===----------------------------------------------------------------------===//
 // Index calculation
 //===----------------------------------------------------------------------===//
 
 Value mlir::spirv::linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
-                                  int64_t offset, Location loc,
-                                  OpBuilder &builder) {
+                                  int64_t offset, Type integerType,
+                                  Location loc, OpBuilder &builder) {
   assert(indices.size() == strides.size() &&
          "must provide indices for all dimensions");
 
-  auto indexType = SPIRVTypeConverter::getIndexType(builder.getContext());
-
   // TODO: Consider moving to use affine.apply and patterns converting
   // affine.apply to standard ops. This needs converting to SPIR-V passes to be
   // broken down into progressive small steps so we can have intermediate steps
   // using other dialects. At the moment SPIR-V is the final sink.
 
   Value linearizedIndex = builder.create<spirv::ConstantOp>(
-      loc, indexType, IntegerAttr::get(indexType, offset));
+      loc, integerType, IntegerAttr::get(integerType, offset));
   for (auto index : llvm::enumerate(indices)) {
     Value strideVal = builder.create<spirv::ConstantOp>(
-        loc, indexType, IntegerAttr::get(indexType, strides[index.index()]));
+        loc, integerType,
+        IntegerAttr::get(integerType, strides[index.index()]));
     Value update = builder.create<spirv::IMulOp>(loc, strideVal, index.value());
     linearizedIndex =
         builder.create<spirv::IAddOp>(loc, linearizedIndex, update);
   }
   return linearizedIndex;
 }
 
 spirv::AccessChainOp mlir::spirv::getElementPtr(
     SPIRVTypeConverter &typeConverter, MemRefType baseType, Value basePtr,
     ValueRange indices, Location loc, OpBuilder &builder) {
   // Get base and offset of the MemRefType and verify they are static.
 
   int64_t offset;
   SmallVector<int64_t, 4> strides;
   if (failed(getStridesAndOffset(baseType, strides, offset)) ||
       llvm::is_contained(strides, MemRefType::getDynamicStrideOrOffset()) ||
       offset == MemRefType::getDynamicStrideOrOffset()) {
     return nullptr;
   }
 
-  auto indexType = typeConverter.getIndexType(builder.getContext());
+  auto indexType = typeConverter.getIndexType();
 
   SmallVector<Value, 2> linearizedIndices;
   auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
 
   // Add a '0' at the start to index into the struct.
   linearizedIndices.push_back(zero);
 
   if (baseType.getRank() == 0) {
     linearizedIndices.push_back(zero);
   } else {
     linearizedIndices.push_back(
-        linearizeIndex(indices, strides, offset, loc, builder));
+        linearizeIndex(indices, strides, offset, indexType, loc, builder));
   }
   return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
 }
 
 //===----------------------------------------------------------------------===//
 // SPIR-V ConversionTarget
 //===----------------------------------------------------------------------===//