[Mlir-commits] [mlir] [mlir] Add subbyte emulation support for `memref.store`. (PR #73174)

[Han-Chung Wang]

================
@@ -301,6 +360,60 @@ struct ConvertMemRefReinterpretCast final
   }
 };
 
+//===----------------------------------------------------------------------===//
+// ConvertMemrefStore
+//===----------------------------------------------------------------------===//
+
+struct ConvertMemrefStore final : OpConversionPattern<memref::StoreOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::StoreOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto convertedType = adaptor.getMemref().getType().cast<MemRefType>();
+    int srcBits = op.getMemRefType().getElementTypeBitWidth();
+    int dstBits = convertedType.getElementTypeBitWidth();
+    auto dstIntegerType = rewriter.getIntegerType(dstBits);
+    if (dstBits % srcBits != 0) {
+      return rewriter.notifyMatchFailure(
+          op, "only dstBits % srcBits == 0 supported");
+    }
+
+    Location loc = op.getLoc();
+    Value extendedInput = rewriter.create<arith::ExtUIOp>(loc, dstIntegerType,
+                                                          adaptor.getValue());
+
+    // Special case 0-rank memref stores. We can compute the mask at compile
+    // time.
+    if (convertedType.getRank() == 0) {
+      // Create mask to clear destination bits
+      auto writeMaskValAttr =
+          rewriter.getIntegerAttr(dstIntegerType, ~(1 << (srcBits)) - 1);
+      Value writeMask = rewriter.create<arith::ConstantOp>(loc, dstIntegerType,
+                                                           writeMaskValAttr);
+
+      replaceStoreWithAtomics(rewriter, op, writeMask, extendedInput,
+                              adaptor.getMemref(), ValueRange{});
----------------
hanhanW wrote:

I think we can just create a atomic assign op (i.e., using `arith::AtomicRMWKind::assign` for the case. Then we don't need the `replaceStoreWithAtomics` util. It is better to remove the util because it is very simple and the context is very limited to the store op emulation.

https://github.com/llvm/llvm-project/pull/73174