[VectorExt] Add TransferScatterOp

keshavvinayak01 · keshavvinayak01 · commit 0adaf2ef92db · 2026-03-06T20:07:51.000Z
Signed-off-by: Keshav Vinayak Jha &lt;keshavvinayakjha@gmail.com&gt;
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.cpp b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.cpp
@@ -62,35 +62,35 @@ void TransferGatherOp::getEffects(
   }
 }
 
-// Verifier.
-
-LogicalResult TransferGatherOp::verify() {
-  OperandRange indexVecs = getIndexVecs();
-  TypedValue<VectorType> vector = getVector();
-  Value mask = getMask();
-  SmallVector<AffineMap> indexingMaps = getIndexingMapsArray();
+// Shared verifier for TransferGatherOp and TransferScatterOp.
 
+static LogicalResult
+verifyTransferGatherScatterLikeOp(Operation *op, VectorType vectorType,
+                                  OperandRange indexVecs, Value mask,
+                                  ArrayRef<AffineMap> indexingMaps) {
   // Check that we have the correct number of indexing maps.
   int64_t expectedNumIndexingMaps =
-      /*sourceIndexingMap=*/1 + /*indexVecIndexingMaps=*/indexVecs.size() +
+      /*baseIndexingMap=*/1 + /*indexVecIndexingMaps=*/indexVecs.size() +
       /*maskIndexingMap=*/(mask ? 1 : 0);
   if (expectedNumIndexingMaps != static_cast<int64_t>(indexingMaps.size())) {
-    return emitOpError("expected ")
+    return op->emitOpError("expected ")
            << expectedNumIndexingMaps
            << " indexing maps, got: " << indexingMaps.size();
   }
 
-  int64_t vectorRank = vector.getType().getRank();
+  int64_t vectorRank = vectorType.getRank();
   int64_t indexSyms = indexVecs.size();
   for (AffineMap map : indexingMaps) {
     if (map.getNumDims() != vectorRank) {
-      return emitOpError("expected all indexing maps to have number of dims "
-                         "equal to vector rank. expected: ")
+      return op->emitOpError(
+                 "expected all indexing maps to have number of dims "
+                 "equal to vector rank. expected: ")
              << vectorRank << ", got: " << map.getNumDims() << " dims";
     }
     if (map.getNumSymbols() != indexSyms) {
-      return emitOpError("expected all indexing maps to have number of symbols "
-                         "equal to number of index vecs. expected: ")
+      return op->emitOpError(
+                 "expected all indexing maps to have number of symbols "
+                 "equal to number of index vecs. expected: ")
              << indexSyms << ", got: " << map.getNumSymbols() << " syms";
     }
     for (AffineExpr expr : map.getResults()) {
@@ -99,18 +99,18 @@ LogicalResult TransferGatherOp::verify() {
       }
       if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {
         if (constExpr.getValue() != 0) {
-          return emitOpError("expected constant 0 in indexing map, got: ")
+          return op->emitOpError("expected constant 0 in indexing map, got: ")
                  << constExpr.getValue();
         }
         continue;
       }
-      return emitOpError(
+      return op->emitOpError(
           "expected indexing map results to only be a dim, symbol, or 0");
     }
   }
 
   // Extra verification for index vecs.
-  ArrayRef<int64_t> vectorShape = vector.getType().getShape();
+  ArrayRef<int64_t> vectorShape = vectorType.getShape();
   ArrayRef<AffineMap> vectorIndexingMaps =
       ArrayRef(indexingMaps).slice(1, indexSyms);
   for (auto [i, map] : llvm::enumerate(vectorIndexingMaps)) {
@@ -119,23 +119,25 @@ LogicalResult TransferGatherOp::verify() {
       if (auto dim = dyn_cast<AffineDimExpr>(expr)) {
         expectedShape.push_back(vectorShape[dim.getPosition()]);
       } else {
-        return emitOpError(
+        return op->emitOpError(
             "expected vector indexing maps to not have any symbols");
       }
     }
     // Scalar index: map must have 0 results and type must be plain index.
     if (isa<IndexType>(indexVecs[i].getType())) {
       if (!expectedShape.empty()) {
-        return emitOpError("expected empty indexing map for scalar index vec "
-                           "at position ")
+        return op->emitOpError(
+                   "expected empty indexing map for scalar index vec "
+                   "at position ")
                << i;
       }
       continue;
     }
     ArrayRef<int64_t> actualShape =
         cast<VectorType>(indexVecs[i].getType()).getShape();
     if (ArrayRef<int64_t>(expectedShape) != actualShape) {
-      return emitOpError("Mismatched vector shape for index vec at position ")
+      return op->emitOpError(
+                 "Mismatched vector shape for index vec at position ")
              << i << ". Expected: [" << expectedShape << "]" << ", got: ["
              << actualShape << "]";
     }
@@ -149,13 +151,13 @@ LogicalResult TransferGatherOp::verify() {
       if (auto dim = dyn_cast<AffineDimExpr>(expr)) {
         expectedShape.push_back(vectorShape[dim.getPosition()]);
       } else {
-        return emitOpError(
+        return op->emitOpError(
             "expected mask indexing map to not have any symbols");
       }
     }
     ArrayRef<int64_t> actualShape = cast<VectorType>(mask.getType()).getShape();
     if (ArrayRef<int64_t>(expectedShape) != actualShape) {
-      return emitOpError("Mismatched mask shape")
+      return op->emitOpError("Mismatched mask shape")
              << ". Expected: [" << expectedShape << "]" << ", got: ["
              << actualShape << "]";
     }
@@ -164,6 +166,12 @@ LogicalResult TransferGatherOp::verify() {
   return success();
 }
 
+LogicalResult TransferGatherOp::verify() {
+  return verifyTransferGatherScatterLikeOp(
+      getOperation(), getVector().getType(), getIndexVecs(), getMask(),
+      getIndexingMapsArray());
+}
+
 // Fold and canonicalization helpers.
 
 static int64_t getVectorRank(Type type) {
@@ -663,6 +671,53 @@ void TransferGatherOp::getCanonicalizationPatterns(RewritePatternSet &results,
           ctx);
 }
 
+//===----------------------------------------------------------------------===//
+// TransferScatterOp
+//===----------------------------------------------------------------------===//
+
+Speculation::Speculatability TransferScatterOp::getSpeculatability() {
+  if (isa<RankedTensorType>(getBase().getType())) {
+    return Speculation::Speculatable;
+  }
+  return Speculation::NotSpeculatable;
+}
+
+void TransferScatterOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  if (isa<MemRefType>(getBase().getType())) {
+    effects.emplace_back(MemoryEffects::Read::get(), &getBaseMutable(),
+                         SideEffects::DefaultResource::get());
+    effects.emplace_back(MemoryEffects::Write::get(), &getBaseMutable(),
+                         SideEffects::DefaultResource::get());
+  }
+}
+
+LogicalResult TransferScatterOp::verify() {
+  if (failed(verifyTransferGatherScatterLikeOp(getOperation(), getVectorType(),
+                                               getIndexVecs(), getMask(),
+                                               getIndexingMapsArray()))) {
+    return failure();
+  }
+
+  // Verify result type matches base type for tensor semantics.
+  if (hasTensorSemantics()) {
+    if (!getResult()) {
+      return emitOpError("expected result for tensor operand");
+    }
+    if (getResult().getType() != getBase().getType()) {
+      return emitOpError("result type must match base type");
+    }
+  } else {
+    // Memref semantics: no result expected.
+    if (getResult()) {
+      return emitOpError("unexpected result for memref operand");
+    }
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.td b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.td
@@ -224,6 +224,127 @@ def IREEVectorExt_TransferGatherOp : IREEVectorExt_PureOp<"transfer_gather", [
   let hasVerifier = 1;
 }
 
+def IREEVectorExt_TransferScatterOp : IREEVectorExt_PureOp<"transfer_scatter", [
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+    DeclareOpInterfaceMethods<ConditionallySpeculatable>,
+    AttrSizedOperandSegments,
+    AllElementTypesMatch<["base", "vector"]>
+  ]> {
+  let arguments = (ins AnyShaped:$base,
+                       AnyVectorOfAnyRank:$vector,
+                       Variadic<Index>:$offsets,
+                       Variadic<AnyTypeOf<[Index, VectorOfAnyRankOf<[Index]>]>>:$index_vecs,
+                       AffineMapArrayAttr:$indexing_maps,
+                       Optional<VectorOfAnyRankOf<[I1]>>:$mask);
+  let results = (outs Optional<AnyShaped>:$result);
+
+  let summary = [{Scatters a supervector from an SSA vector value into a shaped destination.}];
+
+  let description = [{
+    The `transfer_scatter` operation is the write counterpart of
+    `transfer_gather`. It writes elements from a vector into a shaped
+    destination (memref or tensor), where each destination dimension can be
+    independently contiguous, scattered, or broadcast.
+
+    The scatter indices are expected to be unique. If multiple vector elements
+    map to the same destination location, there may be data races.
+
+    For tensor operands, the operation returns the modified tensor. For memref
+    operands, the operation has no result.
+
+    Example — scatter write: writing values into scattered rows of a 2D dest:
+
+    ```
+    // dest[indices[i], j] = vector[i, j]
+    %result = iree_vector_ext.transfer_scatter %vector into %dest[%c0, %c0]
+      [%indices : vector<16xindex>] {
+        indexing_maps = [
+          affine_map<(d0, d1)[s0] -> (s0, d1)>,
+          affine_map<(d0, d1)[s0] -> (d0)>
+        ]
+      } : vector<16x8xf16>, tensor<4096x8xf16>
+    ```
+
+    Semantically, for each position in the source vector:
+
+    ```
+    dest[offsets[0] + f0(d, s), offsets[1] + f1(d, s), ...] = vector[d0, d1, ...]
+    ```
+
+    where each `fi` is the i-th result of the dest indexing map evaluated at
+    the vector position `d = (d0, d1, ...)` and scattered index values
+    `s = (s0, s1, ...)`.
+
+    The `indexing_maps` attribute follows the same structure as
+    `transfer_gather`. See `transfer_gather` documentation for details. The
+    only difference is that Map 0 describes the destination indexing rather
+    than the source indexing.
+  }];
+
+  let extraClassDeclaration = [{
+    SmallVector<AffineMap> getIndexingMapsArray() {
+      return llvm::to_vector(getIndexingMaps().getAsValueRange<AffineMapAttr>());
+    }
+
+    AffineMap getDestIndexingMap() {
+      return getIndexingMapsArray().front();
+    }
+
+    SmallVector<AffineMap> getIndexVecIndexingMaps() {
+      auto maps = getIndexingMapsArray();
+      return SmallVector<AffineMap>(
+          maps.begin() + 1, maps.begin() + 1 + getIndexVecs().size());
+    }
+
+    std::optional<AffineMap> getMaskIndexingMap() {
+      if (!getMask()) return std::nullopt;
+      return getIndexingMapsArray().back();
+    }
+
+    VectorType getVectorType() {
+      return getVector().getType();
+    }
+
+    ShapedType getBaseType() {
+      return cast<ShapedType>(getBase().getType());
+    }
+
+    bool hasTensorSemantics() {
+      return isa<RankedTensorType>(getBase().getType());
+    }
+
+    /// Invert the dest indexing map to get a permutation map suitable for
+    /// vector.transfer_write. The dest map is
+    /// (vector_dims)[symbols] -> (dest_dims). This returns
+    /// (dest_dims) -> (vector_dims), where non-dim results (scattered
+    /// symbols, broadcast constants) become constant 0.
+    AffineMap getPermutationMap() {
+      AffineMap destMap = getDestIndexingMap();
+      MLIRContext *ctx = getContext();
+      SmallVector<AffineExpr> exprs(destMap.getNumDims(),
+                                    getAffineConstantExpr(0, ctx));
+      for (auto [i, expr] : llvm::enumerate(destMap.getResults())) {
+        if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
+          exprs[dimExpr.getPosition()] = getAffineDimExpr(i, ctx);
+        }
+      }
+      return AffineMap::get(destMap.getNumResults(), /*symbolCount=*/0,
+                            exprs, ctx);
+    }
+  }];
+
+  let assemblyFormat = [{
+    $vector `into` $base `[` $offsets `]`
+    (`[` $index_vecs^ `:` type($index_vecs) `]`)?
+    (`,` $mask^)?
+    attr-dict `:` type($vector) `,` type($base)
+    (`,` type($mask)^)?
+    (`->` type($result)^)?
+  }];
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Terminator ops.
 //===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/test/invalid.mlir b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/test/invalid.mlir
@@ -85,6 +85,59 @@ func.func @index_vec_shape_mismatch(%indices: vector<128x64xindex>,
 
 // -----
 
+func.func @scatter_wrong_num_indexing_maps(%indices: vector<128xindex>,
+  %vector: vector<128xf16>,
+  %dest: tensor<128xf16>)
+  -> tensor<128xf16> {
+
+  %c0 = arith.constant 0 : index
+
+  // expected-error @+1 {{'iree_vector_ext.transfer_scatter' op expected 2 indexing maps, got: 1}}
+  %out = iree_vector_ext.transfer_scatter %vector into %dest[%c0]
+  [%indices : vector<128xindex>] {
+    indexing_maps = [affine_map<(d0)[s0] -> (s0)>]
+  } : vector<128xf16>, tensor<128xf16> -> tensor<128xf16>
+
+  return %out : tensor<128xf16>
+}
+
+// -----
+
+func.func @scatter_index_vec_shape_mismatch(%indices: vector<128x64xindex>,
+  %vector: vector<128x64xf16>,
+  %dest: tensor<128x64xf16>)
+  -> tensor<128x64xf16> {
+
+  %c0 = arith.constant 0 : index
+
+  // expected-error @+1 {{'iree_vector_ext.transfer_scatter' op Mismatched vector shape for index vec at position 0. Expected: [64, 128], got: [128, 64]}}
+  %out = iree_vector_ext.transfer_scatter %vector into %dest[%c0, %c0]
+  [%indices : vector<128x64xindex>] {
+    indexing_maps = [affine_map<(d0, d1)[s0] -> (d0, s0)>,
+                     affine_map<(d0, d1)[s0] -> (d1, d0)>]
+  } : vector<128x64xf16>, tensor<128x64xf16> -> tensor<128x64xf16>
+
+  return %out : tensor<128x64xf16>
+}
+
+// -----
+
+func.func @scatter_memref_with_result(%vector: vector<128xf16>,
+  %dest: memref<128xf16>)
+  -> memref<128xf16> {
+
+  %c0 = arith.constant 0 : index
+
+  // expected-error @+1 {{'iree_vector_ext.transfer_scatter' op unexpected result for memref operand}}
+  %out = iree_vector_ext.transfer_scatter %vector into %dest[%c0] {
+    indexing_maps = [affine_map<(d0) -> (d0)>]
+  } : vector<128xf16>, memref<128xf16> -> memref<128xf16>
+
+  return %out : memref<128xf16>
+}
+
+// -----
+
 func.func @arg_compare_dimension_out_of_bounds(%input: vector<4x128xf32>,
                                                %out_val: vector<4xf32>,
                                                %out_idx: vector<4xi32>)
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/test/roundtrip.mlir b/compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/test/roundtrip.mlir
