[mlir][spirv] Avoid struct wrapping in non-interface storage class

[antiagainst]

Summary

Per the SPIR-V spec "2.16.2. Validation Rules for Shader Capabilities":

Composite objects in the StorageBuffer, PhysicalStorageBuffer,
Uniform, and PushConstant Storage Classes must be explicitly
laid out.

Per the Vulkan spec "15.6. Shader Resource Interface":

The shader variables defined with a storage class of PushConstant
that are statically used by the shader entry points for the
pipeline define the push constant interface. They must be:
* typed as OpTypeStruct.

Variables identified with the Uniform storage class are used to
access transparent buffer backed resources. Such variables must be:
* typed as OpTypeStruct, or an array of this type,

Variables identified with the StorageBuffer storage class are
used to access transparent buffer backed resources. Such
variables must be:
* typed as OpTypeStruct, or an array of this type,

For other cases we don't need to wrap with a struct.

[mravishankar]

Wondering if there is a problem if we have the struct all the time. One of the issues on SPIR-V side is that we cannot do pointer casts. So a memref_cast where the operand is from something that doesnt need layout to something that does need layout would be not possible support. In general, there is no information of where the memref came from, so we might end up in situations where we cannot lower such memref_cast. If there is no reason to not have a struct, it seems more uniform to have a struct by default for all memref types.

[antiagainst]

Good points Mahesh!

In D98507#2627024, @mravishankar wrote:

Wondering if there is a problem if we have the struct all the time.

Main motivation is to provide as clean code as possible to the backend driver to avoid unnecessary overhead there (like no need for SROA etc.). But I don't have concrete numbers to say how much it matters for now.

One of the issues on SPIR-V side is that we cannot do pointer casts. So a memref_cast where the operand is from something that doesnt need layout to something that does need layout would be not possible support. In general, there is no information of where the memref came from, so we might end up in situations where we cannot lower such memref_cast.

Actually this should not be a valid case. Different storage classes potentially mean different hardware functionality blocks. So casting pointers from one storage class (i.e., memref memory space) to another one is not valid in general. One will need to load/store to cross storage class boundaries in general.

If there is no reason to not have a struct, it seems more uniform to have a struct by default for all memref types.

Sounds reasonable to me. I'll keep the struct still for now. Send D100386 to drop the struct offset though.

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

static Value adjustAccessChainForBitwidth(SPIRVTypeConverter &typeConverter,
   assert(targetBits % sourceBits == 0);
   const auto loc = op.getLoc();
   IntegerType targetType = builder.getIntegerType(targetBits);
   IntegerAttr attr =
       builder.getIntegerAttr(targetType, targetBits / sourceBits);
   auto idx = builder.create<spirv::ConstantOp>(loc, targetType, attr);
   auto lastDim = op->getOperand(op.getNumOperands() - 1);
   auto indices = llvm::to_vector<4>(op.indices());
-  // There are two elements if this is a 1-D tensor.
-  assert(indices.size() == 2);
+  // There are at most two elements (one implicit zero for indexing into
+  // wrapping struct in interface storage classes) if this is a 1-D tensor.
+  assert(indices.size() <= 2);
   indices.back() = builder.create<spirv::SDivOp>(loc, lastDim, idx);
   Type t = typeConverter.convertType(op.component_ptr().getType());
   return builder.create<spirv::AccessChainOp>(loc, t, op.base_ptr(), indices);
 }
 
 /// Returns the shifted `targetBits`-bit value with the given offset.
 static Value shiftValue(Location loc, Value value, Value offset, Value mask,
                         int targetBits, OpBuilder &builder) {
if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, loadOperands.memref(),
                            loadOperands.indices(), loc, rewriter);
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
-  auto dstType = typeConverter.convertType(memrefType)
+  Type dstType = typeConverter.convertType(memrefType)
                      .cast<spirv::PointerType>()
