[mlir][spirv] Use memref memory space for storage class

[antiagainst]

Summary

Previously in SPIRVTypeConverter, we always convert memref types
to StorageBuffer regardless of their memory spaces. This commit
fixes that to let the conversion to look into memory space
properly. For this purpose, a mapping between SPIR-V storage class
and memref memory space is introduced. The mapping is arbitary
decided at the moment and the hope is that we can leverage
string memory space later to be more clear.

Now spv.interface_var_abi cannot contain storage class unless it's
attached to a scalar value, where we need the storage class as side
channel information. Verifications and tests are properly adjusted.

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

870–873

nit: ///

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

41

nit: ///

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

15

Is this not included transitively already?

[mravishankar]

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

885–891

Nit : Update the error message

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

48

we should try to keep it close to what GPU dialect uses. There it is using a convention that NVVM adopts (https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#address-spaces).

I think this mapping would be most consistent (there is no direct mapping)

StorageBuffer : 0, 1
Workgroup : 3
PushConstant : 4
Private : 5.

Lets leave 2 out since it is reserved.

[antiagainst]

Address comments

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

885–891

I plan to turn this into a proper SPIR-V dialect attribute too. It's becoming more annoying seeing the complicated internal structure of a dictionary attribute and we'll have more fields later.

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

48

Adjusted. I guess the benefit of having similar mapping is mostly cognitive so that folks familiar with the NVVM mapping is not that surprised. But I'd avoid having non-1-1 mapping given it's not round trippable and also confusing. So for now,

Generic: 1
StorageBuffer: 0
Workgroup: 3
Uniform: 4
Private: 5
Function: 6
PushConstant: 7

We can swap Generic and StorageBuffer later after IREE side is fixed.

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

15

Yes. I think it was added by YouCompleteMe. Updated and tweaked some configuration for my vim plugins, now seeing this.

[bondhugula]

mlir/include/mlir/IR/Types.h

172

Typo: of this -> if this

174

of this -> if this

[antiagainst]

Address comments

mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h

 /// satisfy shader interface requirements: shader interface types must be
 /// pointers to structs.
 class SPIRVTypeConverter : public TypeConverter {
 public:
   SPIRVTypeConverter();
 
   /// Gets the SPIR-V correspondence for the standard index type.
   static Type getIndexType(MLIRContext *context);
+
+  /// 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);
 };
 
 /// Base class to define a conversion pattern to lower `SourceOp` into SPIR-V.
 template <typename SourceOp>
 class SPIRVOpLowering : public OpConversionPattern<SourceOp> {
 public:
   SPIRVOpLowering(MLIRContext *context, SPIRVTypeConverter &typeConverter,
                   PatternBenefit benefit = 1)

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

 enum class StorageClass : uint32_t;
 
 /// Returns the attribute name for specifying argument ABI information.
 StringRef getInterfaceVarABIAttrName();
 
 /// Gets the InterfaceVarABIAttr given its fields.
 InterfaceVarABIAttr getInterfaceVarABIAttr(unsigned descriptorSet,
                                            unsigned binding,
-                                           StorageClass storageClass,
+                                           Optional<StorageClass> storageClass,
                                            MLIRContext *context);
 
 /// Returns the attribute name for specifying entry point information.
 StringRef getEntryPointABIAttrName();
 
 /// Gets the EntryPointABIAttr given its fields.
 EntryPointABIAttr getEntryPointABIAttr(ArrayRef<int32_t> localSize,
                                        MLIRContext *context);

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td

 // shader interface, this attribute specifies the information regarding
 // the global variable:
 // 1) Descriptor Set.
 // 2) Binding number.
 // 3) Storage class.
 def SPV_InterfaceVarABIAttr : StructAttr<"InterfaceVarABIAttr", SPIRV_Dialect, [
     StructFieldAttr<"descriptor_set", I32Attr>,
     StructFieldAttr<"binding", I32Attr>,
-    StructFieldAttr<"storage_class", SPV_StorageClassAttr>
+    StructFieldAttr<"storage_class", OptionalAttr<SPV_StorageClassAttr>>
 ]>;
 
 // For entry functions, this attribute specifies information related to entry
 // points in the generated SPIR-V module:
 // 1) WorkGroup Size.
 def SPV_EntryPointABIAttr : StructAttr<"EntryPointABIAttr", SPIRV_Dialect, [
     StructFieldAttr<"local_size", I32ElementsAttr>
 ]>;

mlir/include/mlir/IR/Types.h

#include "DialectSymbolRegistry.def"
 
   /// Return true if this is a signless integer or index type.
   bool isSignlessIntOrIndex();
   /// Return true if this is a signless integer, index, or float type.
   bool isSignlessIntOrIndexOrFloat();
   /// Return true of this is a signless integer or a float type.
   bool isSignlessIntOrFloat();
 
-  /// Return true of this is an integer(of any signedness) or a float type.
+  /// Return true if this is an integer (of any signedness) or a float type.
   bool isIntOrFloat();
+  /// Return true if this is an integer (of any signedness), index, or float
+  /// type.
+  bool isIntOrIndexOrFloat();
 
   /// Print the current type.
   void print(raw_ostream &os);
   void dump();
 
   friend ::llvm::hash_code hash_value(Type arg);
 
   unsigned getSubclassData() const;

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

 }
 
 PatternMatchResult GPUFuncOpConversion::matchAndRewrite(
     gpu::GPUFuncOp funcOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   if (!gpu::GPUDialect::isKernel(funcOp))
     return matchFailure();
 
+  // TODO(antiagainst): we are dictating the ABI by ourselves here; it should be
+  // specified outside.
   SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
-  for (auto argNum : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
-    argABI.push_back(spirv::getInterfaceVarABIAttr(
-        0, argNum, spirv::StorageClass::StorageBuffer, rewriter.getContext()));
+  for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
+    Optional<spirv::StorageClass> sc;
+    if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
+      sc = spirv::StorageClass::StorageBuffer;
+    argABI.push_back(
+        spirv::getInterfaceVarABIAttr(0, argIndex, sc, rewriter.getContext()));
   }
 
   auto entryPointAttr = spirv::lookupEntryPointABI(funcOp);
   if (!entryPointAttr) {
     funcOp.emitRemark("match failure: missing 'spv.entry_point_abi' attribute");
     return matchFailure();
   }
   spirv::FuncOp newFuncOp = lowerAsEntryFunction(

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

struct SPIRVInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   /// Returns true if the given region 'src' can be inlined into the region
   /// 'dest' that is attached to an operation registered to the current dialect.
   bool isLegalToInline(Region *dest, Region *src,
                        BlockAndValueMapping &) const final {
     // Return true here when inlining into spv.func, spv.selection, and
     // spv.loop operations.
-    auto op = dest->getParentOp();
+    auto *op = dest->getParentOp();
     return isa<spirv::FuncOp>(op) || isa<spirv::SelectionOp>(op) ||
            isa<spirv::LoopOp>(op);
   }
 
   /// Returns true if the given operation 'op', that is registered to this
   /// dialect, can be inlined into the region 'dest' that is attached to an
   /// operation registered to the current dialect.
   bool isLegalToInline(Operation *op, Region *dest,
Optional<uint64_t> parseAndVerify<uint64_t>(SPIRVDialect const &dialect,
   return parseAndVerifyInteger<uint64_t>(dialect, parser);
 }
 
 namespace {
 // Functor object to parse a comma separated list of specs. The function
 // parseAndVerify does the actual parsing and verification of individual
 // elements. This is a functor since parsing the last element of the list
 // (termination condition) needs partial specialization.
-template <typename ParseType, typename... Args> struct parseCommaSeparatedList {
+template <typename ParseType, typename... Args>
+struct ParseCommaSeparatedList {
   Optional<std::tuple<ParseType, Args...>>
   operator()(SPIRVDialect const &dialect, DialectAsmParser &parser) const {
     auto parseVal = parseAndVerify<ParseType>(dialect, parser);
     if (!parseVal)
       return llvm::None;
 
     auto numArgs = std::tuple_size<std::tuple<Args...>>::value;
     if (numArgs != 0 && failed(parser.parseComma()))
       return llvm::None;
-    auto remainingValues = parseCommaSeparatedList<Args...>{}(dialect, parser);
+    auto remainingValues = ParseCommaSeparatedList<Args...>{}(dialect, parser);
     if (!remainingValues)
       return llvm::None;
     return std::tuple_cat(std::tuple<ParseType>(parseVal.getValue()),
                           remainingValues.getValue());
   }
 };
 
 // Partial specialization of the function to parse a comma separated list of
 // specs to parse the last element of the list.
-template <typename ParseType> struct parseCommaSeparatedList<ParseType> {
+template <typename ParseType>
+struct ParseCommaSeparatedList<ParseType> {
   Optional<std::tuple<ParseType>> operator()(SPIRVDialect const &dialect,
                                              DialectAsmParser &parser) const {
     if (auto value = parseAndVerify<ParseType>(dialect, parser))
       return std::tuple<ParseType>(value.getValue());
     return llvm::None;
   }
 };
 } // namespace

 //                              arrayed-info `,` sampling-info `,`
 //                              sampler-use-info `,` format `>`
 static Type parseImageType(SPIRVDialect const &dialect,
                            DialectAsmParser &parser) {
   if (parser.parseLess())
     return Type();
 
   auto value =
-      parseCommaSeparatedList<Type, Dim, ImageDepthInfo, ImageArrayedInfo,
+      ParseCommaSeparatedList<Type, Dim, ImageDepthInfo, ImageArrayedInfo,
                               ImageSamplingInfo, ImageSamplerUseInfo,
                               ImageFormat>{}(dialect, parser);
   if (!value)
     return Type();
 
   if (parser.parseGreater())
     return Type();
   return ImageType::get(value.getValue());
auto printMember = [&](unsigned i) {
     type.getMemberDecorations(i, decorations);
     if (type.hasLayout() || !decorations.empty()) {
       os << " [";
       if (type.hasLayout()) {
         os << type.getOffset(i);
         if (!decorations.empty())
           os << ", ";
       }
-      auto each_fn = [&os](spirv::Decoration decoration) {
+      auto eachFn = [&os](spirv::Decoration decoration) {
         os << stringifyDecoration(decoration);
       };
-      interleaveComma(decorations, os, each_fn);
+      interleaveComma(decorations, os, eachFn);
       os << "]";
     }
   };
   interleaveComma(llvm::seq<unsigned>(0, type.getNumElements()), os,
                   printMember);
   os << ">";
 }
 
LogicalResult SPIRVDialect::verifyOperationAttribute(Operation *op,
   } else {
     return op->emitError("found unsupported '")
            << symbol << "' attribute on operation";
   }
 
   return success();
 }
 
-// Verifies the given SPIR-V `attribute` attached to a region's argument or
-// result and reports error to the given location if invalid.
-static LogicalResult
-verifyRegionAttribute(Location loc, NamedAttribute attribute, bool forArg) {
+/// Verifies the given SPIR-V `attribute` attached to a value of the given
+/// `valueType` is valid.
+static LogicalResult verifyRegionAttribute(Location loc, Type valueType,
+                                           NamedAttribute attribute) {
   StringRef symbol = attribute.first.strref();
   Attribute attr = attribute.second;
 
   if (symbol != spirv::getInterfaceVarABIAttrName())
     return emitError(loc, "found unsupported '")
-           << symbol << "' attribute on region "
-           << (forArg ? "argument" : "result");
+           << symbol << "' attribute on region argument";
 
-  if (!attr.isa<spirv::InterfaceVarABIAttr>())
+  auto varABIAttr = attr.dyn_cast<spirv::InterfaceVarABIAttr>();
+  if (!varABIAttr)
     return emitError(loc, "'")
            << symbol
-           << "' attribute must be a dictionary attribute containing three "
-              "32-bit integer attributes: 'descriptor_set', 'binding', and "
-              "'storage_class'";
+           << "' attribute must be a dictionary attribute containing two or "
+              "three 32-bit integer attributes: 'descriptor_set', 'binding', "
+              "and optional 'storage_class'";
+  if (varABIAttr.storage_class() && !valueType.isIntOrIndexOrFloat())
+    return emitError(loc, "'") << symbol
+                               << "' attribute cannot specify storage class "
+                                  "when attaching to a non-scalar value";
 
   return success();
 }
 
 LogicalResult SPIRVDialect::verifyRegionArgAttribute(Operation *op,
-                                                     unsigned /*regionIndex*/,
-                                                     unsigned /*argIndex*/,
+                                                     unsigned regionIndex,
+                                                     unsigned argIndex,
                                                      NamedAttribute attribute) {
-  return verifyRegionAttribute(op->getLoc(), attribute,
-                               /*forArg=*/true);
+  return verifyRegionAttribute(
+      op->getLoc(),
+      op->getRegion(regionIndex).front().getArgument(argIndex).getType(),
+      attribute);
 }
 
 LogicalResult SPIRVDialect::verifyRegionResultAttribute(
     Operation *op, unsigned /*regionIndex*/, unsigned /*resultIndex*/,
     NamedAttribute attribute) {
-  return verifyRegionAttribute(op->getLoc(), attribute,
-                               /*forArg=*/false);
+  return op->emitError("cannot attach SPIR-V attributes to region result");
 }

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

Type SPIRVTypeConverter::getIndexType(MLIRContext *context) {
   // TODO(ravishankarm): 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(32, context);
 }
 
+/// Mapping between SPIR-V storage classes to memref memory spaces.
+///
+/// Note: memref does not have a defined smenatics for each memory space; it
+/// depends on the context where it is used. There are no particular reasons
+/// behind the number assigments; we try to follow NVVM conventions and largely
+/// give common storage classes a smaller number. The hope is use symbolic
+/// memory space representation eventually after memref supports it.
+// TODO(antiagainst): swap Generic and StorageBuffer assignment to be more akin
+// to NVVM.
+#define STORAGE_SPACE_MAP_LIST(MAP_FN)                                         \
+  MAP_FN(spirv::StorageClass::Generic, 1)                                      \
+  MAP_FN(spirv::StorageClass::StorageBuffer, 0)                                \
+  MAP_FN(spirv::StorageClass::Workgroup, 3)                                    \
+  MAP_FN(spirv::StorageClass::Uniform, 4)                                      \
+  MAP_FN(spirv::StorageClass::Private, 5)                                      \
+  MAP_FN(spirv::StorageClass::Function, 6)                                     \
+  MAP_FN(spirv::StorageClass::PushConstant, 7)                                 \
+  MAP_FN(spirv::StorageClass::UniformConstant, 8)                              \
+  MAP_FN(spirv::StorageClass::Input, 9)                                        \
+  MAP_FN(spirv::StorageClass::Output, 10)                                      \
+  MAP_FN(spirv::StorageClass::CrossWorkgroup, 11)                              \
+  MAP_FN(spirv::StorageClass::AtomicCounter, 12)                               \
+  MAP_FN(spirv::StorageClass::Image, 13)                                       \
+  MAP_FN(spirv::StorageClass::CallableDataNV, 14)                              \
+  MAP_FN(spirv::StorageClass::IncomingCallableDataNV, 15)                      \
+  MAP_FN(spirv::StorageClass::RayPayloadNV, 16)                                \
+  MAP_FN(spirv::StorageClass::HitAttributeNV, 17)                              \
+  MAP_FN(spirv::StorageClass::IncomingRayPayloadNV, 18)                        \
+  MAP_FN(spirv::StorageClass::ShaderRecordBufferNV, 19)                        \
+  MAP_FN(spirv::StorageClass::PhysicalStorageBuffer, 20)
+
+unsigned
+SPIRVTypeConverter::getMemorySpaceForStorageClass(spirv::StorageClass storage) {
+#define STORAGE_SPACE_MAP_FN(storage, space)                                   \
+  case storage:                                                                \
+    return space;
+
+  switch (storage) { STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN) }
+#undef STORAGE_SPACE_MAP_FN
+}
+
+Optional<spirv::StorageClass>
+SPIRVTypeConverter::getStorageClassForMemorySpace(unsigned space) {
+#define STORAGE_SPACE_MAP_FN(storage, space)                                   \
+  case space:                                                                  \
+    return storage;
+
+  switch (space) {
+    STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
+  default:
+    return llvm::None;
+  }
+#undef STORAGE_SPACE_MAP_FN
+}
+
+#undef STORAGE_SPACE_MAP_LIST
+
 // TODO(ravishankarm): 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(Type t) {
   if (spirv::SPIRVDialect::isValidScalarType(t)) {
     auto bitWidth = t.getIntOrFloatBitWidth();
     // According to the SPIR-V spec:
     // "There is no physical size or bit pattern defined for values with boolean
     // type. If they are stored (in conjunction with OpVariable), they can only
SPIRVTypeConverter::SPIRVTypeConverter() {
   addConversion([](Type type) -> Optional<Type> {
     // If the type is already valid in SPIR-V, directly return.
     return spirv::SPIRVDialect::isValidType(type) ? type : Optional<Type>();
   });
   addConversion([](IndexType indexType) {
     return SPIRVTypeConverter::getIndexType(indexType.getContext());
   });
   addConversion([this](MemRefType memRefType) -> Type {
-    // TODO(ravishankarm): For now only support default memory space. The memory
-    // space description is not set is stone within MLIR, i.e. it depends on the
-    // context it is being used. To map this to SPIR-V storage classes, we
-    // should rely on the ABI attributes, and not on the memory space. This is
-    // still evolving, and needs to be revisited when there is more clarity.
-    if (memRefType.getMemorySpace())
-      return Type();
-
     auto elementType = convertType(memRefType.getElementType());
     if (!elementType)
       return Type();
 
     auto elementSize = getTypeNumBytes(elementType);
     if (!elementSize)
       return Type();
 
     // TODO(ravishankarm) : Handle dynamic shapes.
     if (memRefType.hasStaticShape()) {
       auto arraySize = getTypeNumBytes(memRefType);
       if (!arraySize)
         return Type();
 
       auto arrayType = spirv::ArrayType::get(
           elementType, arraySize.getValue() / elementSize.getValue(),
           elementSize.getValue());
+
+      // Wrap in a struct to satisfy Vulkan interface requirements.
       auto structType = spirv::StructType::get(arrayType, 0);
-      // For now initialize the storage class to StorageBuffer. This will be
-      // updated later based on whats passed in w.r.t to the ABI attributes.
-      return spirv::PointerType::get(structType,
-                                     spirv::StorageClass::StorageBuffer);
+      if (auto sc = getStorageClassForMemorySpace(memRefType.getMemorySpace()))
+        return spirv::PointerType::get(structType, *sc);
+      return Type();
     }
     return Type();
   });
   addConversion([this](TensorType tensorType) -> Type {
     // TODO(ravishankarm) : Handle dynamic shapes.
     if (!tensorType.hasStaticShape())
       return Type();
 

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

 //===- TargetAndABI.cpp - SPIR-V target and ABI utilities -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/TargetAndABI.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
 
 using namespace mlir;
 
 StringRef spirv::getInterfaceVarABIAttrName() {
   return "spv.interface_var_abi";
 }
 
 spirv::InterfaceVarABIAttr
 spirv::getInterfaceVarABIAttr(unsigned descriptorSet, unsigned binding,
-                              spirv::StorageClass storageClass,
+                              Optional<spirv::StorageClass> storageClass,
                               MLIRContext *context) {
   Type i32Type = IntegerType::get(32, context);
+  auto scAttr =
+      storageClass
+          ? IntegerAttr::get(i32Type, static_cast<int64_t>(*storageClass))
+          : IntegerAttr();
   return spirv::InterfaceVarABIAttr::get(
       IntegerAttr::get(i32Type, descriptorSet),
-      IntegerAttr::get(i32Type, binding),
-      IntegerAttr::get(i32Type, static_cast<int64_t>(storageClass)), context);
+      IntegerAttr::get(i32Type, binding), scAttr, context);
 }
 
 StringRef spirv::getEntryPointABIAttrName() { return "spv.entry_point_abi"; }
 
 spirv::EntryPointABIAttr
 spirv::getEntryPointABIAttr(ArrayRef<int32_t> localSize, MLIRContext *context) {
   assert(localSize.size() == 3);
   return spirv::EntryPointABIAttr::get(

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

 /// Checks if the `type` is a scalar or vector type. It is assumed that they are
 /// valid for SPIR-V dialect already.
 static bool isScalarOrVectorType(Type type) {
   return spirv::SPIRVDialect::isValidScalarType(type) || type.isa<VectorType>();
 }
 
 /// Creates a global variable for an argument based on the ABI info.
 static spirv::GlobalVariableOp
-createGlobalVariableForArg(spirv::FuncOp funcOp, OpBuilder &builder,
-                           unsigned argNum,
+createGlobalVarForEntryPointArgument(OpBuilder &builder, spirv::FuncOp funcOp,
+                                     unsigned argIndex,
                                      spirv::InterfaceVarABIAttr abiInfo) {
   auto spirvModule = funcOp.getParentOfType<spirv::ModuleOp>();
-  if (!spirvModule) {
+  if (!spirvModule)
     return nullptr;
-  }
+
   OpBuilder::InsertionGuard moduleInsertionGuard(builder);
   builder.setInsertionPoint(funcOp.getOperation());
   std::string varName =
-      funcOp.getName().str() + "_arg_" + std::to_string(argNum);
+      funcOp.getName().str() + "_arg_" + std::to_string(argIndex);
 
   // Get the type of variable. If this is a scalar/vector type and has an ABI
-  // info create a variable of type !spv.ptr<!spv.struct<elementTYpe>>. If not
+  // info create a variable of type !spv.ptr<!spv.struct<elementType>>. If not
   // it must already be a !spv.ptr<!spv.struct<...>>.
-  auto varType = funcOp.getType().getInput(argNum);
+  auto varType = funcOp.getType().getInput(argIndex);
+  if (isScalarOrVectorType(varType)) {
     auto storageClass =
         static_cast<spirv::StorageClass>(abiInfo.storage_class().getInt());
-  if (isScalarOrVectorType(varType)) {
     varType =
         spirv::PointerType::get(spirv::StructType::get(varType), storageClass);
   }
   auto varPtrType = varType.cast<spirv::PointerType>();
   auto varPointeeType = varPtrType.getPointeeType().cast<spirv::StructType>();
 
   // Set the offset information.
   VulkanLayoutUtils::Size size = 0, alignment = 0;
getInterfaceVariables(spirv::FuncOp funcOp,
   llvm::SetVector<Operation *> interfaceVarSet;
 
   // TODO(ravishankarm) : This should in reality traverse the entry function
   // call graph and collect all the interfaces. For now, just traverse the
   // instructions in this function.
   funcOp.walk([&](spirv::AddressOfOp addressOfOp) {
     auto var =
         module.lookupSymbol<spirv::GlobalVariableOp>(addressOfOp.variable());
-    if (var.type().cast<spirv::PointerType>().getStorageClass() !=
-        spirv::StorageClass::StorageBuffer) {
+    // TODO(antiagainst): Per SPIR-V spec: "Before version 1.4, the interface’s
+    // storage classes are limited to the Input and Output storage classes.
+    // Starting with version 1.4, the interface’s storage classes are all
+    // storage classes used in declaring all global variables referenced by the
+    // entry point’s call tree." We should consider the target environment here.
+    switch (var.type().cast<spirv::PointerType>().getStorageClass()) {
+    case spirv::StorageClass::Input:
+    case spirv::StorageClass::Output:
       interfaceVarSet.insert(var.getOperation());
+      break;
+    default:
+      break;
     }
   });
   for (auto &var : interfaceVarSet) {
     interfaceVars.push_back(SymbolRefAttr::get(
         cast<spirv::GlobalVariableOp>(var).sym_name(), funcOp.getContext()));
   }
   return success();
 }
auto abiInfo = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
         argType.index(), attrName);
     if (!abiInfo) {
       // TODO(ravishankarm) : 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
       // produce an error.
       return matchFailure();
     }
-    auto var =
-        createGlobalVariableForArg(funcOp, rewriter, argType.index(), abiInfo);
-    if (!var) {
+    spirv::GlobalVariableOp var = createGlobalVarForEntryPointArgument(
+        rewriter, funcOp, argType.index(), abiInfo);
+    if (!var)
       return matchFailure();
-    }
 
     OpBuilder::InsertionGuard funcInsertionGuard(rewriter);
     rewriter.setInsertionPointToStart(&funcOp.front());
     // Insert spirv::AddressOf and spirv::AccessChain operations.
     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.

mlir/lib/IR/StandardTypes.cpp

 }
 
 bool Type::isSignlessIntOrFloat() {
   return isSignlessInteger() || isa<FloatType>();
 }
 
 bool Type::isIntOrFloat() { return isa<IntegerType>() || isa<FloatType>(); }
 
+bool Type::isIntOrIndexOrFloat() { return isIntOrFloat() || isIndex(); }
+
 //===----------------------------------------------------------------------===//
 // Integer Type
 //===----------------------------------------------------------------------===//
 
 // static constexpr must have a definition (until in C++17 and inline variable).
 constexpr unsigned IntegerType::kMaxWidth;
 
 /// Verify the construction of an integer type.

mlir/test/Conversion/GPUToSPIRV/load-store.mlir

module attributes {gpu.container_module} {
   }
 
   // CHECK-LABEL: spv.module Logical GLSL450
   gpu.module @kernels {
     // CHECK-DAG: spv.globalVariable [[NUMWORKGROUPSVAR:@.*]] built_in("NumWorkgroups") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-DAG: spv.globalVariable [[LOCALINVOCATIONIDVAR:@.*]] built_in("LocalInvocationId") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-DAG: spv.globalVariable [[WORKGROUPIDVAR:@.*]] built_in("WorkgroupId") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-LABEL:    spv.func @load_store_kernel
-    // CHECK-SAME: [[ARG0:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
-    // CHECK-SAME: [[ARG1:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
-    // CHECK-SAME: [[ARG2:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 2 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
+    // CHECK-SAME: [[ARG0:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32{{[}][}]}}
+    // CHECK-SAME: [[ARG1:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32{{[}][}]}}
+    // CHECK-SAME: [[ARG2:%.*]]: !spv.ptr<!spv.struct<!spv.array<48 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 2 : i32, descriptor_set = 0 : i32{{[}][}]}}
     // CHECK-SAME: [[ARG3:%.*]]: i32 {spv.interface_var_abi = {binding = 3 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     // CHECK-SAME: [[ARG4:%.*]]: i32 {spv.interface_var_abi = {binding = 4 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     // CHECK-SAME: [[ARG5:%.*]]: i32 {spv.interface_var_abi = {binding = 5 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     // CHECK-SAME: [[ARG6:%.*]]: i32 {spv.interface_var_abi = {binding = 6 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
     gpu.func @load_store_kernel(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>, %arg3: index, %arg4: index, %arg5: index, %arg6: index)
       attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       // CHECK: [[ADDRESSWORKGROUPID:%.*]] = spv._address_of [[WORKGROUPIDVAR]]
       // CHECK: [[WORKGROUPID:%.*]] = spv.Load "Input" [[ADDRESSWORKGROUPID]]

mlir/test/Conversion/GPUToSPIRV/simple.mlir

 // RUN: mlir-opt -split-input-file -convert-gpu-to-spirv -verify-diagnostics %s -o - | FileCheck %s
 
 module attributes {gpu.container_module} {
   gpu.module @kernels {
     // CHECK:       spv.module Logical GLSL450 {
     // CHECK-LABEL: spv.func @basic_module_structure
     // CHECK-SAME: {{%.*}}: f32 {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
-    // CHECK-SAME: {{%.*}}: !spv.ptr<!spv.struct<!spv.array<12 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32{{[}][}]}}
+    // CHECK-SAME: {{%.*}}: !spv.ptr<!spv.struct<!spv.array<12 x f32 [4]> [0]>, StorageBuffer> {spv.interface_var_abi = {binding = 1 : i32, descriptor_set = 0 : i32{{[}][}]}}
     // CHECK-SAME: spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>}
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32>)
       attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
       // CHECK: spv.Return
       gpu.return
     }
   }
 

mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir

 // memref type
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func @memref_type({{%.*}}: memref<3xi1>)
 func @memref_type(%arg0: memref<3xi1>) {
   return
 }
 
+// CHECK-LABEL: func @memref_mem_space
+// CHECK-SAME: StorageBuffer
+// CHECK-SAME: Uniform
+// CHECK-SAME: Workgroup
+// CHECK-SAME: PushConstant
+// CHECK-SAME: Private
+// CHECK-SAME: Function
+func @memref_mem_space(
+    %arg0: memref<4xf32, 0>,
+    %arg1: memref<4xf32, 4>,
+    %arg2: memref<4xf32, 3>,
+    %arg3: memref<4xf32, 7>,
+    %arg4: memref<4xf32, 5>,
+    %arg5: memref<4xf32, 6>
+) { return }
+
 // CHECK-LABEL: @load_store_zero_rank_float
 // CHECK: [[ARG0:%.*]]: !spv.ptr<!spv.struct<!spv.array<1 x f32 [4]> [0]>, StorageBuffer>,
 // CHECK: [[ARG1:%.*]]: !spv.ptr<!spv.struct<!spv.array<1 x f32 [4]> [0]>, StorageBuffer>)
 func @load_store_zero_rank_float(%arg0: memref<f32>, %arg1: memref<f32>) {
   //      CHECK: [[ZERO1:%.*]] = spv.constant 0 : i32
   //      CHECK: spv.AccessChain [[ARG0]][
   // CHECK-SAME: [[ZERO1]], [[ZERO1]]
   // CHECK-SAME: ] :

mlir/test/Dialect/SPIRV/Transforms/abi-interface.mlir

 // RUN: mlir-opt -spirv-lower-abi-attrs -verify-diagnostics %s -o - | FileCheck %s
 
 // CHECK-LABEL: spv.module
 spv.module Logical GLSL450 {
   // CHECK-DAG:    spv.globalVariable [[VAR0:@.*]] bind(0, 0) : !spv.ptr<!spv.struct<f32 [0]>, StorageBuffer>
   // CHECK-DAG:    spv.globalVariable [[VAR1:@.*]] bind(0, 1) : !spv.ptr<!spv.struct<!spv.array<12 x f32 [4]> [0]>, StorageBuffer>
   // CHECK:    spv.func [[FN:@.*]]()
-  spv.func @kernel(%arg0: f32
-                {spv.interface_var_abi = {binding = 0 : i32,
-                                          descriptor_set = 0 : i32,
-                                          storage_class = 12 : i32}},
-                 %arg1: !spv.ptr<!spv.struct<!spv.array<12 x f32>>, StorageBuffer>
-                 {spv.interface_var_abi = {binding = 1 : i32,
-                                           descriptor_set = 0 : i32,
-                                           storage_class = 12 : i32}}) "None"
+  spv.func @kernel(
+    %arg0: f32
+                                     descriptor_set = 0 : i32}}) "None"
   attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}} {
     // CHECK: [[ARG1:%.*]] = spv._address_of [[VAR1]]
     // CHECK: [[ADDRESSARG0:%.*]] = spv._address_of [[VAR0]]
     // CHECK: [[CONST0:%.*]] = spv.constant 0 : i32
     // CHECK: [[ARG0PTR:%.*]] = spv.AccessChain [[ADDRESSARG0]]{{\[}}[[CONST0]]
     // CHECK: [[ARG0:%.*]] = spv.Load "StorageBuffer" [[ARG0PTR]]
     // CHECK: spv.Return
     spv.Return
   }
   // CHECK: spv.EntryPoint "GLCompute" [[FN]]
   // CHECK: spv.ExecutionMode [[FN]] "LocalSize", 32, 1, 1
 }

mlir/test/Dialect/SPIRV/Transforms/abi-load-store.mlir

spv.module Logical GLSL450 {
   // CHECK-DAG: spv.globalVariable [[VAR3:@.*]] bind(0, 3) : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer>
   // CHECK-DAG: spv.globalVariable [[VAR4:@.*]] bind(0, 4) : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer>
   // CHECK-DAG: spv.globalVariable [[VAR5:@.*]] bind(0, 5) : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer>
   // CHECK-DAG: spv.globalVariable [[VAR6:@.*]] bind(0, 6) : !spv.ptr<!spv.struct<i32 [0]>, StorageBuffer>
   // CHECK: spv.func [[FN:@.*]]()
   spv.func @load_store_kernel(
     %arg0: !spv.ptr<!spv.struct<!spv.array<12 x !spv.array<4 x f32>>>, StorageBuffer>
     {spv.interface_var_abi = {binding = 0 : i32,
-                              descriptor_set = 0 : i32,
-                              storage_class = 12 : i32}},
+                              descriptor_set = 0 : i32}},
     %arg1: !spv.ptr<!spv.struct<!spv.array<12 x !spv.array<4 x f32>>>, StorageBuffer>
     {spv.interface_var_abi = {binding = 1 : i32,
-                              descriptor_set = 0 : i32,
-                              storage_class = 12 : i32}},
+                              descriptor_set = 0 : i32}},
     %arg2: !spv.ptr<!spv.struct<!spv.array<12 x !spv.array<4 x f32>>>, StorageBuffer>
     {spv.interface_var_abi = {binding = 2 : i32,
-                              descriptor_set = 0 : i32,
-                              storage_class = 12 : i32}},
+                              descriptor_set = 0 : i32}},
     %arg3: i32
     {spv.interface_var_abi = {binding = 3 : i32,
                               descriptor_set = 0 : i32,
                               storage_class = 12 : i32}},
     %arg4: i32
     {spv.interface_var_abi = {binding = 4 : i32,
                               descriptor_set = 0 : i32,
                               storage_class = 12 : i32}},

mlir/test/Dialect/SPIRV/target-and-abi.mlir

 // RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
 
 // expected-error @+1 {{found unsupported 'spv.something' attribute on operation}}
 func @unknown_attr_on_op() attributes {
   spv.something = 64
 } { return }
 
 // -----
 
 // expected-error @+1 {{found unsupported 'spv.something' attribute on region argument}}
 func @unknown_attr_on_region(%arg: i32 {spv.something}) {
   return
 }
 
 // -----
 
-// expected-error @+1 {{found unsupported 'spv.something' attribute on region result}}
+// expected-error @+1 {{cannot attach SPIR-V attributes to region result}}
 func @unknown_attr_on_region() -> (i32 {spv.something}) {
   %0 = constant 10.0 : f32
   return %0: f32
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//

 } { return }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv.interface_var_abi
 //===----------------------------------------------------------------------===//
 
-// expected-error @+1 {{'spv.interface_var_abi' attribute must be a dictionary attribute containing three 32-bit integer attributes: 'descriptor_set', 'binding', and 'storage_class'}}
+// expected-error @+1 {{'spv.interface_var_abi' attribute must be a dictionary attribute containing two or three 32-bit integer attributes: 'descriptor_set', 'binding', and optional 'storage_class'}}
 func @interface_var(
   %arg0 : f32 {spv.interface_var_abi = 64}
 ) { return }
 
 // -----
 
-// expected-error @+1 {{'spv.interface_var_abi' attribute must be a dictionary attribute containing three 32-bit integer attributes: 'descriptor_set', 'binding', and 'storage_class'}}
+// expected-error @+1 {{'spv.interface_var_abi' attribute must be a dictionary attribute containing two or three 32-bit integer attributes: 'descriptor_set', 'binding', and optional 'storage_class'}}
 func @interface_var(
   %arg0 : f32 {spv.interface_var_abi = {binding = 0: i32}}
 ) { return }
 
 // -----
 
 // CHECK: {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32}}
 func @interface_var(
   %arg0 : f32 {spv.interface_var_abi = {binding = 0 : i32,
                                         descriptor_set = 0 : i32,
                                         storage_class = 12 : i32}}
 ) { return }
 
 // -----
 
-// expected-error @+1 {{'spv.interface_var_abi' attribute must be a dictionary attribute containing three 32-bit integer attributes: 'descriptor_set', 'binding', and 'storage_class'}}
-func @interface_var() -> (f32 {spv.interface_var_abi = 64})
-{
-  %0 = constant 10.0 : f32
-  return %0: f32
-}
-
-// -----
-
-// expected-error @+1 {{'spv.interface_var_abi' attribute must be a dictionary attribute containing three 32-bit integer attributes: 'descriptor_set', 'binding', and 'storage_class'}}
-func @interface_var() -> (f32 {spv.interface_var_abi = {binding = 0: i32}})
-{
-  %0 = constant 10.0 : f32
-  return %0: f32
-}
-
-// -----
-
-// CHECK: {spv.interface_var_abi = {binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32}}
-func @interface_var() -> (f32 {spv.interface_var_abi = {
-    binding = 0 : i32, descriptor_set = 0 : i32, storage_class = 12 : i32}})
-{
-  %0 = constant 10.0 : f32
-  return %0: f32
-}
+// expected-error @+1 {{'spv.interface_var_abi' attribute cannot specify storage class when attaching to a non-scalar value}}
+func @interface_var(
+  %arg0 : memref<4xf32> {spv.interface_var_abi = {binding = 0 : i32,
+                                                  descriptor_set = 0 : i32,
+                                                  storage_class = 12 : i32}}
+) { return }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv.target_env
 //===----------------------------------------------------------------------===//
 
 func @target_env_missing_limits() attributes {