[mlir][vector] Add support for vector.maskedstore sub-type emulation.

[hanhanW]

The idea is similar to vector.maskedload + vector.store emulation. What the emulation does is:

1. Get a compressed mask and load the data from destination.
2. Bitcast the data to original vector type.
3. Select values between op.valueToStore and the data from load using original mask.
4. Bitcast the new value and store it to destination using compressed masked.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-vector

Author: Han-Chung Wang (hanhanW)

Changes

The idea is similar to vector.maskedload + vector.store emulation. What the emulation does is:

1. Get a compressed mask and load the data from destination.
2. Bitcast the data to original vector type.
3. Select values between op.valueToStore and the data from load using original mask.
4. Bitcast the new value and store it to destination using compressed masked.

---

Full diff: https://github.com/llvm/llvm-project/pull/73871.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp (+168-62)
• (modified) mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir (+72)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index 6aea0343bfc9327..05c98b89e8a94c1 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -32,6 +32,78 @@ using namespace mlir;
 #define DBGSNL() (llvm::dbgs() << "\n")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
+/// Returns a compressed mask. The mask value is set only if any mask is present
+/// in the the scale range. E.g., if `scale` equals to 2, the following mask:
+///
+///   %mask = [1, 1, 1, 0, 0, 0]
+///
+/// will return the following new compressed mask:
+///
+///   %mask = [1, 1, 0]
+static FailureOr<Operation *> getCompressedMaskOp(OpBuilder &rewriter,
+                                                  Location loc, Value mask,
+                                                  int origElements, int scale) {
+  auto numElements = (origElements + scale - 1) / scale;
+
+  auto maskOp = mask.getDefiningOp();
+  SmallVector<vector::ExtractOp, 2> extractOps;
+  // Finding the mask creation operation.
+  while (maskOp && !isa<vector::CreateMaskOp, vector::ConstantMaskOp>(maskOp)) {
+    if (auto extractOp = dyn_cast<vector::ExtractOp>(maskOp)) {
+      maskOp = extractOp.getVector().getDefiningOp();
+      extractOps.push_back(extractOp);
+    }
+  }
+  auto createMaskOp = dyn_cast_or_null<vector::CreateMaskOp>(maskOp);
+  auto constantMaskOp = dyn_cast_or_null<vector::ConstantMaskOp>(maskOp);
+  if (!createMaskOp && !constantMaskOp)
+    return failure();
+
+  // Computing the "compressed" mask. All the emulation logic (i.e. computing
+  // new mask index) only happens on the last dimension of the vectors.
+  Operation *newMask = nullptr;
+  auto shape = llvm::to_vector(
+      maskOp->getResultTypes()[0].cast<VectorType>().getShape().drop_back());
+  shape.push_back(numElements);
+  auto newMaskType = VectorType::get(shape, rewriter.getI1Type());
+  if (createMaskOp) {
+    auto maskOperands = createMaskOp.getOperands();
+    auto numMaskOperands = maskOperands.size();
+    AffineExpr s0;
+    bindSymbols(rewriter.getContext(), s0);
+    s0 = s0 + scale - 1;
+    s0 = s0.floorDiv(scale);
+    OpFoldResult origIndex =
+        getAsOpFoldResult(maskOperands[numMaskOperands - 1]);
+    OpFoldResult maskIndex =
+        affine::makeComposedFoldedAffineApply(rewriter, loc, s0, origIndex);
+    auto newMaskOperands = llvm::to_vector(maskOperands.drop_back());
+    newMaskOperands.push_back(
+        getValueOrCreateConstantIndexOp(rewriter, loc, maskIndex));
+    newMask = rewriter.create<vector::CreateMaskOp>(loc, newMaskType,
+                                                    newMaskOperands);
+  } else if (constantMaskOp) {
+    auto maskDimSizes = constantMaskOp.getMaskDimSizes().getValue();
+    auto numMaskOperands = maskDimSizes.size();
+    auto origIndex =
+        cast<IntegerAttr>(maskDimSizes[numMaskOperands - 1]).getInt();
+    auto maskIndex =
+        rewriter.getI64IntegerAttr((origIndex + scale - 1) / scale);
+    auto newMaskDimSizes = llvm::to_vector(maskDimSizes.drop_back());
+    newMaskDimSizes.push_back(maskIndex);
+    newMask = rewriter.create<vector::ConstantMaskOp>(
+        loc, newMaskType, rewriter.getArrayAttr(newMaskDimSizes));
+  }
+
+  while (!extractOps.empty()) {
+    newMask = rewriter.create<vector::ExtractOp>(
+        loc, newMask->getResults()[0], extractOps.back().getMixedPosition());
+    extractOps.pop_back();
+  }
+
+  return newMask;
+}
+
 namespace {
 
 //===----------------------------------------------------------------------===//
@@ -99,6 +171,94 @@ struct ConvertVectorStore final : OpConversionPattern<vector::StoreOp> {
   }
 };
 