-                     .getPointeeType()
-                     .cast<spirv::StructType>()
-                     .getElementType(0)
-                     .cast<spirv::ArrayType>()
-                     .getElementType();
+                     .getPointeeType();
+  if (dstType.isa<spirv::StructType>())
+    dstType = dstType.cast<spirv::StructType>().getElementType(0);
+  dstType = dstType.cast<spirv::ArrayType>().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) {
     rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp,
                                                accessChainOp.getResult());
     return success();
   }
 
   // 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);
+
+  // There are at most two indices (one implicit zero for indexing into
+  // wrapping struct in interface storage classes).
+  assert(accessChainOp.indices().size() <= 2);
   Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
                                                    srcBits, dstBits, rewriter);
   Value spvLoadOp = rewriter.create<spirv::LoadOp>(
       loc, dstType, adjustedPtr,
       loadOp->getAttrOfType<IntegerAttr>(
           spirv::attributeName<spirv::MemoryAccess>()),
       loadOp->getAttrOfType<IntegerAttr>("alignment"));
 
if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto loc = storeOp.getLoc();
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, storeOperands.memref(),
                            storeOperands.indices(), loc, rewriter);
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
-  auto dstType = typeConverter.convertType(memrefType)
+  Type dstType = typeConverter.convertType(memrefType)
                      .cast<spirv::PointerType>()
-                     .getPointeeType()
-                     .cast<spirv::StructType>()
-                     .getElementType(0)
-                     .cast<spirv::ArrayType>()
-                     .getElementType();
+                     .getPointeeType();
+  if (dstType.isa<spirv::StructType>())
+    dstType = dstType.cast<spirv::StructType>().getElementType(0);
+  dstType = dstType.cast<spirv::ArrayType>().getElementType();
   int dstBits = dstType.getIntOrFloatBitWidth();
   assert(dstBits % srcBits == 0);
 
   if (srcBits == dstBits) {
     rewriter.replaceOpWithNewOp<spirv::StoreOp>(
         storeOp, accessChainOp.getResult(), storeOperands.value());
     return success();
   }
 
   // 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
   // 6) store 32-bit value back
   // The step 1 to step 3 are done by AtomicAnd as one atomic step, and the step
   // 4 to step 6 are done by AtomicOr as another atomic step.
-  assert(accessChainOp.indices().size() == 2);
+
+  // There are at most two indices (one implicit zero for indexing into
+  // wrapping struct in interface storage classes).
+  assert(accessChainOp.indices().size() <= 2);
   Value lastDim = accessChainOp->getOperand(accessChainOp.getNumOperands() - 1);
   Value offset = getOffsetForBitwidth(loc, lastDim, srcBits, dstBits, rewriter);
 
   // Create a mask to clear the destination. E.g., if it is the second i8 in
   // i32, 0xFFFF00FF is created.
   Value mask = rewriter.create<spirv::ConstantOp>(
       loc, dstType, rewriter.getIntegerAttr(dstType, (1 << srcBits) - 1));
   Value clearBitsMask =

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

LLVM_DEBUG(llvm::dbgs()
                << label << " illegal: requires at least one capability in ["
                << llvm::join(capStrings, ", ")
                << "] but none allowed in target environment\n");
     return failure();
   }
   return success();
 }
 
