[mlir][StandardToSPIRV] Emulate bitwidths not supported for store op.

[hanchung]

Summary

As D78974, this patch implements the emulation for store op. The emulation is
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.

[mravishankar]

@antiagainst PTAL at the suggestion here.

@hanchung : Thanks for porting this to MLIR core. This was what we had converged on in IREE where we were trying to avoid the expensive compare exchange. We tried this, but looking at it with fresh eyes this seems buggy.

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

731

Might be easier to just convert the memref element type using the type converter (and do the same at the loadOp) instead of this.

747–757

I think these steps will not give the result that you want in all cases. There will be a race between the other threads that are modifying the same word, but different bits in the word. I dont see a way around it apart from using a atomic compare exchange.

do {

%val = load %basePtr[%baseOffset] : !spv.ptr<i32, ..>
%clearbits = and %val, %mask : i32 // Clear the bits that need to be updated.
%updateval = or %clearbits, %update : i32 // Update the bits as required

} while( %val != atomicCompareExchange(%basePtr, %updateVal, %val));

See semantics of compareExchange here : https://www.khronos.org/registry/spir-v/specs/unified1/SPIRV.html#OpAtomicCompareExchange. If the condition of the while failed, it implies the compareExchange succeeded, i.e. the value at %basePtr was same as the %val that was loaded. If the condition failed, then the value at %basePtr was different from what was loaded indicating a stale load (due to an other write to the same word).

[antiagainst]

Nice! Thanks Hanhan! I just have a few nits. Feel free to land after addressing them.

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

149

Not necessarily 32-bit here?

150

No need to require the whole StoreOpOperandAdaptor here given we just need op.value() in this function? Just passing in a Value? That way it creates cleaner interface.

304

The comment needs to be updated. :)

[antiagainst]

Sorry didn't see Mahesh's comment before commenting. :)

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

731

This was requested by me. Actually depending on the storage class, we might want different capabilities, which determines whether we can use 16-bit or not in that particular storage class. For example, 16-bit in storage buffer is fine if we have StorageBuffer16BitAccess, but 16-bit in push constant is not fine if we don't have StoragePushConstant16. So we need to go through the memref here.

747–757

I might be missing something but the approach here should be fine given each thread is actually only touching strictly disjoint portion of the word? Could you elaborate a bit here regarding the issue?

[mravishankar]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

747–757

I dont think that is the case. Here is the proposed sequence of operations by this patch

load
atomicAnd
load
atomicOr.

Lets just focus on the last two instructions

load
atomicOr

between the load and the atomicOr some other thread might have written to the same word but at different bits. So the load value that is seen might be different than what exists in memory at the time of atomicOr. A thread can overwrite the entire word if it can guarantee that the entire word was same as it saw when it did the load. Otherwise its load value is stale. It has to abort and retry. Thats what the suggested snippet is doing. It ensures that when a thread updates the word, the value it saw at the load is not stale.

[mravishankar]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

747–757

Sorry, got confused by the comment. There is no intermediate load. Then this works. Apologies for the noise.

[antiagainst]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

747–757

Yup the comment can be made clearer by saying "... to step 3 are done by spv.AtomicAnd as one atomic step, ... by spv.AtomicOr as another atomic step." Hanhan could you update here?

[hanchung]

Address comments

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

