[mlir][spirv] Lower allocation/deallocations of workgroup memory.

[mravishankar]

Summary

This allocation of a workgroup memory is lowered to a
spv.globalVariable. Only static size allocation with element type
being int or float is handled. The lowering does account for the
element type that are not supported in the lowered spv.module based on
the extensions/capabilities and adjusts the number of elements to get
the same byte length.

Depends On D80365

[mravishankar]

Rebase

[antiagainst]

Cool, thanks Mahesh! Overall looks good to me. Approval conditional on fixing the workgroup layout annotation issue.

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

46

Returns None

48

Super nit: can we just leave out the t? It does not provide more information here.

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

240

I feel we should place this logic inside the type converter itself. This is basically getting the SPIR-V "equivalent" lowered memref type for the input memref. Letting each pattern figuring out this is not the proper separation of concerns. Fine for now given this is only used by allocation; but can we put a TODO here to remind ourselves?

273

spv.globalVariable

274

Could you document that this pattern requires the pass to work on module instead of function given it injects stuff to the model?

284

Do we need to unique the symbol manually here? I remember there is infrastructure support for it; but cannot recall off the top of my head.

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

61

Might just want to add -canonicalize to simplify here? I know it's kinda tiding two passes together but in general canonicalization is not changing too frequently and we can have more readable tests here. But up to you. :)

mlir/test/Conversion/StandardToSPIRV/alloc.mlir

23

Ah, this should not have layout annotation at all. Variables in Workgroup should not be explicitly laid out. The layout handling is problematic in the type converter. Could you turn off layout annotation for Workgroup memrefs in TypeConverter? Just do not pass in offsets when handling Workgroup in convertMemrefType.

36

You mean fix? This test is not commented out at the moment.

109

We can trim of the but got ... part given we are not particularly interested in verifying the detailed error here. Similarly for the following case.

[mravishankar]

Addressing comments

[mravishankar]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

240

Yup I agree and obviously the type converter already handles this. So removing this from here. Thanks!

284

I verified that we do. I got an error with two allocations if I didnt do this. (Adding a test for that as well).

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

61

True. But this will probably need to make changes to unrelated tests. Would prefer to do this in a later CL.

mlir/test/Conversion/StandardToSPIRV/alloc.mlir

23

THanks. Fixed!

36

Damn. Thanks for catching it. Commented now.
Actually looking at this test again, I realize that if this test worked as expected the generated spv.AtomicOr/spiv.AtomicAnd operations would have the wrong scope. Adjusted the code to handle that, but this cant be tested without the underlying issue being fixed.

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

 /// pointers to structs.
 ///
 /// TODO(antiagainst): We might want to introduce a way to control how
 /// unsupported bitwidth are handled and explicitly fail if wanted.
 class SPIRVTypeConverter : public TypeConverter {
 public:
   explicit SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr);
 
+  /// Gets the number of bytes used for a type when converted to SPIR-V
+  /// type. Note that it doesnt account for whether the type is legal for a
+  /// SPIR-V target (described by spirv::TargetEnvAttr). Returns None on
+  /// failure.
+  static Optional<int64_t> getConvertedTypeNumBytes(Type);
+
   /// 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

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

 }
 
 #define CHECK_UNSIGNED_OP(SPIRVOp)                                             \
   template <>                                                                  \
   bool isUnsignedOp<SPIRVOp>() {                                               \
     return true;                                                               \
   }
 