+//===----------------------------------------------------------------------===//
+// ConvertVectorMaskedStore
+//===----------------------------------------------------------------------===//
+
+struct ConvertVectorMaskedStore final
+    : OpConversionPattern<vector::MaskedStoreOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::MaskedStoreOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    auto loc = op.getLoc();
+    auto convertedType = cast<MemRefType>(adaptor.getBase().getType());
+    Type oldElementType = op.getValueToStore().getType().getElementType();
+    Type newElementType = convertedType.getElementType();
+    int srcBits = oldElementType.getIntOrFloatBitWidth();
+    int dstBits = newElementType.getIntOrFloatBitWidth();
+
+    if (dstBits % srcBits != 0) {
+      return rewriter.notifyMatchFailure(
+          op, "only dstBits % srcBits == 0 supported");
+    }
+
+    int scale = dstBits / srcBits;
+    auto origElements = op.getValueToStore().getType().getNumElements();
+    if (origElements % scale != 0)
+      return failure();
+
+    auto stridedMetadata =
+        rewriter.create<memref::ExtractStridedMetadataOp>(loc, op.getBase());
+    OpFoldResult linearizedIndicesOfr;
+    std::tie(std::ignore, linearizedIndicesOfr) =
+        memref::getLinearizedMemRefOffsetAndSize(
+            rewriter, loc, srcBits, dstBits,
+            stridedMetadata.getConstifiedMixedOffset(),
+            stridedMetadata.getConstifiedMixedSizes(),
+            stridedMetadata.getConstifiedMixedStrides(),
+            getAsOpFoldResult(adaptor.getIndices()));
+    Value linearizedIndices =
+        getValueOrCreateConstantIndexOp(rewriter, loc, linearizedIndicesOfr);
+
+    // Load the whole data and use arith.select to handle the corner cases.
+    // E.g., given these input values:
+    //
+    //   %mask = [1, 1, 1, 0, 0, 0]
+    //   %0[%c0, %c0] contains [0x1, 0x2, 0x3, 0x4, 0x5, 0x6]
+    //   %value_to_store = [0x7, 0x8, 0x9, 0xA, 0xB, 0xC]
+    //
+    // we'll have
+    //
+    //    expected output: [0x7, 0x8, 0x9, 0x4, 0x5, 0x6]
+    //
+    //    %new_mask = [1, 1, 0]
+    //    %maskedload = [0x12, 0x34, 0x0]
+    //    %bitcast = [0x1, 0x2, 0x3, 0x4, 0x0, 0x0]
+    //    %select_using_original_mask = [0x7, 0x8, 0x9, 0x4, 0x0, 0x0]
+    //    %packed_data = [0x78, 0x94, 0x0, 0x0]
+    //
+    // Using the new mask to store %packed_data results in expected output.
+    FailureOr<Operation *> newMask =
+        getCompressedMaskOp(rewriter, loc, op.getMask(), origElements, scale);
+    if (failed(newMask))
+      return failure();
+
+    auto numElements = (origElements + scale - 1) / scale;
+    auto newType = VectorType::get(numElements, newElementType);
+    auto passThru = rewriter.create<arith::ConstantOp>(
+        loc, newType, rewriter.getZeroAttr(newType));
+
+    auto newLoad = rewriter.create<vector::MaskedLoadOp>(
+        loc, newType, adaptor.getBase(), linearizedIndices,
+        newMask.value()->getResult(0), passThru);
+
+    Value valueToStore = rewriter.create<vector::BitCastOp>(
+        loc, op.getValueToStore().getType(), newLoad);
+    valueToStore = rewriter.create<arith::SelectOp>(
+        loc, op.getMask(), op.getValueToStore(), valueToStore);
+    valueToStore =
+        rewriter.create<vector::BitCastOp>(loc, newType, valueToStore);
+
+    rewriter.replaceOpWithNewOp<vector::MaskedStoreOp>(
+        op, adaptor.getBase(), linearizedIndices, newMask.value()->getResult(0),
+        valueToStore);
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // ConvertVectorLoad
 //===----------------------------------------------------------------------===//
