[Mlir-commits] [mlir] [MLIR] Implement emulation of static indexing subbyte type vector stores (PR #115922)

[Han-Chung Wang]

================
@@ -306,6 +307,73 @@ emulatedVectorLoad(OpBuilder &rewriter, Location loc, Value base,
       newLoad);
 }
 
+/// Atomically store a subbyte-sized value to memory, with a mask.
+static void atomicStore(OpBuilder &rewriter, Location loc,
+                        TypedValue<MemRefType> emulatedMemref,
+                        Value emulatedIndex, TypedValue<VectorType> value,
+                        Value mask, int64_t scale) {
+  auto atomicOp = rewriter.create<memref::GenericAtomicRMWOp>(
+      loc, emulatedMemref, ValueRange{emulatedIndex});
+  OpBuilder builder =
+      OpBuilder::atBlockEnd(atomicOp.getBody(), rewriter.getListener());
+  Value origValue = atomicOp.getCurrentValue();
+
+  // i8 -> vector type <1xi8> then <1xi8> -> <scale x i.>
+  auto oneVectorType = VectorType::get({1}, origValue.getType());
+  auto fromElem = builder.create<vector::FromElementsOp>(loc, oneVectorType,
+                                                         ValueRange{origValue});
+  auto vectorBitCast =
+      builder.create<vector::BitCastOp>(loc, value.getType(), fromElem);
+
+  auto select =
+      builder.create<arith::SelectOp>(loc, mask, value, vectorBitCast);
+  auto bitcast2 = builder.create<vector::BitCastOp>(loc, oneVectorType, select);
+  auto extract = builder.create<vector::ExtractOp>(loc, bitcast2, 0);
+  builder.create<memref::AtomicYieldOp>(loc, extract.getResult());
----------------
hanhanW wrote:

IIUC, this path is for the non data-contention cases. In this case, can we use few `memref::AtomicRMWOp` ops to get rid of `GenericAtomicRMWOp`? (I assumed that the `GenericAtomicRMWOp` is more expensive.)

Here is an example from `memref.store`. It uses `arith::AtomicRMWKind::andi` to clear the bits, and `arith::AtomicRMWKind::ori` to write the bits.

https://github.com/llvm/llvm-project/blob/d5032b9f4b6aa415e7fd39701f29edb93028d8b3/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp#L435-L450

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