-CHECK_UNSIGNED_OP(spirv::AtomicUMaxOp);
-CHECK_UNSIGNED_OP(spirv::AtomicUMinOp);
-CHECK_UNSIGNED_OP(spirv::BitFieldUExtractOp);
-CHECK_UNSIGNED_OP(spirv::ConvertUToFOp);
-CHECK_UNSIGNED_OP(spirv::GroupNonUniformUMaxOp);
-CHECK_UNSIGNED_OP(spirv::GroupNonUniformUMinOp);
-CHECK_UNSIGNED_OP(spirv::UConvertOp);
-CHECK_UNSIGNED_OP(spirv::UDivOp);
-CHECK_UNSIGNED_OP(spirv::UGreaterThanEqualOp);
-CHECK_UNSIGNED_OP(spirv::UGreaterThanOp);
-CHECK_UNSIGNED_OP(spirv::ULessThanEqualOp);
-CHECK_UNSIGNED_OP(spirv::ULessThanOp);
-CHECK_UNSIGNED_OP(spirv::UModOp);
+CHECK_UNSIGNED_OP(spirv::AtomicUMaxOp)
+CHECK_UNSIGNED_OP(spirv::AtomicUMinOp)
+CHECK_UNSIGNED_OP(spirv::BitFieldUExtractOp)
+CHECK_UNSIGNED_OP(spirv::ConvertUToFOp)
+CHECK_UNSIGNED_OP(spirv::GroupNonUniformUMaxOp)
+CHECK_UNSIGNED_OP(spirv::GroupNonUniformUMinOp)
+CHECK_UNSIGNED_OP(spirv::UConvertOp)
+CHECK_UNSIGNED_OP(spirv::UDivOp)
+CHECK_UNSIGNED_OP(spirv::UGreaterThanEqualOp)
+CHECK_UNSIGNED_OP(spirv::UGreaterThanOp)
+CHECK_UNSIGNED_OP(spirv::ULessThanEqualOp)
+CHECK_UNSIGNED_OP(spirv::ULessThanOp)
+CHECK_UNSIGNED_OP(spirv::UModOp)
 
 #undef CHECK_UNSIGNED_OP
 