@@ -236,7 +396,6 @@ struct ConvertVectorMaskedLoad final
     // TODO: Currently, only the even number of elements loading is supported.
     // To deal with the odd number of elements, one has to extract the
     // subvector at the proper offset after bit-casting.
-
     auto origType = op.getVectorType();
     auto origElements = origType.getNumElements();
     if (origElements % scale != 0)
@@ -244,7 +403,6 @@ struct ConvertVectorMaskedLoad final
 
     auto stridedMetadata =
         rewriter.create<memref::ExtractStridedMetadataOp>(loc, op.getBase());
-
     OpFoldResult linearizedIndices;
     std::tie(std::ignore, linearizedIndices) =
         memref::getLinearizedMemRefOffsetAndSize(
@@ -254,66 +412,13 @@ struct ConvertVectorMaskedLoad final
             stridedMetadata.getConstifiedMixedStrides(),
             getAsOpFoldResult(adaptor.getIndices()));
 
-    auto numElements = (origElements + scale - 1) / scale;
-    auto newType = VectorType::get(numElements, newElementType);
-
-    auto maskOp = op.getMask().getDefiningOp();
-    SmallVector<vector::ExtractOp, 2> extractOps;
-    // Finding the mask creation operation.
-    while (maskOp &&
-           !isa<vector::CreateMaskOp, vector::ConstantMaskOp>(maskOp)) {
-      if (auto extractOp = dyn_cast<vector::ExtractOp>(maskOp)) {
-        maskOp = extractOp.getVector().getDefiningOp();
-        extractOps.push_back(extractOp);
-      }
-    }
-    auto createMaskOp = dyn_cast_or_null<vector::CreateMaskOp>(maskOp);
-    auto constantMaskOp = dyn_cast_or_null<vector::ConstantMaskOp>(maskOp);
-    if (!createMaskOp && !constantMaskOp)
+    FailureOr<Operation *> newMask =
+        getCompressedMaskOp(rewriter, loc, op.getMask(), origElements, scale);
+    if (failed(newMask))
       return failure();
 
