[mlir][spirv] Memref to SPIRV lowering for OCL

[raikonenfnu]

Summary

-Add awareness to Kernel vs Shader capability for memref to spv lowering.
-Add lowering using spv.PtrAccessChain
-Enable lowering from scalar pointee types for kernel capabilities.

[antiagainst]

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

67

This should not be in Options. Options allows a user to control how the type converter works. Whether we have kernel capability is not something a user can control; it is a characteristic of the target environment.

We should actually expose a getter method in the SPIRVTypeConverter to return the targetEnv to allow calling its allows(spirv::Capabilty) method. This is also generic enough to avoid special casing Kernel capabilities.

172

Hmm, I think we can fold the different code path selection for Vulkan and OpenCL inside the existing entry point getElementPtr. Under the hood we call into getVulkanElementPtr (which is the existing getElementPtr impl) and getOpenCLElementPtr, depending on whether allows(spirv::Capability::Kernel).

The benefits of doing this is that we don't need duplicating all the patterns for Vulkan and OpenCL. (I know I said it's preferable to have separate parallel patterns; but here we are seeing too much similar code duplication across many patterns. Unifying the different part behind a common helper function is better.)

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

581

This is not the proper way to add patterns. Patterns should be safe to be _populated_ no matter what target environment we have. See my previous comments regarding how to have the kernel capability probing and selection logic inside getElementPtr.

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

341

This should query spirv::TargetEnv::allows(spirv::Capability::Kernel).

Also please put some comments before saying " For OpenCL Kernel we can just emit a pointer pointing to the element." and after saying " For Vulkan we need extra wrapping struct and array to satisfy interface needs."

406

Similarly here.

435

This is not the proper way. Options are for user controls. It is different from target environment.

[raikonenfnu]

Refactored code to address issues.

[raikonenfnu]

minor clean up

[raikonenfnu]

another minor clean up

[nicolasvasilache]

how does this intercept the ongoing work with extract_strided_metadata ?

[raikonenfnu]

In D132714#3794664, @nicolasvasilache wrote:

how does this intercept the ongoing work with extract_strided_metadata ?

Hey Nicolas, I am not too familiar with the ongoing work with extract_strided_metadata, however, this current patch is just to check if the spirv is of kernel flavor and introducing the usage and correct accessing (spv.ptrAccessChain) of spv.ptr<ElemType> rather than (spv.AccessChain) + spv.ptr<spv.struct<spv.array<elementType>>> which is needed for vulkan/shader type.

Since the other logic is not really touched I assume the works being done in extract_strided_metadata that's made to fit the current pipeline(vulkan/shader only) can also be used by (ocl/kernel) after this patch lands.

[nicolasvasilache]

In D132714#3796390, @raikonenfnu wrote:

In D132714#3794664, @nicolasvasilache wrote:

how does this intercept the ongoing work with extract_strided_metadata ?

Hey Nicolas, I am not too familiar with the ongoing work with extract_strided_metadata, however, this current patch is just to check if the spirv is of kernel flavor and introducing the usage and correct accessing (spv.ptrAccessChain) of spv.ptr<ElemType> rather than (spv.AccessChain) + spv.ptr<spv.struct<spv.array<elementType>>> which is needed for vulkan/shader type.

Since the other logic is not really touched I assume the works being done in extract_strided_metadata that's made to fit the current pipeline(vulkan/shader only) can also be used by (ocl/kernel) after this patch lands.

SGTM, I just wanted to attract your collective attention on that line of work that may be useful for SPIRV too.
I do not have deep SPIRV knowledge myself but the descriptor being only unpacked at the level of LLVM was/is a common problem that we are addressing actively.
The logic that creates spv.IMul spv.IAdd spv.PtrAccessChain and friends should also benefit from the new patterns we are adding (cc @qcolombet ).

[antiagainst]

In D132714#3799206, @nicolasvasilache wrote:

In D132714#3796390, @raikonenfnu wrote:

In D132714#3794664, @nicolasvasilache wrote:

how does this intercept the ongoing work with extract_strided_metadata ?

Hey Nicolas, I am not too familiar with the ongoing work with extract_strided_metadata, however, this current patch is just to check if the spirv is of kernel flavor and introducing the usage and correct accessing (spv.ptrAccessChain) of spv.ptr<ElemType> rather than (spv.AccessChain) + spv.ptr<spv.struct<spv.array<elementType>>> which is needed for vulkan/shader type.