+/// Returns true if the allocations of type `t` can be lowered to SPIR-V.
+static bool isAllocationSupported(MemRefType t) {
+  // Currently only support workgroup local memory allocations with static
+  // shape and int or float element type.
+  return t.hasStaticShape() &&
+         SPIRVTypeConverter::getMemorySpaceForStorageClass(
+             spirv::StorageClass::Workgroup) == t.getMemorySpace() &&
+         t.getElementType().isIntOrFloat();
+}
+
+/// Returns the scope to use for atomic operations use for emulating store
+/// operations of unsupported integer bitwidths, based on the memref
+/// type. Returns None on failure.
+static Optional<spirv::Scope> getAtomicOpScope(MemRefType t) {
+  Optional<spirv::StorageClass> storageClass =
+      SPIRVTypeConverter::getStorageClassForMemorySpace(t.getMemorySpace());
+  if (!storageClass)
+    return {};
+  switch (*storageClass) {
+  case spirv::StorageClass::StorageBuffer:
+    return spirv::Scope::Device;
+  case spirv::StorageClass::Workgroup:
+    return spirv::Scope::Workgroup;
+  default: {
+  }
+  }
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // Operation conversion
 //===----------------------------------------------------------------------===//
 
 // Note that DRR cannot be used for the patterns in this file: we may need to
 // convert type along the way, which requires ConversionPattern. DRR generates
 // normal RewritePattern.
 
 namespace {
 
+/// Converts an allocation operation to SPIR-V. Currently only supports lowering
+/// to Workgroup memory when the size is constant.  Note that this pattern needs
+/// to be applied in a pass that runs at least at spv.module scope since it wil
+/// ladd global variables into the spv.module.
+class AllocOpPattern final : public SPIRVOpLowering<AllocOp> {
+public:
+  using SPIRVOpLowering<AllocOp>::SPIRVOpLowering;
+
+  LogicalResult
+  matchAndRewrite(AllocOp operation, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefType allocType = operation.getType();
+    if (!isAllocationSupported(allocType))
+      return operation.emitError("unhandled allocation type");
+
+    // Get the SPIR-V type for the allocation.
+    Type spirvType = typeConverter.convertType(allocType);
+
+    // Insert spv.globalVariable for this allocation.
+    Operation *parent =
+        SymbolTable::getNearestSymbolTable(operation.getParentOp());
+    if (!parent)
+      return failure();
+    Location loc = operation.getLoc();
+    spirv::GlobalVariableOp varOp;
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      Block &entryBlock = *parent->getRegion(0).begin();
+      rewriter.setInsertionPointToStart(&entryBlock);
+      auto varOps = entryBlock.getOps<spirv::GlobalVariableOp>();
+      std::string varName =
+          std::string("__workgroup_mem__") +
+          std::to_string(std::distance(varOps.begin(), varOps.end()));
+      varOp = rewriter.create<spirv::GlobalVariableOp>(
+          loc, TypeAttr::get(spirvType), varName,
+          /*initializer = */ nullptr);
+    }
+
+    // Get pointer to global variable at the current scope.
+    rewriter.replaceOpWithNewOp<spirv::AddressOfOp>(operation, varOp);
+    return success();
+  }
+};
+
+/// Removed a deallocation if it is a supported allocation. Currently only
+/// removes deallocation if the memory space is workgroup memory.
+class DeallocOpPattern final : public SPIRVOpLowering<DeallocOp> {
+public:
+  using SPIRVOpLowering<DeallocOp>::SPIRVOpLowering;
+
+  LogicalResult
+  matchAndRewrite(DeallocOp operation, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefType deallocType = operation.memref().getType().cast<MemRefType>();
+    if (!isAllocationSupported(deallocType))
+      return operation.emitError("unhandled deallocation type");
+    rewriter.eraseOp(operation);
+    return success();
+  }
+};
+
 /// Converts unary and binary standard operations to SPIR-V operations.
 template <typename StdOp, typename SPIRVOp>
 class UnaryAndBinaryOpPattern final : public SPIRVOpLowering<StdOp> {
 public:
   using SPIRVOpLowering<StdOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(StdOp operation, ArrayRef<Value> operands,
IntStoreOpPattern::matchAndRewrite(StoreOp storeOp, ArrayRef<Value> operands,
   Value clearBitsMask =
       rewriter.create<spirv::ShiftLeftLogicalOp>(loc, dstType, mask, offset);
   clearBitsMask = rewriter.create<spirv::NotOp>(loc, dstType, clearBitsMask);
 
   Value storeVal =
       shiftValue(loc, storeOperands.value(), offset, mask, dstBits, rewriter);
   Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
                                                    srcBits, dstBits, rewriter);
+  Optional<spirv::Scope> scope = getAtomicOpScope(memrefType);
+  if (!scope)
+    return failure();
   Value result = rewriter.create<spirv::AtomicAndOp>(
-      loc, dstType, adjustedPtr, spirv::Scope::Device,
-      spirv::MemorySemantics::AcquireRelease, clearBitsMask);
+      loc, dstType, adjustedPtr, *scope, spirv::MemorySemantics::AcquireRelease,
+      clearBitsMask);
   result = rewriter.create<spirv::AtomicOrOp>(
-      loc, dstType, adjustedPtr, spirv::Scope::Device,
-      spirv::MemorySemantics::AcquireRelease, storeVal);
+      loc, dstType, adjustedPtr, *scope, spirv::MemorySemantics::AcquireRelease,
+      storeVal);
 
   // The AtomicOrOp has no side effect. Since it is already inserted, we can
   // just remove the original StoreOp. Note that rewriter.replaceOp()
   // doesn't work because it only accepts that the numbers of result are the
   // same.
   rewriter.eraseOp(storeOp);
 
   assert(accessChainOp.use_empty());
patterns.insert<
       UnaryAndBinaryOpPattern<SinOp, spirv::GLSLSinOp>,
       UnaryAndBinaryOpPattern<SqrtOp, spirv::GLSLSqrtOp>,
       UnaryAndBinaryOpPattern<SubFOp, spirv::FSubOp>,
       UnaryAndBinaryOpPattern<SubIOp, spirv::ISubOp>,
       UnaryAndBinaryOpPattern<TanhOp, spirv::GLSLTanhOp>,
       UnaryAndBinaryOpPattern<UnsignedDivIOp, spirv::UDivOp>,
       UnaryAndBinaryOpPattern<UnsignedRemIOp, spirv::UModOp>,
       UnaryAndBinaryOpPattern<UnsignedShiftRightOp, spirv::ShiftRightLogicalOp>,
+      AllocOpPattern, DeallocOpPattern,
       BitwiseOpPattern<AndOp, spirv::LogicalAndOp, spirv::BitwiseAndOp>,
       BitwiseOpPattern<OrOp, spirv::LogicalOrOp, spirv::BitwiseOrOp>,
       BoolCmpIOpPattern, ConstantCompositeOpPattern, ConstantScalarOpPattern,
       CmpFOpPattern, CmpIOpPattern, IntLoadOpPattern, LoadOpPattern,
       ReturnOpPattern, SelectOpPattern, IntStoreOpPattern, StoreOpPattern,
       TypeCastingOpPattern<IndexCastOp, spirv::SConvertOp>,
       TypeCastingOpPattern<SIToFPOp, spirv::ConvertSToFOp>,
       TypeCastingOpPattern<ZeroExtendIOp, spirv::UConvertOp>,
       TypeCastingOpPattern<TruncateIOp, spirv::SConvertOp>,
       TypeCastingOpPattern<FPToSIOp, spirv::ConvertFToSOp>,
       TypeCastingOpPattern<FPExtOp, spirv::FConvertOp>,
       TypeCastingOpPattern<FPTruncOp, spirv::FConvertOp>, XOrOpPattern>(
       context, typeConverter);
 }
 } // namespace mlir

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