-    // Computing the "compressed" mask. All the emulation logic (i.e. computing
-    // new mask index) only happens on the last dimension of the vectors.
-    Operation *newMask = nullptr;
-    auto shape = llvm::to_vector(
-        maskOp->getResultTypes()[0].cast<VectorType>().getShape().drop_back());
-    shape.push_back(numElements);
-    auto newMaskType = VectorType::get(shape, rewriter.getI1Type());
-    if (createMaskOp) {
-      auto maskOperands = createMaskOp.getOperands();
-      auto numMaskOperands = maskOperands.size();
-      AffineExpr s0;
-      bindSymbols(rewriter.getContext(), s0);
-      s0 = s0 + scale - 1;
-      s0 = s0.floorDiv(scale);
-      OpFoldResult origIndex =
-          getAsOpFoldResult(maskOperands[numMaskOperands - 1]);
-      OpFoldResult maskIndex =
-          affine::makeComposedFoldedAffineApply(rewriter, loc, s0, origIndex);
-      auto newMaskOperands = llvm::to_vector(maskOperands.drop_back());
-      newMaskOperands.push_back(
-          getValueOrCreateConstantIndexOp(rewriter, loc, maskIndex));
-      newMask = rewriter.create<vector::CreateMaskOp>(loc, newMaskType,
-                                                      newMaskOperands);
-    } else if (constantMaskOp) {
-      auto maskDimSizes = constantMaskOp.getMaskDimSizes().getValue();
-      auto numMaskOperands = maskDimSizes.size();
-      auto origIndex =
-          cast<IntegerAttr>(maskDimSizes[numMaskOperands - 1]).getInt();
-      auto maskIndex =
-          rewriter.getI64IntegerAttr((origIndex + scale - 1) / scale);
-      auto newMaskDimSizes = llvm::to_vector(maskDimSizes.drop_back());
-      newMaskDimSizes.push_back(maskIndex);
-      newMask = rewriter.create<vector::ConstantMaskOp>(
-          loc, newMaskType, rewriter.getArrayAttr(newMaskDimSizes));
-    }
-
-    while (!extractOps.empty()) {
-      newMask = rewriter.create<vector::ExtractOp>(
-          loc, newMask->getResults()[0], extractOps.back().getMixedPosition());
-      extractOps.pop_back();
-    }
-
+    auto numElements = (origElements + scale - 1) / scale;
+    auto newType = VectorType::get(numElements, newElementType);
     auto newPassThru =
         rewriter.create<vector::BitCastOp>(loc, newType, op.getPassThru());
 
@@ -321,7 +426,7 @@ struct ConvertVectorMaskedLoad final
     auto newLoad = rewriter.create<vector::MaskedLoadOp>(
         loc, newType, adaptor.getBase(),
         getValueOrCreateConstantIndexOp(rewriter, loc, linearizedIndices),
-        newMask->getResult(0), newPassThru);
+        newMask.value()->getResult(0), newPassThru);
 
     // Setting the part that originally was not effectively loaded from memory
     // to pass through.