static Value adjustAccessChainForBitwidth(SPIRVTypeConverter &typeConverter,
   auto indices = llvm::to_vector<4>(op.indices());
   // There are two elements 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.
+Value shiftValue(Location loc, Value value, Value offset, Value mask,
+                 int targetBits, OpBuilder &builder) {
+  Type targetType = builder.getIntegerType(targetBits);
+  Value result = builder.create<spirv::BitwiseAndOp>(loc, value, mask);
+  return builder.create<spirv::ShiftLeftLogicalOp>(loc, targetType, result,
+                                                   offset);
+}
+
 //===----------------------------------------------------------------------===//
 // 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.
 

 class SelectOpPattern final : public SPIRVOpLowering<SelectOp> {
 public:
   using SPIRVOpLowering<SelectOp>::SPIRVOpLowering;
   LogicalResult
   matchAndRewrite(SelectOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts std.store to spv.Store on integers.
+class IntStoreOpPattern final : public SPIRVOpLowering<StoreOp> {
+public:
+  using SPIRVOpLowering<StoreOp>::SPIRVOpLowering;
+
+  LogicalResult
+  matchAndRewrite(StoreOp storeOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Converts std.store to spv.Store.
 class StoreOpPattern final : public SPIRVOpLowering<StoreOp> {
 public:
   using SPIRVOpLowering<StoreOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(StoreOp storeOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
SelectOpPattern::matchAndRewrite(SelectOp op, ArrayRef<Value> operands,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
+IntStoreOpPattern::matchAndRewrite(StoreOp storeOp, ArrayRef<Value> operands,
+                                   ConversionPatternRewriter &rewriter) const {
+  StoreOpOperandAdaptor storeOperands(operands);
+  auto memrefType = storeOp.memref().getType().cast<MemRefType>();
+  if (!memrefType.getElementType().isSignlessInteger())
+    return failure();
+
+  auto loc = storeOp.getLoc();
+  spirv::AccessChainOp accessChainOp =
+      spirv::getElementPtr(typeConverter, memrefType, storeOperands.memref(),
+                           storeOperands.indices(), loc, rewriter);
+  int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
+  auto dstType = typeConverter.convertType(memrefType)
+                     .cast<spirv::PointerType>()
+                     .getPointeeType()
+                     .cast<spirv::StructType>()
+                     .getElementType(0)
+                     .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);
+  Value lastDim = accessChainOp.getOperation()->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 =
+      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);
+  Value result = rewriter.create<spirv::AtomicAndOp>(
+      loc, dstType, adjustedPtr, spirv::Scope::Device,
+      spirv::MemorySemantics::AcquireRelease, clearBitsMask);
+  result = rewriter.create<spirv::AtomicOrOp>(
+      loc, dstType, adjustedPtr, spirv::Scope::Device,
+      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());
+  rewriter.eraseOp(accessChainOp);
+
+  return success();
+}
+
+LogicalResult
 StoreOpPattern::matchAndRewrite(StoreOp storeOp, ArrayRef<Value> operands,
                                 ConversionPatternRewriter &rewriter) const {
   StoreOpOperandAdaptor storeOperands(operands);
-  auto storePtr = spirv::getElementPtr(
-      typeConverter, storeOp.memref().getType().cast<MemRefType>(),
-      storeOperands.memref(), storeOperands.indices(), storeOp.getLoc(),
-      rewriter);
+  auto memrefType = storeOp.memref().getType().cast<MemRefType>();
+  if (memrefType.getElementType().isSignlessInteger())
+    return failure();
+  auto storePtr =
+      spirv::getElementPtr(typeConverter, memrefType, storeOperands.memref(),
+                           storeOperands.indices(), storeOp.getLoc(), rewriter);
   rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, storePtr,
                                               storeOperands.value());
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // XorOp
 //===----------------------------------------------------------------------===//
patterns.insert<
       UnaryAndBinaryOpPattern<TanhOp, spirv::GLSLTanhOp>,
       UnaryAndBinaryOpPattern<UnsignedDivIOp, spirv::UDivOp>,
       UnaryAndBinaryOpPattern<UnsignedRemIOp, spirv::UModOp>,
       UnaryAndBinaryOpPattern<UnsignedShiftRightOp, spirv::ShiftRightLogicalOp>,
       BitwiseOpPattern<AndOp, spirv::LogicalAndOp, spirv::BitwiseAndOp>,
       BitwiseOpPattern<OrOp, spirv::LogicalOrOp, spirv::BitwiseOrOp>,
       BoolCmpIOpPattern, ConstantCompositeOpPattern, ConstantScalarOpPattern,
       CmpFOpPattern, CmpIOpPattern, IntLoadOpPattern, LoadOpPattern,
-      ReturnOpPattern, SelectOpPattern, StoreOpPattern,
+      ReturnOpPattern, SelectOpPattern, IntStoreOpPattern, StoreOpPattern,
       TypeCastingOpPattern<SIToFPOp, spirv::ConvertSToFOp>,
       TypeCastingOpPattern<ZeroExtendIOp, spirv::UConvertOp>,
       TypeCastingOpPattern<TruncateIOp, spirv::SConvertOp>,
       TypeCastingOpPattern<FPExtOp, spirv::FConvertOp>,
       TypeCastingOpPattern<FPTruncOp, spirv::FConvertOp>, XOrOpPattern>(
       context, typeConverter);
 }
 } // namespace mlir

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

 }
 
 } // end module
 
 // -----
 
 // Check that access chain indices are properly adjusted if non-32-bit types are
 // emulated via 32-bit types.
-// TODO: Test i64 type.
+// TODO: Test i1 and i64 types.
 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>}>
 } {
 
 // CHECK-LABEL: @load_i8

 func @load_f32(%arg0: memref<f32>) {
   // CHECK-NOT: spv.SDiv
   //     CHECK: spv.Load
   // CHECK-NOT: spv.ShiftRightArithmetic
   %0 = load %arg0[] : memref<f32>
   return
 }
 
+// CHECK-LABEL: @store_i8
+//       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i32)
+func @store_i8(%arg0: memref<i8>, %value: i8) {
+  //     CHECK: %[[ZERO:.+]] = spv.constant 0 : i32
+  //     CHECK: %[[FOUR:.+]] = spv.constant 4 : i32
+  //     CHECK: %[[EIGHT:.+]] = spv.constant 8 : i32
+  //     CHECK: %[[IDX:.+]] = spv.SMod %[[ZERO]], %[[FOUR]] : i32
+  //     CHECK: %[[OFFSET:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
+  //     CHECK: %[[MASK1:.+]] = spv.constant 255 : i32
+  //     CHECK: %[[TMP1:.+]] = spv.ShiftLeftLogical %[[MASK1]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[MASK:.+]] = spv.Not %[[TMP1]] : i32
+  //     CHECK: %[[CLAMPED_VAL:.+]] = spv.BitwiseAnd %[[ARG1]], %[[MASK1]] : i32
+  //     CHECK: %[[STORE_VAL:.+]] = spv.ShiftLeftLogical %[[CLAMPED_VAL]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[FOUR2:.+]] = spv.constant 4 : i32
+  //     CHECK: %[[ACCESS_IDX:.+]] = spv.SDiv %[[ZERO]], %[[FOUR2]] : i32
+  //     CHECK: %[[PTR:.+]] = spv.AccessChain %[[ARG0]][%[[ZERO]], %[[ACCESS_IDX]]]
+  //     CHECK: spv.AtomicAnd "Device" "AcquireRelease" %[[PTR]], %[[MASK]]
+  //     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: %[[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
+  //     CHECK: %[[STORE_VAL:.+]] = spv.ShiftLeftLogical %[[CLAMPED_VAL]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[TWO2:.+]] = spv.constant 2 : i32
+  //     CHECK: %[[ACCESS_IDX:.+]] = spv.SDiv %[[FLAT_IDX]], %[[TWO2]] : i32
+  //     CHECK: %[[PTR:.+]] = spv.AccessChain %[[ARG0]][%[[ZERO]], %[[ACCESS_IDX]]]
+  //     CHECK: spv.AtomicAnd "Device" "AcquireRelease" %[[PTR]], %[[MASK]]
+  //     CHECK: spv.AtomicOr "Device" "AcquireRelease" %[[PTR]], %[[STORE_VAL]]
+  store %value, %arg0[%index] : memref<10xi16>
+  return
+}
+
+// CHECK-LABEL: @store_i32
+func @store_i32(%arg0: memref<i32>, %value: i32) {
+  //     CHECK: spv.Store
+  // CHECK-NOT: spv.AtomicAnd
+  // CHECK-NOT: spv.AtomicOr
+  store %value, %arg0[] : memref<i32>
+  return
+}
+
+// CHECK-LABEL: @store_f32
+func @store_f32(%arg0: memref<f32>, %value: f32) {
+  //     CHECK: spv.Store
+  // CHECK-NOT: spv.AtomicAnd
+  // CHECK-NOT: spv.AtomicOr
+  store %value, %arg0[] : memref<f32>
+  return
+}
+
 } // end module
 
 // -----
 
 // Check that access chain indices are properly adjusted if non-16/32-bit types
 // are emulated via 32-bit types.
 module attributes {
   spv.target_env = #spv.target_env<

 func @load_i16(%arg0: memref<i16>) {
   // CHECK-NOT: spv.SDiv
   //     CHECK: spv.Load
   // CHECK-NOT: spv.ShiftRightArithmetic
   %0 = load %arg0[] : memref<i16>
   return
 }
 
+// CHECK-LABEL: @store_i8
+//       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i32)
+func @store_i8(%arg0: memref<i8>, %value: i8) {
+  //     CHECK: %[[ZERO:.+]] = spv.constant 0 : i32
+  //     CHECK: %[[FOUR:.+]] = spv.constant 4 : i32
+  //     CHECK: %[[EIGHT:.+]] = spv.constant 8 : i32
+  //     CHECK: %[[IDX:.+]] = spv.SMod %[[ZERO]], %[[FOUR]] : i32
+  //     CHECK: %[[OFFSET:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
+  //     CHECK: %[[MASK1:.+]] = spv.constant 255 : i32
+  //     CHECK: %[[TMP1:.+]] = spv.ShiftLeftLogical %[[MASK1]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[MASK:.+]] = spv.Not %[[TMP1]] : i32
+  //     CHECK: %[[CLAMPED_VAL:.+]] = spv.BitwiseAnd %[[ARG1]], %[[MASK1]] : i32
+  //     CHECK: %[[STORE_VAL:.+]] = spv.ShiftLeftLogical %[[CLAMPED_VAL]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[FOUR2:.+]] = spv.constant 4 : i32
+  //     CHECK: %[[ACCESS_IDX:.+]] = spv.SDiv %[[ZERO]], %[[FOUR2]] : i32
+  //     CHECK: %[[PTR:.+]] = spv.AccessChain %[[ARG0]][%[[ZERO]], %[[ACCESS_IDX]]]
+  //     CHECK: spv.AtomicAnd "Device" "AcquireRelease" %[[PTR]], %[[MASK]]
+  //     CHECK: spv.AtomicOr "Device" "AcquireRelease" %[[PTR]], %[[STORE_VAL]]
+  store %value, %arg0[] : memref<i8>
+  return
+}
+
+// CHECK-LABEL: @store_i16
+func @store_i16(%arg0: memref<10xi16>, %index: index, %value: i16) {
+  //     CHECK: spv.Store
+  // CHECK-NOT: spv.AtomicAnd
+  // CHECK-NOT: spv.AtomicOr
+  store %value, %arg0[%index] : memref<10xi16>
+  return
+}
+
 } // end module