for (auto shape : tensorType.getShape()) {
       size *= shape;
     }
     return size;
   }
   // TODO: Add size computation for other types.
   return llvm::None;
 }
 
+Optional<int64_t> SPIRVTypeConverter::getConvertedTypeNumBytes(Type t) {
+  return getTypeNumBytes(t);
+}
+
 /// Converts a scalar `type` to a suitable type under the given `targetEnv`.
 static Optional<Type>
 convertScalarType(const spirv::TargetEnv &targetEnv, spirv::ScalarType type,
                   Optional<spirv::StorageClass> storageClass = {}) {
   // Get extension and capability requirements for the given type.
   SmallVector<ArrayRef<spirv::Extension>, 1> extensions;
   SmallVector<ArrayRef<spirv::Capability>, 2> capabilities;
   type.getExtensions(extensions, storageClass);
if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return llvm::None;
   }
 
   auto arrayType =
       spirv::ArrayType::get(*arrayElemType, arrayElemCount, *arrayElemSize);
 
-  // Wrap in a struct to satisfy Vulkan interface requirements.
-  auto structType = spirv::StructType::get(arrayType, 0);
+  // Wrap in a struct to satisfy Vulkan interface requirements. Memrefs with
+  // workgroup storage class do not need the struct to be laid out explicitly.
+  auto structType = *storageClass == spirv::StorageClass::Workgroup
+                        ? spirv::StructType::get(arrayType)
+                        : spirv::StructType::get(arrayType, 0);
   return spirv::PointerType::get(structType, *storageClass);
 }
 
 SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr)
     : targetEnv(targetAttr) {
   // Add conversions. The order matters here: later ones will be tried earlier.
 
   // All other cases failed. Then we cannot convert this type.

 //===----------------------------------------------------------------------===//
 // Index calculation
 //===----------------------------------------------------------------------===//
 
 spirv::AccessChainOp mlir::spirv::getElementPtr(
     SPIRVTypeConverter &typeConverter, MemRefType baseType, Value basePtr,
     ArrayRef<Value> 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())) {
+      llvm::is_contained(strides, MemRefType::getDynamicStrideOrOffset()) ||
+      offset == MemRefType::getDynamicStrideOrOffset()) {
     return nullptr;
   }
 
   auto indexType = typeConverter.getIndexType(builder.getContext());
+  SmallVector<Value, 2> linearizedIndices;
+  // Add a '0' at the start to index into the struct.
+  auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
+  linearizedIndices.push_back(zero);
 