@@ -821,7 +926,8 @@ void vector::populateVectorNarrowTypeEmulationPatterns(
 
   // Populate `vector.*` conversion patterns.
   patterns.add<ConvertVectorLoad, ConvertVectorMaskedLoad, ConvertVectorStore,
-               ConvertVectorTransferRead>(typeConverter, patterns.getContext());
+               ConvertVectorMaskedStore, ConvertVectorTransferRead>(
+      typeConverter, patterns.getContext());
 }
 
 void vector::populateVectorNarrowTypeRewritePatterns(
diff --git a/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir b/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir
index af0f98a1c447de0..cba299b2a1d9567 100644
--- a/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir
+++ b/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir
@@ -428,3 +428,75 @@ func.func @vector_store_i4_dynamic(%arg0: vector<8xi4>, %arg1: index, %arg2: ind
 //      CHECK32: %[[INDEX:.+]] = affine.apply #[[MAP1]]()[%[[ARG3]], %[[ARG2]], %[[ARG4]]]
 //      CHECK32: %[[VEC_I8:.+]] = vector.bitcast %[[ARG0]] : vector<8xi4> to vector<1xi32>
 //      CHECK32: vector.store %[[VEC_I8:.+]], %[[ALLOC:.+]][%[[INDEX:.+]]] : memref<?xi32>, vector<1xi32>
+
+// -----
+
+func.func @vector_maskedstore_i8(%arg0: index, %arg1: index, %arg2: index, %value: vector<8xi8>) {
+  %0 = memref.alloc() : memref<3x8xi8>
+  %mask = vector.create_mask %arg2 : vector<8xi1>
+  vector.maskedstore %0[%arg0, %arg1], %mask, %value : memref<3x8xi8>, vector<8xi1>, vector<8xi8>
+  return
+}
+// Expect no conversions, i8 is supported.
+//      CHECK: func @vector_maskedstore_i8(
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]
+// CHECK-SAME:     %[[VAL:[a-zA-Z0-9]+]]
+// CHECK-NEXT:   %[[ALLOC:.+]] = memref.alloc() : memref<3x8xi8>
+// CHECK-NEXT:   %[[MASK:.+]] = vector.create_mask %[[ARG2]] : vector<8xi1>
+// CHECK-NEXT:   vector.maskedstore %[[ALLOC]][%[[ARG0]], %[[ARG1]]], %[[MASK]], %[[VAL]]
+// CHECK-NEXT:   return
+
+// CHECK32-DAG: #[[LOAD_IDX_MAP:.+]] = affine_map<()[s0, s1] -> (s0 * 2 + s1 floordiv 4)>
+// CHECK32-DAG: #[[MASK_IDX_MAP:.+]] = affine_map<()[s0] -> ((s0 + 3) floordiv 4)>
+// CHECK32:     func @vector_maskedstore_i8(
+// CHECK32-SAME:     %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK32-SAME:     %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK32-SAME:     %[[ARG2:[a-zA-Z0-9]+]]
+// CHECK32-SAME:     %[[VAL:[a-zA-Z0-9]+]]
+// CHECK32:        %[[ALLOC:.+]] = memref.alloc() : memref<6xi32>
+// CHECK32:        %[[ORIG_MASK:.+]] = vector.create_mask %[[ARG2]] : vector<8xi1>
+// CHECK32:        %[[LIDX:.+]] = affine.apply #[[LOAD_IDX_MAP]]()[%[[ARG0]], %[[ARG1]]]
+// CHECK32:        %[[MASK_IDX:.+]] = affine.apply #[[MASK_IDX_MAP]]()[%[[ARG2]]]
+// CHECK32:        %[[NEW_MASK:.+]] = vector.create_mask %[[MASK_IDX]] : vector<2xi1>
+// CHECK32:        %[[PASS_THRU:.+]] = arith.constant dense<0> : vector<2xi32>
+// CHECK32:        %[[LOAD:.+]] = vector.maskedload %[[ALLOC]][%[[LIDX]]], %[[NEW_MASK]], %[[PASS_THRU]]
+// CHECK32:        %[[BITCAST:.+]] = vector.bitcast %[[LOAD]] : vector<2xi32> to vector<8xi8>
+// CHECK32:        %[[SELECT:.+]] = arith.select %[[ORIG_MASK]], %[[VAL]], %[[BITCAST]] : vector<8xi1>, vector<8xi8>
+// CHECK32:        %[[NEW_VAL:.+]] = vector.bitcast %[[SELECT]] : vector<8xi8> to vector<2xi32>
+// CHECK32:        vector.maskedstore %[[ALLOC]][%[[LIDX]]], %[[NEW_MASK]], %[[NEW_VAL]]
+
+// -----
+
+func.func @vector_cst_maskedstore_i8(%arg0: index, %arg1: index, %value: vector<8xi8>) {
+  %0 = memref.alloc() : memref<3x8xi8>
+  %mask = vector.constant_mask [4] : vector<8xi1>
+  vector.maskedstore %0[%arg0, %arg1], %mask, %value : memref<3x8xi8>, vector<8xi1>, vector<8xi8>
+  return
+}
+// Expect no conversions, i8 is supported.
+//      CHECK: func @vector_cst_maskedstore_i8(
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK-SAME:     %[[VAL:[a-zA-Z0-9]+]]
+// CHECK-NEXT:   %[[ALLOC:.+]] = memref.alloc() : memref<3x8xi8>
+// CHECK-NEXT:   %[[MASK:.+]] = vector.constant_mask [4] : vector<8xi1>
+// CHECK-NEXT:   vector.maskedstore %[[ALLOC]][%[[ARG0]], %[[ARG1]]], %[[MASK]], %[[VAL]]
+// CHECK-NEXT:   return
+
+// CHECK32-DAG: #[[LOAD_IDX_MAP:.+]] = affine_map<()[s0, s1] -> (s0 * 2 + s1 floordiv 4)>
+// CHECK32:     func @vector_cst_maskedstore_i8(
+// CHECK32-SAME:     %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK32-SAME:     %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK32-SAME:     %[[VAL:[a-zA-Z0-9]+]]
+// CHECK32:        %[[ALLOC:.+]] = memref.alloc() : memref<6xi32>
+// CHECK32:        %[[ORIG_MASK:.+]] = vector.constant_mask [4] : vector<8xi1>
+// CHECK32:        %[[LIDX:.+]] = affine.apply #[[LOAD_IDX_MAP]]()[%[[ARG0]], %[[ARG1]]]
+// CHECK32:        %[[NEW_MASK:.+]] = vector.constant_mask [1] : vector<2xi1>
+// CHECK32:        %[[PASS_THRU:.+]] = arith.constant dense<0> : vector<2xi32>
+// CHECK32:        %[[LOAD:.+]] = vector.maskedload %[[ALLOC]][%[[LIDX]]], %[[NEW_MASK]], %[[PASS_THRU]]
+// CHECK32:        %[[BITCAST:.+]] = vector.bitcast %[[LOAD]] : vector<2xi32> to vector<8xi8>
+// CHECK32:        %[[SELECT:.+]] = arith.select %[[ORIG_MASK]], %[[VAL]], %[[BITCAST]] : vector<8xi1>, vector<8xi8>
+// CHECK32:        %[[NEW_VAL:.+]] = vector.bitcast %[[SELECT]] : vector<8xi8> to vector<2xi32>
+// CHECK32:        vector.maskedstore %[[ALLOC]][%[[LIDX]]], %[[NEW_MASK]], %[[NEW_VAL]]

[dcaballe]

LGTM, thanks! Just minor comments

[Max191]

I think this can still result in a race condition like in memref.store emulation, e.g., if you have the following stores running on different threads:

//    %memref = [0x12, 0x34, 0x56]
// Masked store
//    %new_mask = [1, 1, 0]
//    %maskedload = [0x12, 0x34, 0x0]
//    %bitcast = [0x1, 0x2, 0x3, 0x4, 0x0, 0x0]
//    %select_using_original_mask = [0x7, 0x8, 0x9, 0x4, 0x0, 0x0]
//    %packed_data = [0x78, 0x94, 0x00]
//    vector.maskedstore %memref, %new_mask, %packed_data
// Other store
//    %rmw_mask = [0xF0]
//    %rmw_val = [0x0D]
//    memref.atomic_rmw andi %rmw_mask, %memref[1]
//    memref.atomic_rmw ori %rmw_val, %memref[1]

If the memref.atomic_rmw ops happen after the masked_load, but before the masked_store, then the masked_store will overwrite what was written by the atomic_rmw ops.

A potential solution to this race condition would be to split off the corner cases from the masked store, and rewrite them the same way as memref.store emulation (i.e. with two atomic_rmw ops like above). However, I don't think this would be a very common occurrence, since masked_store would mostly be used for tiling with masking, but it is potentially possible. Maybe at least a TODO here would be warranted.

[hanhanW]

I'm not very sure if this is the case. I need to think more about it. The indices for store ops are not overlapping with each other when we distribute the work to multi-threads. The race condition issue happens when store ops index on the same "byte". After the emulation, different memref.store ops could index on the same byte. In this context, we need atomic ops.

For vector.maskedstore and vector.store, the whole bytes are taken into account during emulation. So the "index ranges" of each thread does not overlap with others. In this context, we do'nt need atomic ops.

Does it make sense?

[Max191]

Maybe we would never distribute this way, but it is possible to have two threads with non-overlapping indices that could result in IR with a similar problem to the above. For example, if there is a tensor of shape 6xi4, and the work was distributed into 2 threads, storing into the first 3xi4 values and the second 3xi4 values respectively. Then these could be lowered into 2 vector.maskedstore/vector.store ops that overlap on the middle byte after narrow type emulation, since the 6xi4 would become 3xi8.

This is probably a moot point, though, because I don't think we would distribute into non powers of 2 in this way.

[hanhanW]

it is possible to have two threads with non-overlapping indices that could result in IR with a similar problem to the above.

That already has race condition issue, and it is undefined-behavior. It is not introduced by emulation.

The example you provided is reasonable to me! We bail out the case out in line 202:

    int scale = dstBits / srcBits;
    int origElements = op.getValueToStore().getType().getNumElements();
    if (origElements % scale != 0)
      return failure();

I agree that we might need further support for the case, good catch!

It is doing correct emulation under these assumptions and checks, so I'm going to land the PR.

EDIT: I keep the conversation open because the discussion is useful.