Since the other logic is not really touched I assume the works being done in extract_strided_metadata that's made to fit the current pipeline(vulkan/shader only) can also be used by (ocl/kernel) after this patch lands.

SGTM, I just wanted to attract your collective attention on that line of work that may be useful for SPIRV too.
I do not have deep SPIRV knowledge myself but the descriptor being only unpacked at the level of LLVM was/is a common problem that we are addressing actively.
The logic that creates spv.IMul spv.IAdd spv.PtrAccessChain and friends should also benefit from the new patterns we are adding (cc @qcolombet ).

Thanks for the heads up, @nicolasvasilache ! Yup, memref.extract_strided_metadata would certainly help the spirv path! Unlike llvm, we don't have the flexibility to always represent memref descriptors as structs inside spirv (spirv is used for various client APIs with different capabilities; not all client API or target environment supports such structs. Notably vulkan by default does not.) So on the spirv path we "linearize" the n-D memref descriptor first before converting to spirv. With memref.extract_strided_metadata that would be easier given now we can use that op first and relying on simple arith op lowering afterwards, instead of baking the logic under converting memref.load/memref.store!

I'll convert spirv side to use it later afterwards, which should handle both vulkan and opencl side. (converting to OpAccessChain or OpPtrAccessChain for vulkan/opencl is later to that; so should be orthogonal.)

[antiagainst]

Cool, overall LGTM; just a few remaining nits. Thanks! :)

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

162

Nit: In comment please write as "SPIR-V", as that's the official spelling of it. (In code we cannot because of the -, but in comment we don't subject to this restriction.)

167

Same here.

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

322

Just use Value here to be clear.

362

... unsupported for Kernel capability / spv.PtrAccessChain.

363

Style: no need to use { ... } for trivial if conditions.

367

We can write as accessChain.getDefiningOp<spirv::AccessChainOp>().

449

Value here.

490

Same here.

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

342

No {...} here.

407

Same here.

[raikonenfnu]

fix styling nits

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

 //===- SPIRVConversion.h - SPIR-V Conversion Utilities ----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Defines utilities to use while converting to the SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_TRANSFORMS_SPIRVCONVERSION_H
 #define MLIR_DIALECT_SPIRV_TRANSFORMS_SPIRVCONVERSION_H
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SmallSet.h"
 
 namespace mlir {
 
 //===----------------------------------------------------------------------===//

 ///
 /// For memref types, this converter additionally performs type wrapping to
 /// satisfy shader interface requirements: shader interface types must be
 /// pointers to structs.
 class SPIRVTypeConverter : public TypeConverter {
 public:
   explicit SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
                               const SPIRVConversionOptions &options = {});
 
   /// Gets the SPIR-V correspondence for the standard index type.
   Type getIndexType() const;
 
   const spirv::TargetEnv &getTargetEnv() const { return targetEnv; }
 
   /// Returns the options controlling the SPIR-V type converter.
   const SPIRVConversionOptions &getOptions() const { return options; }
 
+  /// Checks if the SPIR-V capability inquired is supported.
+  bool allows(spirv::Capability capability);
+
 private:
   spirv::TargetEnv targetEnv;
   SPIRVConversionOptions options;
 
   MLIRContext *getContext() const;
 };
 
 //===----------------------------------------------------------------------===//
Value linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
                      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);