-  Value ptrLoc = nullptr;
+  if (baseType.getRank() == 0) {
+    linearizedIndices.push_back(zero);
+  } else {
+    // TODO: Instead of this logic, use affine.apply and add patterns for
+    // lowering affine.apply to standard ops. These will get lowered to SPIR-V
+    // ops by the DialectConversion framework.
+    Value ptrLoc = builder.create<spirv::ConstantOp>(
+        loc, indexType, IntegerAttr::get(indexType, offset));
     assert(indices.size() == strides.size() &&
            "must provide indices for all dimensions");
     for (auto index : enumerate(indices)) {
       Value strideVal = builder.create<spirv::ConstantOp>(
           loc, indexType, IntegerAttr::get(indexType, strides[index.index()]));
-    Value update = builder.create<spirv::IMulOp>(loc, strideVal, index.value());
-    ptrLoc =
-        (ptrLoc ? builder.create<spirv::IAddOp>(loc, ptrLoc, update).getResult()
-                : update);
-  }
-  SmallVector<Value, 2> linearizedIndices;
-  // Add a '0' at the start to index into the struct.
-  auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
-  linearizedIndices.push_back(zero);
-  // If it is a zero-rank memref type, extract the element directly.
-  if (!ptrLoc) {
-    ptrLoc = zero;
+      Value update =
+          builder.create<spirv::IMulOp>(loc, strideVal, index.value());
+      ptrLoc = builder.create<spirv::IAddOp>(loc, ptrLoc, update);
     }
     linearizedIndices.push_back(ptrLoc);
+  }
   return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
 }
 
 //===----------------------------------------------------------------------===//
 // Set ABI attributes for lowering entry functions.
 //===----------------------------------------------------------------------===//
 
 LogicalResult

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

gpu.func @load_store_kernel(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>, %arg3: index, %arg4: index, %arg5: index, %arg6: index) kernel
       %8 = "gpu.grid_dim"() {dimension = "z"} : () -> index
       %9 = "gpu.block_dim"() {dimension = "x"} : () -> index
       %10 = "gpu.block_dim"() {dimension = "y"} : () -> index
       %11 = "gpu.block_dim"() {dimension = "z"} : () -> index
       // CHECK: %[[INDEX1:.*]] = spv.IAdd %[[ARG3]], %[[WORKGROUPIDX]]
       %12 = addi %arg3, %0 : index
       // CHECK: %[[INDEX2:.*]] = spv.IAdd %[[ARG4]], %[[LOCALINVOCATIONIDX]]
       %13 = addi %arg4, %3 : index
+      // CHECK: %[[ZERO:.*]] = spv.constant 0 : i32
+      // CHECK: %[[OFFSET1_0:.*]] = spv.constant 0 : i32
       // CHECK: %[[STRIDE1_1:.*]] = spv.constant 4 : i32
-      // CHECK: %[[OFFSET1_1:.*]] = spv.IMul %[[STRIDE1_1]], %[[INDEX1]] : i32
+      // CHECK: %[[UPDATE1_1:.*]] = spv.IMul %[[STRIDE1_1]], %[[INDEX1]] : i32
+      // CHECK: %[[OFFSET1_1:.*]] = spv.IAdd %[[OFFSET1_0]], %[[UPDATE1_1]] : i32
       // CHECK: %[[STRIDE1_2:.*]] = spv.constant 1 : i32
       // CHECK: %[[UPDATE1_2:.*]] = spv.IMul %[[STRIDE1_2]], %[[INDEX2]] : i32
       // CHECK: %[[OFFSET1_2:.*]] = spv.IAdd %[[OFFSET1_1]], %[[UPDATE1_2]] : i32