+/// Returns true if the given `storageClass` needs explicit layout when used in
+/// Shader environments.
+static bool needsExplicitLayout(spirv::StorageClass storageClass) {
+  switch (storageClass) {
+  case spirv::StorageClass::PhysicalStorageBuffer:
+  case spirv::StorageClass::PushConstant:
+  case spirv::StorageClass::StorageBuffer:
+  case spirv::StorageClass::Uniform:
+    return true;
+  default:
+    return false;
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
 Type SPIRVTypeConverter::getIndexType(MLIRContext *context) {
   // Convert to 32-bit integers for now. Might need a way to control this in
   // future.
   // TODO: It is probably better to make it 64-bit integers. To
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. 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);
+  Type resultType = arrayType;
+  if (needsExplicitLayout(*storageClass)) {
+    // Wrap in a struct to satisfy Vulkan interface requirements.
+    resultType = spirv::StructType::get(resultType, 0);
+  }
+  return spirv::PointerType::get(resultType, *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.
   addConversion([](Type type) { return llvm::None; });
if (failed(getStridesAndOffset(baseType, strides, offset)) ||
     return nullptr;
   }
 
   auto indexType = typeConverter.getIndexType(builder.getContext());
 
   SmallVector<Value, 2> linearizedIndices;
   auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
 
+  auto storageClass = SPIRVTypeConverter::getStorageClassForMemorySpace(
+      baseType.getMemorySpaceAsInt());
+  if (!storageClass)
+    return {};
+
+  if (needsExplicitLayout(*storageClass)) {
     // Add a '0' at the start to index into the struct.
     linearizedIndices.push_back(zero);
+  }
 
   if (baseType.getRank() == 0) {
     linearizedIndices.push_back(zero);
   } else {
     linearizedIndices.push_back(
         linearizeIndex(indices, strides, offset, loc, builder));
   }
   return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);

mlir/test/Conversion/GPUToSPIRV/module-structure-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], []>, {}>
 } {
   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, stride=4> [0])>, CrossWorkgroup>
+    //  CHECK-SAME:     {{%.*}}: !spv.ptr<!spv.array<12 x f32, stride=4>, CrossWorkgroup>
     //   CHECK-NOT:     spv.interface_var_abi
     //  CHECK-SAME:     spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>}
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, 11>) kernel
         attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
       gpu.return
     }
   }
 

mlir/test/Conversion/StandardToSPIRV/alloc.mlir

module attributes {
   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: spv.GlobalVariable @[[VAR:.+]] : !spv.ptr<!spv.array<20 x f32, stride=4>, Workgroup>
 //     CHECK: func @alloc_dealloc_workgroup_mem
 // CHECK-NOT:   alloc
 //     CHECK:   %[[PTR:.+]] = spv.mlir.addressof @[[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
func @alloc_dealloc_workgroup_mem(%arg0 : index, %arg1 : index) {
     %1 = load %0[%arg0, %arg1] : memref<4x5xi16, 3>
     store %1, %0[%arg0, %arg1] : memref<4x5xi16, 3>
     dealloc %0 : memref<4x5xi16, 3>
     return
   }
 }
 
 //       CHECK: spv.GlobalVariable @__workgroup_mem__{{[0-9]+}}
-//  CHECK-SAME:   !spv.ptr<!spv.struct<(!spv.array<20 x i32, stride=4>)>, Workgroup>
+//  CHECK-SAME:   !spv.ptr<!spv.array<20 x i32, stride=4>, Workgroup>
 // CHECK_LABEL: spv.func @alloc_dealloc_workgroup_mem
 //       CHECK:   %[[VAR:.+]] = spv.mlir.addressof @__workgroup_mem__0
 //       CHECK:   %[[LOC:.+]] = spv.SDiv
-//       CHECK:   %[[PTR:.+]] = spv.AccessChain %[[VAR]][%{{.+}}, %[[LOC]]]
+//       CHECK:   %[[PTR:.+]] = spv.AccessChain %[[VAR]][%[[LOC]]]
 //       CHECK:   %{{.+}} = spv.Load "Workgroup" %[[PTR]] : i32
 //       CHECK:   %[[LOC:.+]] = spv.SDiv