+Value getElementPtr(SPIRVTypeConverter &typeConverter, MemRefType baseType,
+                    Value basePtr, ValueRange indices, Location loc,
 
+// GetElementPtr implementation for Kernel/OpenCL flavored SPIR-V.
+Value getOpenCLElementPtr(SPIRVTypeConverter &typeConverter,
+                          MemRefType baseType, Value basePtr,
+                          ValueRange indices, Location loc, OpBuilder &builder);
+
+// GetElementPtr implementation for Vulkan/Shader flavored SPIR-V.
+Value getVulkanElementPtr(SPIRVTypeConverter &typeConverter,
+                          MemRefType baseType, Value basePtr,
+                          ValueRange indices, Location loc, OpBuilder &builder);
+
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_TRANSFORMS_SPIRVCONVERSION_H

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 public:
   using OpConversionPattern<memref::DeallocOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::DeallocOp operation, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
-/// Converts memref.load to spv.Load.
+/// Converts memref.load to spv.Load + spv.AccessChain on integers.
 class IntLoadOpPattern final : public OpConversionPattern<memref::LoadOp> {
 public:
   using OpConversionPattern<memref::LoadOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
-/// Converts memref.load to spv.Load.
+/// Converts memref.load to spv.Load + spv.AccessChain.
 class LoadOpPattern final : public OpConversionPattern<memref::LoadOp> {
 public:
   using OpConversionPattern<memref::LoadOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };

 IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
   auto loc = loadOp.getLoc();
   auto memrefType = loadOp.getMemref().getType().cast<MemRefType>();
   if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
-  spirv::AccessChainOp accessChainOp =
+  Value accessChain =
       spirv::getElementPtr(typeConverter, memrefType, adaptor.getMemref(),
                            adaptor.getIndices(), loc, rewriter);
 
-  if (!accessChainOp)
+  if (!accessChain)
     return failure();
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
   Type pointeeType = typeConverter.convertType(memrefType)
                          .cast<spirv::PointerType>()
                          .getPointeeType();
-  Type structElemType = pointeeType.cast<spirv::StructType>().getElementType(0);
   Type dstType;
+  if (typeConverter.allows(spirv::Capability::Kernel)) {
+    // For OpenCL Kernel, pointer will be directly pointing to the element.
+    dstType = pointeeType;
+  } else {
+    // For Vulkan we need to extract element from wrapping struct and array.
+    Type structElemType =
+        pointeeType.cast<spirv::StructType>().getElementType(0);
     if (auto arrayType = structElemType.dyn_cast<spirv::ArrayType>())
       dstType = arrayType.getElementType();
     else
       dstType = structElemType.cast<spirv::RuntimeArrayType>().getElementType();
-
+  }
   int dstBits = dstType.getIntOrFloatBitWidth();
   assert(dstBits % srcBits == 0);
 
   // If the rewrited load op has the same bit width, use the loading value
   // directly.
   if (srcBits == dstBits) {
-    Value loadVal =
-        rewriter.create<spirv::LoadOp>(loc, accessChainOp.getResult());
+    Value loadVal = rewriter.create<spirv::LoadOp>(loc, accessChain);
     if (isBool)
       loadVal = castIntNToBool(loc, loadVal, rewriter);
     rewriter.replaceOp(loadOp, loadVal);
     return success();
   }
 
+  // Bitcasting is currently unsupported for Kernel capability /
+  // spv.PtrAccessChain.
+  if (typeConverter.allows(spirv::Capability::Kernel))
+    return failure();
+
+  auto accessChainOp = accessChain.getDefiningOp<spirv::AccessChainOp>();
+  if (!accessChainOp)
+    return failure();
+
   // Assume that getElementPtr() works linearizely. If it's a scalar, the method
   // still returns a linearized accessing. If the accessing is not linearized,
   // there will be offset issues.
   assert(accessChainOp.indices().size() == 2);
   Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
                                                    srcBits, dstBits, rewriter);
   Value spvLoadOp = rewriter.create<spirv::LoadOp>(
       loc, dstType, adjustedPtr,

 IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                                    ConversionPatternRewriter &rewriter) const {
   auto memrefType = storeOp.getMemref().getType().cast<MemRefType>();
   if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto loc = storeOp.getLoc();
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
-  spirv::AccessChainOp accessChainOp =
+  Value accessChain =
       spirv::getElementPtr(typeConverter, memrefType, adaptor.getMemref(),
                            adaptor.getIndices(), loc, rewriter);
 
-  if (!accessChainOp)
+  if (!accessChain)
     return failure();
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
 
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
 
   Type pointeeType = typeConverter.convertType(memrefType)
                          .cast<spirv::PointerType>()
                          .getPointeeType();
-  Type structElemType = pointeeType.cast<spirv::StructType>().getElementType(0);
   Type dstType;
+  if (typeConverter.allows(spirv::Capability::Kernel)) {
+    // For OpenCL Kernel, pointer will be directly pointing to the element.
+    dstType = pointeeType;
+  } else {
+    // For Vulkan we need to extract element from wrapping struct and array.
+    Type structElemType =
+        pointeeType.cast<spirv::StructType>().getElementType(0);
     if (auto arrayType = structElemType.dyn_cast<spirv::ArrayType>())
       dstType = arrayType.getElementType();
     else
       dstType = structElemType.cast<spirv::RuntimeArrayType>().getElementType();
+  }
 
   int dstBits = dstType.getIntOrFloatBitWidth();
   assert(dstBits % srcBits == 0);
 
   if (srcBits == dstBits) {
     Value storeVal = adaptor.getValue();
     if (isBool)
       storeVal = castBoolToIntN(loc, storeVal, dstType, rewriter);
-    rewriter.replaceOpWithNewOp<spirv::StoreOp>(
-        storeOp, accessChainOp.getResult(), storeVal);
+    rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, accessChain, storeVal);
     return success();
   }
 
+  // Bitcasting is currently unsupported for Kernel capability /
+  // spv.PtrAccessChain.
+  if (typeConverter.allows(spirv::Capability::Kernel))
+    return failure();
+
+  auto accessChainOp = accessChain.getDefiningOp<spirv::AccessChainOp>();
+  if (!accessChainOp)
+    return failure();
+
   // Since there are multi threads in the processing, the emulation will be done
   // with atomic operations. E.g., if the storing value is i8, rewrite the
   // StoreOp to
   // 1) load a 32-bit integer
   // 2) clear 8 bits in the loading value
   // 3) store 32-bit value back
   // 4) load a 32-bit integer
   // 5) modify 8 bits in the loading value

 namespace mlir {
 void populateMemRefToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
   patterns
       .add<AllocaOpPattern, AllocOpPattern, DeallocOpPattern, IntLoadOpPattern,
            IntStoreOpPattern, LoadOpPattern, StoreOpPattern>(
           typeConverter, patterns.getContext());
 }
 } // namespace mlir

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

 Type SPIRVTypeConverter::getIndexType() const {
   return IntegerType::get(getContext(), options.use64bitIndex ? 64 : 32);
 }
 
 MLIRContext *SPIRVTypeConverter::getContext() const {
   return targetEnv.getAttr().getContext();
 }
 