-      // CHECK: %[[ZERO1:.*]] = spv.constant 0 : i32
-      // CHECK: %[[PTR1:.*]] = spv.AccessChain %[[ARG0]]{{\[}}%[[ZERO1]], %[[OFFSET1_2]]{{\]}}
+      // CHECK: %[[PTR1:.*]] = spv.AccessChain %[[ARG0]]{{\[}}%[[ZERO]], %[[OFFSET1_2]]{{\]}}
       // CHECK-NEXT: %[[VAL1:.*]] = spv.Load "StorageBuffer" %[[PTR1]]
       %14 = load %arg0[%12, %13] : memref<12x4xf32>
       // CHECK: %[[PTR2:.*]] = spv.AccessChain %[[ARG1]]{{\[}}{{%.*}}, {{%.*}}{{\]}}
       // CHECK-NEXT: %[[VAL2:.*]] = spv.Load "StorageBuffer" %[[PTR2]]
       %15 = load %arg1[%12, %13] : memref<12x4xf32>
       // CHECK: %[[VAL3:.*]] = spv.FAdd %[[VAL1]], %[[VAL2]]
       %16 = addf %14, %15 : f32
       // CHECK: %[[PTR3:.*]] = spv.AccessChain %[[ARG2]]{{\[}}{{%.*}}, {{%.*}}{{\]}}
       // CHECK-NEXT: spv.Store "StorageBuffer" %[[PTR3]], %[[VAL3]]
       store %16, %arg2[%12, %13] : memref<12x4xf32>
       gpu.return
     }
   }
 }

mlir/test/Conversion/GPUToSPIRV/loop.mlir

attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       // CHECK: %[[STEP:.*]] = spv.constant 2 : i32
       %step = constant 2 : index
       // CHECK:      spv.loop {
       // CHECK-NEXT:   spv.Branch ^[[HEADER:.*]](%[[LB]] : i32)
       // CHECK:      ^[[HEADER]](%[[INDVAR:.*]]: i32):
       // CHECK:        %[[CMP:.*]] = spv.SLessThan %[[INDVAR]], %[[UB]] : i32
       // CHECK:        spv.BranchConditional %[[CMP]], ^[[BODY:.*]], ^[[MERGE:.*]]
       // CHECK:      ^[[BODY]]:
-      // CHECK:        %[[STRIDE1:.*]] = spv.constant 1 : i32
-      // CHECK:        %[[INDEX1:.*]] = spv.IMul %[[STRIDE1]], %[[INDVAR]] : i32
       // CHECK:        %[[ZERO1:.*]] = spv.constant 0 : i32
+      // CHECK:        %[[OFFSET1:.*]] = spv.constant 0 : i32
+      // CHECK:        %[[STRIDE1:.*]] = spv.constant 1 : i32
+      // CHECK:        %[[UPDATE1:.*]] = spv.IMul %[[STRIDE1]], %[[INDVAR]] : i32
+      // CHECK:        %[[INDEX1:.*]] = spv.IAdd %[[OFFSET1]], %[[UPDATE1]] : i32
       // CHECK:        spv.AccessChain {{%.*}}{{\[}}%[[ZERO1]], %[[INDEX1]]{{\]}}
-      // CHECK:        %[[STRIDE2:.*]] = spv.constant 1 : i32
-      // CHECK:        %[[INDEX2:.*]] = spv.IMul %[[STRIDE2]], %[[INDVAR]] : i32
       // CHECK:        %[[ZERO2:.*]] = spv.constant 0 : i32
+      // CHECK:        %[[OFFSET2:.*]] = spv.constant 0 : i32
+      // CHECK:        %[[STRIDE2:.*]] = spv.constant 1 : i32
+      // CHECK:        %[[UPDATE2:.*]] = spv.IMul %[[STRIDE2]], %[[INDVAR]] : i32
+      // CHECK:        %[[INDEX2:.*]] = spv.IAdd %[[OFFSET2]], %[[UPDATE2]] : i32
       // CHECK:        spv.AccessChain {{%.*}}[%[[ZERO2]], %[[INDEX2]]]
       // CHECK:        %[[INCREMENT:.*]] = spv.IAdd %[[INDVAR]], %[[STEP]] : i32
       // CHECK:        spv.Branch ^[[HEADER]](%[[INCREMENT]] : i32)
       // CHECK:      ^[[MERGE]]
       // CHECK:        spv._merge
       // CHECK:      }
       scf.for %arg4 = %lb to %ub step %step {
         %1 = load %arg2[%arg4] : memref<10xf32>
         store %1, %arg3[%arg4] : memref<10xf32>
       }
       gpu.return
     }
   }
 }