-//       CHECK:   %[[PTR:.+]] = spv.AccessChain %[[VAR]][%{{.+}}, %[[LOC]]]
+//       CHECK:   %[[PTR:.+]] = spv.AccessChain %[[VAR]][%[[LOC]]]
 //       CHECK:   %{{.+}} = spv.AtomicAnd "Workgroup" "AcquireRelease" %[[PTR]], %{{.+}} : !spv.ptr<i32, Workgroup>
 //       CHECK:   %{{.+}} = spv.AtomicOr "Workgroup" "AcquireRelease" %[[PTR]], %{{.+}} : !spv.ptr<i32, Workgroup>
 
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   }
 {
   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-SAME:   !spv.ptr<!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-SAME:   !spv.ptr<!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]>, {}>
   }
 {
   func @two_allocs_vector() {
     %0 = alloc() : memref<4xvector<4xf32>, 3>
     %1 = alloc() : memref<2xvector<2xi32>, 3>
     return
   }
 }
 
 //  CHECK-DAG: spv.GlobalVariable @__workgroup_mem__{{[0-9]+}}
-// CHECK-SAME:   !spv.ptr<!spv.struct<(!spv.array<2 x vector<2xi32>, stride=8>)>, Workgroup>
+// CHECK-SAME:   !spv.ptr<!spv.array<2 x vector<2xi32>, stride=8>, Workgroup>
 //  CHECK-DAG: spv.GlobalVariable @__workgroup_mem__{{[0-9]+}}
-// CHECK-SAME:   !spv.ptr<!spv.struct<(!spv.array<4 x vector<4xf32>, stride=16>)>, Workgroup>
+// CHECK-SAME:   !spv.ptr<!spv.array<4 x vector<4xf32>, stride=16>, Workgroup>
 //      CHECK: spv.func @two_allocs_vector()
 //      CHECK: spv.Return
 
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<

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

 // CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x si32, stride=4> [0])>, Uniform>
 func @memref_16bit_Uniform(%arg0: memref<16xsi16, 4>) { return }
 
 // CHECK-LABEL: spv.func @memref_16bit_PushConstant
 // CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x ui32, stride=4> [0])>, PushConstant>
 func @memref_16bit_PushConstant(%arg0: memref<16xui16, 7>) { return }
 
 // CHECK-LABEL: spv.func @memref_16bit_Input
-// CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, Input>
-func @memref_16bit_Input(%arg3: memref<16xf16, 9>) { return }
+// CHECK-SAME: !spv.ptr<!spv.array<4 x f32, stride=4>, Input>
+func @memref_16bit_Input(%arg3: memref<4xf16, 9>) { return }
 
 // CHECK-LABEL: spv.func @memref_16bit_Output
-// CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, Output>
-func @memref_16bit_Output(%arg4: memref<16xf16, 10>) { return }
+// CHECK-SAME: !spv.ptr<!spv.array<4 x f32, stride=4>, Output>
+func @memref_16bit_Output(%arg4: memref<4xf16, 10>) { return }
 
 } // end module
 
 // -----
 
 // Check that using non-32-bit scalar types in interface storage classes
 // requires special capability and extension: keep as-is when the capability
 // and extension is available.

 // requires special capability and extension: keep as-is when the capability
 // and extension is available.
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [StorageInputOutput16], [SPV_KHR_16bit_storage]>, {}>
 } {
 
 // CHECK-LABEL: spv.func @memref_16bit_Input
-// CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x f16, stride=2> [0])>, Input>
-func @memref_16bit_Input(%arg3: memref<16xf16, 9>) { return }
+// CHECK-SAME: !spv.ptr<!spv.array<4 x f16, stride=2>, Input>
+func @memref_16bit_Input(%arg3: memref<4xf16, 9>) { return }
 
 // CHECK-LABEL: spv.func @memref_16bit_Output
-// CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x i16, stride=2> [0])>, Output>
-func @memref_16bit_Output(%arg4: memref<16xi16, 10>) { return }
+// CHECK-SAME: !spv.ptr<!spv.array<4 x i16, stride=2>, Output>
+func @memref_16bit_Output(%arg4: memref<4xi16, 10>) { return }
 
 } // end module
 
 // -----
 
 // Check that memref offset and strides affect the array size.
 module attributes {
   spv.target_env = #spv.target_env<