+bool SPIRVTypeConverter::allows(spirv::Capability capability) {
+  return targetEnv.allows(capability);
+}
+
 // 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 SPIRVConversionOptions &options,
                                          Type type) {
   if (type.isa<spirv::ScalarType>()) {
     auto bitWidth = type.getIntOrFloatBitWidth();
     // According to the SPIR-V spec:
     // "There is no physical size or bit pattern defined for values with boolean
if (!arrayElemType)
     return nullptr;
   Optional<int64_t> arrayElemSize = getTypeNumBytes(options, arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return nullptr;
   }
 
+  // For OpenCL Kernel we can just emit a pointer pointing to the element.
+  if (targetEnv.allows(spirv::Capability::Kernel))
+    return spirv::PointerType::get(arrayElemType, storageClass);
+
+  // For Vulkan we need extra wrapping struct and array to satisfy interface
+  // needs.
   if (!type.hasStaticShape()) {
     int64_t stride = needsExplicitLayout(storageClass) ? *arrayElemSize : 0;
     auto arrayType = spirv::RuntimeArrayType::get(arrayElemType, stride);
     return wrapInStructAndGetPointer(arrayType, storageClass);
   }
 
   int64_t memrefSize = (type.getNumElements() * numBoolBits + 7) / 8;
   auto arrayElemCount = llvm::divideCeil(memrefSize, *arrayElemSize);
static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
 
   Optional<int64_t> arrayElemSize = getTypeNumBytes(options, arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return nullptr;
   }
 
+  // For OpenCL Kernel we can just emit a pointer pointing to the element.
+  if (targetEnv.allows(spirv::Capability::Kernel))
+    return spirv::PointerType::get(arrayElemType, storageClass);
+
+  // For Vulkan we need extra wrapping struct and array to satisfy interface
+  // needs.
   if (!type.hasStaticShape()) {
     int64_t stride = needsExplicitLayout(storageClass) ? *arrayElemSize : 0;
     auto arrayType = spirv::RuntimeArrayType::get(arrayElemType, stride);
     return wrapInStructAndGetPointer(arrayType, storageClass);
   }
 
   Optional<int64_t> memrefSize = getTypeNumBytes(options, type);
   if (!memrefSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce element count\n");
     return nullptr;
   }
 
   auto arrayElemCount = llvm::divideCeil(*memrefSize, *arrayElemSize);
   int64_t stride = needsExplicitLayout(storageClass) ? *arrayElemSize : 0;
   auto arrayType = spirv::ArrayType::get(arrayElemType, arrayElemCount, stride);
 
   return wrapInStructAndGetPointer(arrayType, storageClass);
 }
 
 SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
                                        const SPIRVConversionOptions &options)
     : targetEnv(targetAttr), options(options) {
   // Add conversions. The order matters here: later ones will be tried earlier.
 
   // Allow all SPIR-V dialect specific types. This assumes all builtin types
   // adopted in the SPIR-V dialect (i.e., IntegerType, FloatType, VectorType)
   // 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
Value strideVal = builder.create<spirv::ConstantOp>(
         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) {
+Value mlir::spirv::getVulkanElementPtr(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;
if (baseType.getRank() == 0) {
     linearizedIndices.push_back(zero);
   } else {
     linearizedIndices.push_back(
         linearizeIndex(indices, strides, offset, indexType, loc, builder));
   }
   return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
 }
 
+Value mlir::spirv::getOpenCLElementPtr(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();
+
+  SmallVector<Value, 2> linearizedIndices;
+  auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
+
+  Value linearIndex;
+  if (baseType.getRank() == 0) {
+    linearIndex = zero;
+  } else {
+    linearIndex =
+        linearizeIndex(indices, strides, offset, indexType, loc, builder);
+  }
+  return builder.create<spirv::PtrAccessChainOp>(loc, basePtr, linearIndex,
+                                                 linearizedIndices);
+}
+
+Value mlir::spirv::getElementPtr(SPIRVTypeConverter &typeConverter,
+                                 MemRefType baseType, Value basePtr,
+                                 ValueRange indices, Location loc,
+                                 OpBuilder &builder) {
+
+  if (typeConverter.allows(spirv::Capability::Kernel)) {
+    return getOpenCLElementPtr(typeConverter, baseType, basePtr, indices, loc,
+                               builder);
+  }
+
+  return getVulkanElementPtr(typeConverter, baseType, basePtr, indices, loc,
+                             builder);
+}
+
 //===----------------------------------------------------------------------===//
 // SPIR-V ConversionTarget
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<SPIRVConversionTarget>
 SPIRVConversionTarget::get(spirv::TargetEnvAttr targetAttr) {
   std::unique_ptr<SPIRVConversionTarget> target(
       // std::make_unique does not work here because the constructor is private.

mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -convert-gpu-to-spirv -verify-diagnostics %s -o - | FileCheck %s
 
 module attributes {
   gpu.container_module,
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Kernel, Addresses], []>, #spv.resource_limits<>>
 } {
   gpu.module @kernels {
     // CHECK-LABEL: spv.module @{{.*}} Physical64 OpenCL
     //       CHECK:   spv.func
     //  CHECK-SAME:     {{%.*}}: f32
     //   CHECK-NOT:     spv.interface_var_abi
-    //  CHECK-SAME:     {{%.*}}: !spv.ptr<!spv.struct<(!spv.array<12 x f32>)>, CrossWorkgroup>
+    //  CHECK-SAME:     {{%.*}}: !spv.ptr<f32, CrossWorkgroup>
     //   CHECK-NOT:     spv.interface_var_abi
     //  CHECK-SAME:     spv.entry_point_abi = #spv.entry_point_abi<local_size = dense<[32, 4, 1]> : vector<3xi32>>
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spv.storage_class<CrossWorkgroup>>) kernel
         attributes {spv.entry_point_abi = #spv.entry_point_abi<local_size = dense<[32, 4, 1]>: vector<3xi32>>} {
       gpu.return
     }
   }
 

mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir

func.func @store_i1(%dst: memref<4xi1, #spv.storage_class<StorageBuffer>>, %i: index) {
   memref.store %true, %dst[%i]: memref<4xi1, #spv.storage_class<StorageBuffer>>
   return
 }
 
 } // end module
 
 // -----
 
+// Check for Kernel capability, that with proper compute and storage extensions, we don't need to
+// perform special tricks.
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0,
+      [
+        Kernel, Addresses, Int8, Int16, Int64, Float16, Float64], []>, #spv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @load_store_zero_rank_float
+func.func @load_store_zero_rank_float(%arg0: memref<f32, #spv.storage_class<CrossWorkgroup>>, %arg1: memref<f32, #spv.storage_class<CrossWorkgroup>>) {
+  //      CHECK: [[ARG0:%.*]] = builtin.unrealized_conversion_cast {{.+}} : memref<f32, #spv.storage_class<CrossWorkgroup>> to !spv.ptr<f32, CrossWorkgroup>
+  //      CHECK: [[ARG1:%.*]] = builtin.unrealized_conversion_cast {{.+}} : memref<f32, #spv.storage_class<CrossWorkgroup>> to !spv.ptr<f32, CrossWorkgroup>
+  //      CHECK: [[ZERO1:%.*]] = spv.Constant 0 : i32
+  //      CHECK: spv.PtrAccessChain [[ARG0]][
+  // CHECK-SAME: [[ZERO1]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Load "CrossWorkgroup" %{{.*}} : f32
+  %0 = memref.load %arg0[] : memref<f32, #spv.storage_class<CrossWorkgroup>>
+  //      CHECK: [[ZERO2:%.*]] = spv.Constant 0 : i32
+  //      CHECK: spv.PtrAccessChain [[ARG1]][
+  // CHECK-SAME: [[ZERO2]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Store "CrossWorkgroup" %{{.*}} : f32
+  memref.store %0, %arg1[] : memref<f32, #spv.storage_class<CrossWorkgroup>>
+  return
+}
+
+// CHECK-LABEL: @load_store_zero_rank_int
+func.func @load_store_zero_rank_int(%arg0: memref<i32, #spv.storage_class<CrossWorkgroup>>, %arg1: memref<i32, #spv.storage_class<CrossWorkgroup>>) {
+  //      CHECK: [[ARG0:%.*]] = builtin.unrealized_conversion_cast {{.+}} : memref<i32, #spv.storage_class<CrossWorkgroup>> to  !spv.ptr<i32, CrossWorkgroup>
+  //      CHECK: [[ARG1:%.*]] = builtin.unrealized_conversion_cast {{.+}} : memref<i32, #spv.storage_class<CrossWorkgroup>> to  !spv.ptr<i32, CrossWorkgroup>
+  //      CHECK: [[ZERO1:%.*]] = spv.Constant 0 : i32
+  //      CHECK: spv.PtrAccessChain [[ARG0]][
+  // CHECK-SAME: [[ZERO1]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Load "CrossWorkgroup" %{{.*}} : i32
+  %0 = memref.load %arg0[] : memref<i32, #spv.storage_class<CrossWorkgroup>>
+  //      CHECK: [[ZERO2:%.*]] = spv.Constant 0 : i32
+  //      CHECK: spv.PtrAccessChain [[ARG1]][
+  // CHECK-SAME: [[ZERO2]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Store "CrossWorkgroup" %{{.*}} : i32
+  memref.store %0, %arg1[] : memref<i32, #spv.storage_class<CrossWorkgroup>>
+  return
+}
+
+// CHECK-LABEL: func @load_store_unknown_dim
+func.func @load_store_unknown_dim(%i: index, %source: memref<?xi32, #spv.storage_class<CrossWorkgroup>>, %dest: memref<?xi32, #spv.storage_class<CrossWorkgroup>>) {
+  // CHECK: %[[SRC:.+]] = builtin.unrealized_conversion_cast {{.+}} : memref<?xi32, #spv.storage_class<CrossWorkgroup>> to !spv.ptr<i32, CrossWorkgroup>
+  // CHECK: %[[DST:.+]] = builtin.unrealized_conversion_cast {{.+}} : memref<?xi32, #spv.storage_class<CrossWorkgroup>> to !spv.ptr<i32, CrossWorkgroup>
+  // CHECK: %[[AC0:.+]] = spv.PtrAccessChain %[[SRC]]
+  // CHECK: spv.Load "CrossWorkgroup" %[[AC0]]
+  %0 = memref.load %source[%i] : memref<?xi32, #spv.storage_class<CrossWorkgroup>>
+  // CHECK: %[[AC1:.+]] = spv.PtrAccessChain %[[DST]]
+  // CHECK: spv.Store "CrossWorkgroup" %[[AC1]]
+  memref.store %0, %dest[%i]: memref<?xi32, #spv.storage_class<CrossWorkgroup>>
+  return
+}
+
+// CHECK-LABEL: func @load_i1
+//  CHECK-SAME: (%[[SRC:.+]]: memref<4xi1, #spv.storage_class<CrossWorkgroup>>, %[[IDX:.+]]: index)
+func.func @load_i1(%src: memref<4xi1, #spv.storage_class<CrossWorkgroup>>, %i : index) -> i1 {
+  // CHECK-DAG: %[[SRC_CAST:.+]] = builtin.unrealized_conversion_cast %[[SRC]] : memref<4xi1, #spv.storage_class<CrossWorkgroup>> to !spv.ptr<i8, CrossWorkgroup>
+  // CHECK-DAG: %[[IDX_CAST:.+]] = builtin.unrealized_conversion_cast %[[IDX]]
+  // CHECK: %[[ZERO_0:.+]] = spv.Constant 0 : i32
+  // CHECK: %[[ZERO_1:.+]] = spv.Constant 0 : i32
+  // CHECK: %[[ONE:.+]] = spv.Constant 1 : i32
+  // CHECK: %[[MUL:.+]] = spv.IMul %[[ONE]], %[[IDX_CAST]] : i32
+  // CHECK: %[[ADD:.+]] = spv.IAdd %[[ZERO_1]], %[[MUL]] : i32
+  // CHECK: %[[ADDR:.+]] = spv.PtrAccessChain %[[SRC_CAST]][%[[ADD]]]
+  // CHECK: %[[VAL:.+]] = spv.Load "CrossWorkgroup" %[[ADDR]] : i8
+  // CHECK: %[[ONE_I8:.+]] = spv.Constant 1 : i8
+  // CHECK: %[[BOOL:.+]] = spv.IEqual %[[VAL]], %[[ONE_I8]] : i8
+  %0 = memref.load %src[%i] : memref<4xi1, #spv.storage_class<CrossWorkgroup>>
+  // CHECK: return %[[BOOL]]
+  return %0: i1
+}
+
+// CHECK-LABEL: func @store_i1
+//  CHECK-SAME: %[[DST:.+]]: memref<4xi1, #spv.storage_class<CrossWorkgroup>>,
+//  CHECK-SAME: %[[IDX:.+]]: index
+func.func @store_i1(%dst: memref<4xi1, #spv.storage_class<CrossWorkgroup>>, %i: index) {
+  %true = arith.constant true
+  // CHECK-DAG: %[[DST_CAST:.+]] = builtin.unrealized_conversion_cast %[[DST]] : memref<4xi1, #spv.storage_class<CrossWorkgroup>> to !spv.ptr<i8, CrossWorkgroup>
+  // CHECK-DAG: %[[IDX_CAST:.+]] = builtin.unrealized_conversion_cast %[[IDX]]
+  // CHECK: %[[ZERO_0:.+]] = spv.Constant 0 : i32
+  // CHECK: %[[ZERO_1:.+]] = spv.Constant 0 : i32
+  // CHECK: %[[ONE:.+]] = spv.Constant 1 : i32
+  // CHECK: %[[MUL:.+]] = spv.IMul %[[ONE]], %[[IDX_CAST]] : i32
+  // CHECK: %[[ADD:.+]] = spv.IAdd %[[ZERO_1]], %[[MUL]] : i32
+  // CHECK: %[[ADDR:.+]] = spv.PtrAccessChain %[[DST_CAST]][%[[ADD]]]
+  // CHECK: %[[ZERO_I8:.+]] = spv.Constant 0 : i8
+  // CHECK: %[[ONE_I8:.+]] = spv.Constant 1 : i8
+  // CHECK: %[[RES:.+]] = spv.Select %{{.+}}, %[[ONE_I8]], %[[ZERO_I8]] : i1, i8
+  // CHECK: spv.Store "CrossWorkgroup" %[[ADDR]], %[[RES]] : i8
+  memref.store %true, %dst[%i]: memref<4xi1, #spv.storage_class<CrossWorkgroup>>
+  return
+}
+
+} // end module
+
+// -----
+
 // Check that access chain indices are properly adjusted if non-32-bit types are
 // emulated via 32-bit types.
 // TODO: Test i64 types.
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spv.resource_limits<>>
 } {