mlir/test/Conversion/StandardToSPIRV/alloc.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -convert-std-to-spirv -canonicalize -verify-diagnostics %s -o - | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// std allocation/deallocation ops
+//===----------------------------------------------------------------------===//
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+  }
+{
+  func @alloc_dealloc_workgroup_mem(%arg0 : index, %arg1 : index) {
+    %0 = alloc() : memref<4x5xf32, 3>
+    %1 = load %0[%arg0, %arg1] : memref<4x5xf32, 3>
+    store %1, %0[%arg0, %arg1] : memref<4x5xf32, 3>
+    dealloc %0 : memref<4x5xf32, 3>
+    return
+  }
+}
+//     CHECK: spv.globalVariable @[[VAR:.+]] : !spv.ptr<!spv.struct<!spv.array<20 x f32, stride=4>>, Workgroup>
+//     CHECK: func @alloc_dealloc_workgroup_mem
+// CHECK-NOT:   alloc
+//     CHECK:   %[[PTR:.+]] = spv._address_of @[[VAR]]
+//     CHECK:   %[[LOADPTR:.+]] = spv.AccessChain %[[PTR]]
+//     CHECK:   %[[VAL:.+]] = spv.Load "Workgroup" %[[LOADPTR]] : f32
+//     CHECK:   %[[STOREPTR:.+]] = spv.AccessChain %[[PTR]]
+//     CHECK:   spv.Store "Workgroup" %[[STOREPTR]], %[[VAL]] : f32
+// CHECK-NOT:   dealloc
+//     CHECK:   spv.Return
+
+// -----
+
+// TODO: Uncomment this test when the extension handling correctly
+// converts an i16 type to i32 type and handles the load/stores
+// correctly.
+
+// module attributes {
+//   spv.target_env = #spv.target_env<
+//     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+//     {max_compute_workgroup_invocations = 128 : i32,
+//      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+//   }
+// {
+//   func @alloc_dealloc_workgroup_mem(%arg0 : index, %arg1 : index) {
+//     %0 = alloc() : memref<4x5xi16, 3>
+//     %1 = load %0[%arg0, %arg1] : memref<4x5xi16, 3>
+//     store %1, %0[%arg0, %arg1] : memref<4x5xi16, 3>
+//     dealloc %0 : memref<4x5xi16, 3>
+//     return
+//   }
+// }
+
+// -----
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+  }
+{
+  func @two_allocs() {
+    %0 = alloc() : memref<4x5xf32, 3>
+    %1 = alloc() : memref<2x3xi32, 3>
+    return
+  }
+}
+
+//  CHECK-DAG: spv.globalVariable @__workgroup_mem__{{[0-9]+}}
+// CHECK-SAME:   !spv.ptr<!spv.struct<!spv.array<6 x i32, stride=4>>, Workgroup>
+//  CHECK-DAG: spv.globalVariable @__workgroup_mem__{{[0-9]+}}
+// CHECK-SAME:   !spv.ptr<!spv.struct<!spv.array<20 x f32, stride=4>>, Workgroup>
+//      CHECK: spv.func @two_allocs()
+//      CHECK: spv.Return
+
+// -----
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+  }
+{
+  func @alloc_dealloc_dynamic_workgroup_mem(%arg0 : index) {
+    // expected-error @+2 {{unhandled allocation type}}
+    // expected-error @+1 {{'std.alloc' op operand #0 must be index}}
+    %0 = alloc(%arg0) : memref<4x?xf32, 3>
+    return
+  }
+}
+
+// -----
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+  }
+{
+  func @alloc_dealloc_mem() {
+    // expected-error @+1 {{unhandled allocation type}}
+    %0 = alloc() : memref<4x5xf32>
+    return
+  }
+}
+
+
+// -----
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+  }
+{
+  func @alloc_dealloc_dynamic_workgroup_mem(%arg0 : memref<4x?xf32, 3>) {
+    // expected-error @+2 {{unhandled deallocation type}}
+    // expected-error @+1 {{'std.dealloc' op operand #0 must be memref of any type values}}
+    dealloc %arg0 : memref<4x?xf32, 3>
+    return
+  }
+}
+
+// -----
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+  }
+{
+  func @alloc_dealloc_mem(%arg0 : memref<4x5xf32>) {
+    // expected-error @+2 {{unhandled deallocation type}}
+    // expected-error @+1 {{op operand #0 must be memref of any type values}}
+    dealloc %arg0 : memref<4x5xf32>
+    return
+  }
+}

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

func @load_i8(%arg0: memref<i8>) {
   //     CHECK: spv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
   %0 = load %arg0[] : memref<i8>
   return
 }
 
 // CHECK-LABEL: @load_i16
 //       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i32)
 func @load_i16(%arg0: memref<10xi16>, %index : index) {
-  //     CHECK: %[[ONE:.+]] = spv.constant 1 : i32
-  //     CHECK: %[[FLAT_IDX:.+]] = spv.IMul %[[ONE]], %[[ARG1]] : i32
   //     CHECK: %[[ZERO:.+]] = spv.constant 0 : i32
+  //     CHECK: %[[OFFSET:.+]] = spv.constant 0 : i32
+  //     CHECK: %[[ONE:.+]] = spv.constant 1 : i32
+  //     CHECK: %[[UPDATE:.+]] = spv.IMul %[[ONE]], %[[ARG1]] : i32
+  //     CHECK: %[[FLAT_IDX:.+]] = spv.IAdd %[[OFFSET]], %[[UPDATE]] : i32
   //     CHECK: %[[TWO1:.+]] = spv.constant 2 : i32
   //     CHECK: %[[QUOTIENT:.+]] = spv.SDiv %[[FLAT_IDX]], %[[TWO1]] : i32
   //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
   //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
   //     CHECK: %[[TWO2:.+]] = spv.constant 2 : i32
   //     CHECK: %[[SIXTEEN:.+]] = spv.constant 16 : i32
   //     CHECK: %[[IDX:.+]] = spv.SMod %[[FLAT_IDX]], %[[TWO2]] : i32
   //     CHECK: %[[BITS:.+]] = spv.IMul %[[IDX]], %[[SIXTEEN]] : i32
func @store_i8(%arg0: memref<i8>, %value: i8) {
   //     CHECK: spv.AtomicOr "Device" "AcquireRelease" %[[PTR]], %[[STORE_VAL]]
   store %value, %arg0[] : memref<i8>
   return
 }
 
 // CHECK-LABEL: @store_i16
 //       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32)
 func @store_i16(%arg0: memref<10xi16>, %index: index, %value: i16) {
-  //     CHECK: %[[ONE:.+]] = spv.constant 1 : i32
-  //     CHECK: %[[FLAT_IDX:.+]] = spv.IMul %[[ONE]], %[[ARG1]] : i32
   //     CHECK: %[[ZERO:.+]] = spv.constant 0 : i32
+  //     CHECK: %[[OFFSET:.+]] = spv.constant 0 : i32
+  //     CHECK: %[[ONE:.+]] = spv.constant 1 : i32
+  //     CHECK: %[[UPDATE:.+]] = spv.IMul %[[ONE]], %[[ARG1]] : i32
+  //     CHECK: %[[FLAT_IDX:.+]] = spv.IAdd %[[OFFSET]], %[[UPDATE]] : i32
   //     CHECK: %[[TWO:.+]] = spv.constant 2 : i32
   //     CHECK: %[[SIXTEEN:.+]] = spv.constant 16 : i32
   //     CHECK: %[[IDX:.+]] = spv.SMod %[[FLAT_IDX]], %[[TWO]] : i32
   //     CHECK: %[[OFFSET:.+]] = spv.IMul %[[IDX]], %[[SIXTEEN]] : i32
   //     CHECK: %[[MASK1:.+]] = spv.constant 65535 : i32
   //     CHECK: %[[TMP1:.+]] = spv.ShiftLeftLogical %[[MASK1]], %[[OFFSET]] : i32, i32
   //     CHECK: %[[MASK:.+]] = spv.Not %[[TMP1]] : i32
   //     CHECK: %[[CLAMPED_VAL:.+]] = spv.BitwiseAnd %[[ARG2]], %[[MASK1]] : i32