[mlir][Vector] Add support for Value indices to vector.extract/insert

[dcaballe]

Summary

vector.extract/insert ops only support constant indices. This PR is
extending them so that arbitrary values can be used instead.

This work is part of the RFC: https://discourse.llvm.org/t/rfc-psa-remove-vector-extractelement-and-vector-insertelement-ops-in-favor-of-vector-extract-and-vector-insert-ops

[dcaballe]

Fix bugs

[jsetoain]

Hi Diego,

Thanks for this. As a start, I'll take care of the tests I see failing under Dialect/Vector:

MLIR :: Dialect/Vector/canonicalize.mlir
MLIR :: Dialect/Vector/constant-fold.mlir
MLIR :: Dialect/Vector/invalid.mlir
MLIR :: Dialect/Vector/ops.mlir
MLIR :: Dialect/Vector/scalar-vector-transfer-to-memref.mlir
MLIR :: Dialect/Vector/vector-broadcast-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-matvec-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-dot-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir
MLIR :: Dialect/Vector/vector-dropleadunitdim-transforms.mlir
MLIR :: Dialect/Vector/vector-emulate-narrow-type.mlir
MLIR :: Dialect/Vector/vector-extract-strided-slice-lowering.mlir
MLIR :: Dialect/Vector/vector-gather-lowering.mlir
MLIR :: Dialect/Vector/vector-mask-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-multi-reduction-lowering.mlir
MLIR :: Dialect/Vector/vector-multi-reduction-outer-lowering.mlir
MLIR :: Dialect/Vector/vector-outerproduct-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-scan-transforms.mlir
MLIR :: Dialect/Vector/vector-shape-cast-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-transfer-to-vector-load-store.mlir
MLIR :: Dialect/Vector/vector-transforms.mlir
MLIR :: Dialect/Vector/vector-transpose-lowering.mlir
MLIR :: Dialect/Vector/vector-warp-distribute.mlir

[awarzynski]

Hey Diego, happy to help! I can update the Vector integration tests:

MLIR :: Integration/Dialect/Vector/CPU/test-broadcast.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-compress.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-contraction.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-expand.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-extract-strided-slice.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-flat-transpose-col.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-flat-transpose-row.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-gather.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-index-vectors.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-maskedload.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-maskedstore.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-matrix-multiply-col.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-matrix-multiply-row.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-outerproduct-f32.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-outerproduct-i64.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f32-reassoc.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f32.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f64-reassoc.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f64.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-i32.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-i64.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-scan.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-scatter.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-shape-cast.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-sparse-dot-matvec.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-sparse-saxpy-jagged-matvec.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-read-1d.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-read-2d.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-read-3d.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-to-loops.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-write.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transpose.mlir

-Andrzej

[jsetoain]

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

722

Why not use AnySignlessIntegerOrIndex here?

[jsetoain]

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

722

I chose the wrong type. What I meant is, why not allow for attributes as well as values? That should get rid of the issues with not accepting things like [3, 3, 3] as indices, no?

[dcaballe]

Thanks a lot for the help! Please, before you spend more time on this, please, give me some time to try mixing attributes and values, as @jsetoain suggested. We can do something similar to what tensor.extract_slice/insert_slice do.

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

722

I made that decision based on other ops like tensor.extract to be consistent. Using only the index type also simplifies things. We don't have to specify the type use for the indices in the assembly format, we don't have to verify that all the indices have the same type, we don't have to check for the type in transformations that are manipulating/replacing/adding more indices, etc.

Regarding mixing attributes and values, that's actually a very good point. I see that tensor.extract_slice is actually doing something like that:

let arguments = (ins
  AnyRankedTensor:$source,
  Variadic<Index>:$offsets,
  Variadic<Index>:$sizes,
  Variadic<Index>:$strides,
  DenseI64ArrayAttr:$static_offsets,
  DenseI64ArrayAttr:$static_sizes,
  DenseI64ArrayAttr:$static_strides
);

Let me give it a try!

[dcaballe]

Progress... Not ready yet!

[dcaballe]

Final rework

[dcaballe]

This should be ready to review and land. I added support for values and constant attributes to extract and insert ops as requested.

[c-rhodes]

Thanks for this work Diego! this will useful for ArmSME. We have ops to load/store slices of a tile (2-D scalable vector) and this support should enable us to lower vector.insert / vector.extract to these ops. #66691 adds support to lower vector.print to ArmSME, this could be done generically in VectorToSCF once this lands, after which we could add our target specific lowerings :) Also be useful for places where we can't fully unroll.

I've left a few comments/questions, mostly minor. I also noticed there's no tests, it would be good to have some in mlir/test/Dialect/Vector/ops.mlir. I also tried lowering a simple example to LLVM but hit an assert, is this expected to work yet? Example below

; mlir-opt -convert-vector-to-llvm

func.func @extract_element_from_vec_2d(%arg0: vector<2x16xf32>) -> vector<16xf32> {
  %c1 = arith.constant 1 : index
  %0 = vector.extract %arg0[%c1]: vector<2x16xf32>
  return %0 : vector<16xf32>
}

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

1246–1247

nit: formatting, these can go on the same line

1478–1480

move above valueToExtractFrom

mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp

98–99

remove cast?

124–125

remove cast?

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

1047–1048

it's not clear to me why is this an assert? (And below for insert)

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

711

spelling: Unexpected

[awarzynski]

+1 to adding some tests in mlir/test/Dialect/Vector/ops.mlir. Otherwise a few nits. Thanks a ton for resurrecting this work :)

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

729

Could you add examples with static sizes?

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

38–39

This comment is out of date :)

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

226–227

[nit] There's an identical comment in the header file - I'd only keep one (otherwise they might get out of sync).

1395–1396

What about dynamic positions? If unsupported, then perhaps add an assert?

[dcaballe]

Addressed feedback

[dcaballe]

Thanks for the feedback! I only addressed @c-rhodes comments for now. I'll take a look at the remaining ones tomorrow.

In D155034#4648233, @c-rhodes wrote:

I also tried lowering a simple example to LLVM but hit an assert, is this expected to work yet? Example below

Sorry about that. My main goal with this patch was about having the representation working at vector level (and adjust all the impacted tests, which were a lot at the beginning) but not so much about adding support for the new value-based functionality end to end. I just added support for 1-D cases to the LLVM conversion to demonstrate that it works but there are quite a few TODOs, esp. wrt canonicalizations. Hopefully that makes sense!

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

1047–1048

Replaced with failure. It's more like a personal preference: I usually lean towards asserting instead of just silently bailing out, esp. when I don't know the pass or the consequences of bailing out very well.

[c-rhodes]

In D155034#4648233, @c-rhodes wrote:

I also tried lowering a simple example to LLVM but hit an assert, is this expected to work yet? Example below

Sorry about that. My main goal with this patch was about having the representation working at vector level (and adjust all the impacted tests, which were a lot at the beginning) but not so much about adding support for the new value-based functionality end to end. I just added support for 1-D cases to the LLVM conversion to demonstrate that it works but there are quite a few TODOs, esp. wrt canonicalizations. Hopefully that makes sense!

No worries, thanks for clarifying

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

729

could you also add an example with mixed static/dynamic indices?

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

1047–1048

Replaced with failure. It's more like a personal preference: I usually lean towards asserting instead of just silently bailing out, esp. when I don't know the pass or the consequences of bailing out very well.

That's fair, an assert to me indicates a bug so I was wondering why a rank-1 extract with dynamic index is considered a bug in this particular rewrite, but I understand this change touches a lot of code and it's difficult to understand the implication for all scenarios.

mlir/test/Dialect/Vector/ops.mlir

223 ↗

(On Diff #557091)

test with mixed static/dynamic indices?

[dcaballe]

Addressed feedback

[dcaballe]

I think I addressed everything. Thanks!

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

226–227

I think that's common practice in MLIR. It has been pointed out to me that we should actually add doc to both decl and def. Extra doc is not harmful, I guess.

1395–1396

We can't add an assert here because this constructor is used by all kind of extract ops but the checks below should gracefully fail when they have dynamic operands.

[c-rhodes]

LGTM cheers

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

590

%5

mlir/include/mlir/Dialect/Vector/IR/VectorOps.h

 //===- VectorOps.h - MLIR Vector Dialect Operations -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the Vector dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_VECTOR_IR_VECTOROPS_H
 #define MLIR_DIALECT_VECTOR_IR_VECTOROPS_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.h"
 #include "mlir/Dialect/Vector/Interfaces/MaskingOpInterface.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
inline bool isParallelIterator(Attribute attr) {
   return cast<IteratorTypeAttr>(attr).getValue() == IteratorType::parallel;
 }
 
 /// Returns true if `attr` has "reduction" iterator type semantics.
 inline bool isReductionIterator(Attribute attr) {
   return cast<IteratorTypeAttr>(attr).getValue() == IteratorType::reduction;
 }
 
+/// Returns the integer numbers in `values`. `values` are expected to be
+/// constant operations.
+SmallVector<int64_t> getAsIntegers(ArrayRef<Value> values);
+
+/// Returns the integer numbers in `foldResults`. `foldResults` are expected to
+/// be constant operations.
+SmallVector<int64_t> getAsIntegers(ArrayRef<OpFoldResult> foldResults);
+
+/// Convert `foldResults` into Values. Integer attributes are converted to
+/// constant op.
+SmallVector<Value> getAsValues(OpBuilder &builder, Location loc,
+                               ArrayRef<OpFoldResult> foldResults);
+
+/// Returns the constant index ops in `values`. `values` are expected to be
+/// constant operations.
+SmallVector<arith::ConstantIndexOp>
+getAsConstantIndexOps(ArrayRef<Value> values);
+
 //===----------------------------------------------------------------------===//
 // Vector Masking Utilities
 //===----------------------------------------------------------------------===//
 
 /// Create the vector.yield-ended region of a vector.mask op with `maskableOp`
 /// as masked operation.
 void createMaskOpRegion(OpBuilder &builder, Operation *maskableOp);
 

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

def Vector_ExtractElementOp :
   let hasVerifier = 1;
   let hasFolder = 1;
 }
 
 def Vector_ExtractOp :
   Vector_Op<"extract", [Pure,
      PredOpTrait<"operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
-     InferTypeOpAdaptorWithIsCompatible]>,
-    Arguments<(ins AnyVectorOfAnyRank:$vector, DenseI64ArrayAttr:$position)>,
-    Results<(outs AnyType)> {
+     InferTypeOpAdaptorWithIsCompatible]> {
   let summary = "extract operation";
   let description = [{
     Takes an n-D vector and a k-D position and extracts the (n-k)-D vector at
     the proper position. Degenerates to an element type if n-k is zero.
 
     Example:
 
     ```mlir
     %1 = vector.extract %0[3]: vector<4x8x16xf32>
     %2 = vector.extract %0[3, 3, 3]: vector<4x8x16xf32>
     %3 = vector.extract %1[]: vector<f32>
     ```
   }];
+
+  let arguments = (ins
+    AnyVectorOfAnyRank:$vector,
+    Variadic<Index>:$dynamic_position,
+    DenseI64ArrayAttr:$static_position
+  );
+  let results = (outs AnyType:$result);
+
   let builders = [
-    // Convenience builder which assumes the values in `position` are defined by
-    // ConstantIndexOp.
-    OpBuilder<(ins "Value":$source, "ValueRange":$position)>
+    OpBuilder<(ins "Value":$source, "int64_t":$position)>,
+    OpBuilder<(ins "Value":$source, "OpFoldResult":$position)>,
+    OpBuilder<(ins "Value":$source, "ArrayRef<int64_t>":$position)>,
+    OpBuilder<(ins "Value":$source, "ArrayRef<OpFoldResult>":$position)>,
   ];
+
   let extraClassDeclaration = [{
     VectorType getSourceVectorType() {
       return ::llvm::cast<VectorType>(getVector().getType());
     }
+
+    /// Return a vector with all the static and dynamic position indices.
+    SmallVector<OpFoldResult> getMixedPosition() {
+      OpBuilder builder(getContext());
+      return getMixedValues(getStaticPosition(), getDynamicPosition(), builder);
+    }
+
+    unsigned getNumIndices() {
+      return getStaticPosition().size();
+    }
+
+    bool hasDynamicPosition() {
+      auto dynPos = getDynamicPosition();
+      return std::any_of(dynPos.begin(), dynPos.end(),
+                         [](Value operand) { return operand != nullptr; });
+    }
   }];
-  let assemblyFormat = "$vector `` $position attr-dict `:` type($vector)";
+
+  let assemblyFormat = [{
+    $vector ``
+    custom<DynamicIndexList>($dynamic_position, $static_position)
+    attr-dict `:` type($vector)
+  }];
+
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def Vector_FMAOp :
   Op<Vector_Dialect, "fma", [
        Pure, AllTypesMatch<["lhs", "rhs", "acc", "result"]>,
def Vector_InsertElementOp :
   let hasVerifier = 1;
   let hasFolder = 1;
 }
 
 def Vector_InsertOp :
   Vector_Op<"insert", [Pure,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
-     AllTypesMatch<["dest", "res"]>]>,
-     Arguments<(ins AnyType:$source, AnyVectorOfAnyRank:$dest, DenseI64ArrayAttr:$position)>,
-     Results<(outs AnyVectorOfAnyRank:$res)> {
+     AllTypesMatch<["dest", "result"]>]> {
   let summary = "insert operation";
   let description = [{
     Takes an n-D source vector, an (n+k)-D destination vector and a k-D position
     and inserts the n-D source into the (n+k)-D destination at the proper
     position. Degenerates to a scalar or a 0-d vector source type when n = 0.
 
     Example:
 
     ```mlir
-    %2 = vector.insert %0, %1[3] : vector<8x16xf32> into vector<4x8x16xf32>
-    %5 = vector.insert %3, %4[3, 3, 3] : f32 into vector<4x8x16xf32>
+    %2 = vector.insert %0, %1[%c3] : vector<8x16xf32> into vector<4x8x16xf32>
+    %5 = vector.insert %3, %4[%c3, %c3, %c3] : f32 into vector<4x8x16xf32>
     %8 = vector.insert %6, %7[] : f32 into vector<f32>
-    %11 = vector.insert %9, %10[3, 3, 3] : vector<f32> into vector<4x8x16xf32>
+    %11 = vector.insert %9, %10[%c3, %c3, %c3] : vector<f32> into vector<4x8x16xf32>
     ```
   }];
-  let assemblyFormat = [{
-    $source `,` $dest $position attr-dict `:` type($source) `into` type($dest)
-  }];
+
+  let arguments = (ins
+    AnyType:$source,
+    AnyVectorOfAnyRank:$dest,
+    Variadic<Index>:$dynamic_position,
+    DenseI64ArrayAttr:$static_position
+  );
+  let results = (outs AnyVectorOfAnyRank:$result);
 
   let builders = [
-    // Convenience builder which assumes all values are constant indices.
-    OpBuilder<(ins "Value":$source, "Value":$dest, "ValueRange":$position)>
+    OpBuilder<(ins "Value":$source, "Value":$dest, "int64_t":$position)>,
+    OpBuilder<(ins "Value":$source, "Value":$dest, "OpFoldResult":$position)>,
+    OpBuilder<(ins "Value":$source, "Value":$dest, "ArrayRef<int64_t>":$position)>,
+    OpBuilder<(ins "Value":$source, "Value":$dest, "ArrayRef<OpFoldResult>":$position)>,
   ];
+
   let extraClassDeclaration = [{
     Type getSourceType() { return getSource().getType(); }
     VectorType getDestVectorType() {
       return ::llvm::cast<VectorType>(getDest().getType());
     }
+
+    /// Return a vector with all the static and dynamic position indices.
+    SmallVector<OpFoldResult> getMixedPosition() {
+      OpBuilder builder(getContext());
+      return getMixedValues(getStaticPosition(), getDynamicPosition(), builder);
+    }
+
+    unsigned getNumIndices() {
+      return getStaticPosition().size();
+    }
+
+    bool hasDynamicPosition() {
+      return llvm::any_of(getDynamicPosition(),
+                          [](Value operand) { return operand != nullptr; });
+    }
+  }];
+
+  let assemblyFormat = [{
+    $source `,` $dest custom<DynamicIndexList>($dynamic_position, $static_position)
+    attr-dict `:` type($source) `into` type($dest)
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def Vector_ScalableInsertOp :

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

static Value castDataPtr(ConversionPatternRewriter &rewriter, Location loc,
   if (converter.useOpaquePointers())
     return ptr;
 
   unsigned addressSpace = *converter.getMemRefAddressSpace(memRefType);
   auto pType = LLVM::LLVMPointerType::get(vt, addressSpace);
   return rewriter.create<LLVM::BitcastOp>(loc, pType, ptr);
 }
 
+/// Convert `foldResult` into a Value. Integer attribute is converted to
+/// an LLVM constant op.
+static Value getAsLLVMValue(OpBuilder &builder, Location loc,
+                            OpFoldResult foldResult) {
+  if (auto attr = foldResult.dyn_cast<Attribute>()) {
+    auto intAttr = cast<IntegerAttr>(attr);
+    return builder.create<LLVM::ConstantOp>(loc, intAttr).getResult();
+  }
+
+  return foldResult.get<Value>();
+}
+
 namespace {
 
 /// Trivial Vector to LLVM conversions
 using VectorScaleOpConversion =
     OneToOneConvertToLLVMPattern<vector::VectorScaleOp, LLVM::vscale>;
 
 /// Conversion pattern for a vector.bitcast.
 class VectorBitCastOpConversion
public:
   using ConvertOpToLLVMPattern<vector::ExtractOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = extractOp->getLoc();
     auto resultType = extractOp.getResult().getType();
     auto llvmResultType = typeConverter->convertType(resultType);
-    ArrayRef<int64_t> positionArray = extractOp.getPosition();
+    SmallVector<OpFoldResult> positionVec = extractOp.getMixedPosition();
+    ArrayRef<OpFoldResult> position(positionVec);
 
     // Bail if result type cannot be lowered.
     if (!llvmResultType)
       return failure();
 
     // Extract entire vector. Should be handled by folder, but just to be safe.
-    if (positionArray.empty()) {
+    if (position.empty()) {
       rewriter.replaceOp(extractOp, adaptor.getVector());
       return success();
     }
 
     // One-shot extraction of vector from array (only requires extractvalue).
     if (isa<VectorType>(resultType)) {
       Value extracted = rewriter.create<LLVM::ExtractValueOp>(
-          loc, adaptor.getVector(), positionArray);
+          loc, adaptor.getVector(), getAsIntegers(position));
       rewriter.replaceOp(extractOp, extracted);
       return success();
     }
 
     // Potential extraction of 1-D vector from array.
     Value extracted = adaptor.getVector();
-    if (positionArray.size() > 1) {
-      extracted = rewriter.create<LLVM::ExtractValueOp>(
-          loc, extracted, positionArray.drop_back());
+    if (position.size() > 1) {
+      SmallVector<int64_t> nMinusOnePosition =
+          getAsIntegers(position.drop_back());
+      extracted = rewriter.create<LLVM::ExtractValueOp>(loc, extracted,
+                                                        nMinusOnePosition);
     }
 
-    // Remaining extraction of element from 1-D LLVM vector
-    auto i64Type = IntegerType::get(rewriter.getContext(), 64);
-    auto constant =
-        rewriter.create<LLVM::ConstantOp>(loc, i64Type, positionArray.back());
-    extracted =
-        rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant);
-    rewriter.replaceOp(extractOp, extracted);
-
+    Value lastPosition = getAsLLVMValue(rewriter, loc, position.back());
+    // Remaining extraction of element from 1-D LLVM vector.
+    rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>(extractOp, extracted,
+                                                        lastPosition);
     return success();
   }
 };
 
 /// Conversion pattern that turns a vector.fma on a 1-D vector
 /// into an llvm.intr.fmuladd. This is a trivial 1-1 conversion.
 /// This does not match vectors of n >= 2 rank.
 ///
public:
 
   LogicalResult
   matchAndRewrite(vector::InsertOp insertOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = insertOp->getLoc();
     auto sourceType = insertOp.getSourceType();
     auto destVectorType = insertOp.getDestVectorType();
     auto llvmResultType = typeConverter->convertType(destVectorType);
-    ArrayRef<int64_t> positionArray = insertOp.getPosition();
+    SmallVector<OpFoldResult> positionVec = insertOp.getMixedPosition();
+    ArrayRef<OpFoldResult> position(positionVec);
 
     // Bail if result type cannot be lowered.
     if (!llvmResultType)
       return failure();
 
     // Overwrite entire vector with value. Should be handled by folder, but
     // just to be safe.
-    if (positionArray.empty()) {
+    if (position.empty()) {
       rewriter.replaceOp(insertOp, adaptor.getSource());
       return success();
     }
 
     // One-shot insertion of a vector into an array (only requires insertvalue).
     if (isa<VectorType>(sourceType)) {
       Value inserted = rewriter.create<LLVM::InsertValueOp>(
-          loc, adaptor.getDest(), adaptor.getSource(), positionArray);
+          loc, adaptor.getDest(), adaptor.getSource(), getAsIntegers(position));
       rewriter.replaceOp(insertOp, inserted);
       return success();
     }
 
     // Potential extraction of 1-D vector from array.
     Value extracted = adaptor.getDest();
     auto oneDVectorType = destVectorType;
-    if (positionArray.size() > 1) {
+    if (position.size() > 1) {
       oneDVectorType = reducedVectorTypeBack(destVectorType);
       extracted = rewriter.create<LLVM::ExtractValueOp>(
-          loc, extracted, positionArray.drop_back());
+          loc, extracted, getAsIntegers(position.drop_back()));
     }
 
     // Insertion of an element into a 1-D LLVM vector.
-    auto i64Type = IntegerType::get(rewriter.getContext(), 64);
-    auto constant =
-        rewriter.create<LLVM::ConstantOp>(loc, i64Type, positionArray.back());
     Value inserted = rewriter.create<LLVM::InsertElementOp>(
         loc, typeConverter->convertType(oneDVectorType), extracted,
-        adaptor.getSource(), constant);
+        adaptor.getSource(), getAsLLVMValue(rewriter, loc, position.back()));
 
     // Potential insertion of resulting 1-D vector into array.
-    if (positionArray.size() > 1) {
+    if (position.size() > 1) {
       inserted = rewriter.create<LLVM::InsertValueOp>(
-          loc, adaptor.getDest(), inserted, positionArray.drop_back());
+          loc, adaptor.getDest(), inserted,
+          getAsIntegers(position.drop_back()));
     }
 
     rewriter.replaceOp(insertOp, inserted);
     return success();
   }
 };
 
 /// Lower vector.scalable.insert ops to LLVM vector.insert

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

vector::InsertOp getInsertOp(TransferReadOp xferOp) const {
     }
 
     return vector::InsertOp();
   }
 
   /// If the result of the TransferReadOp has exactly one user, which is a
   /// vector::InsertOp, return that operation's indices.
   void getInsertionIndices(TransferReadOp xferOp,
-                           SmallVector<int64_t, 8> &indices) const {
-    if (auto insertOp = getInsertOp(xferOp))
-      indices.assign(insertOp.getPosition().begin(),
-                     insertOp.getPosition().end());
+                           SmallVectorImpl<OpFoldResult> &indices) const {
+    if (auto insertOp = getInsertOp(xferOp)) {
+      auto pos = insertOp.getMixedPosition();
+      indices.append(pos.begin(), pos.end());
+    }
   }
 
   /// Rewrite the op: Unpack one dimension. Can handle masks, out-of-bounds
   /// accesses, and broadcasts and transposes in permutation maps.
   LogicalResult matchAndRewrite(TransferReadOp xferOp,
                                 PatternRewriter &rewriter) const override {
     if (xferOp.getVectorType().getRank() <= options.targetRank)
       return failure();
for (int64_t i = 0; i < dimSize; ++i) {
           rewriter, xferOp, iv, unpackedDim(xferOp), TypeRange(vecType),
           /*inBoundsCase=*/
           [&](OpBuilder &b, Location loc) {
             // Indices for the new transfer op.
             SmallVector<Value, 8> xferIndices;
             getXferIndices(b, xferOp, iv, xferIndices);
 
             // Indices for the new vector.insert op.
-            SmallVector<int64_t, 8> insertionIndices;
+            SmallVector<OpFoldResult, 8> insertionIndices;
             getInsertionIndices(xferOp, insertionIndices);
-            insertionIndices.push_back(i);
+            insertionIndices.push_back(rewriter.getIndexAttr(i));
 
             auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
             auto newXferOp = b.create<vector::TransferReadOp>(
                 loc, newXferVecType, xferOp.getSource(), xferIndices,
                 AffineMapAttr::get(unpackedPermutationMap(b, xferOp)),
                 xferOp.getPadding(), Value(), inBoundsAttr);
             maybeAssignMask(b, xferOp, newXferOp, i);
             return b.create<vector::InsertOp>(loc, newXferOp, vec,
vector::ExtractOp getExtractOp(TransferWriteOp xferOp) const {
     if (auto *op = xferOp.getVector().getDefiningOp())
       return dyn_cast<vector::ExtractOp>(op);
     return vector::ExtractOp();
   }
 
   /// If the input of the given TransferWriteOp is an ExtractOp, return its
   /// indices.
   void getExtractionIndices(TransferWriteOp xferOp,
-                            SmallVector<int64_t, 8> &indices) const {
-    if (auto extractOp = getExtractOp(xferOp))
-      indices.assign(extractOp.getPosition().begin(),
-                     extractOp.getPosition().end());
+                            SmallVectorImpl<OpFoldResult> &indices) const {
+    if (auto extractOp = getExtractOp(xferOp)) {
+      auto pos = extractOp.getMixedPosition();
+      indices.append(pos.begin(), pos.end());
+    }
   }
 
   /// Rewrite the op: Unpack one dimension. Can handle masks, out-of-bounds
   /// accesses, and broadcasts and transposes in permutation maps.
   LogicalResult matchAndRewrite(TransferWriteOp xferOp,
                                 PatternRewriter &rewriter) const override {
     if (xferOp.getVectorType().getRank() <= options.targetRank)
       return failure();
for (int64_t i = 0; i < dimSize; ++i) {
           isTensorOp(xferOp) ? TypeRange(sourceType) : TypeRange(),
           /*inBoundsCase=*/
           [&](OpBuilder &b, Location loc) {
             // Indices for the new transfer op.
             SmallVector<Value, 8> xferIndices;
             getXferIndices(b, xferOp, iv, xferIndices);
 
             // Indices for the new vector.extract op.
-            SmallVector<int64_t, 8> extractionIndices;
+            SmallVector<OpFoldResult, 8> extractionIndices;
             getExtractionIndices(xferOp, extractionIndices);
-            extractionIndices.push_back(i);
+            extractionIndices.push_back(b.getI64IntegerAttr(i));
 
             auto extracted =
                 b.create<vector::ExtractOp>(loc, vec, extractionIndices);
             auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
             auto newXferOp = b.create<vector::TransferWriteOp>(
                 loc, sourceType, extracted, source, xferIndices,
                 AffineMapAttr::get(unpackedPermutationMap(b, xferOp)), Value(),
                 inBoundsAttr);

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

 #include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/Support/FormatVariadic.h"
 #include <cassert>
 #include <cstdint>
 #include <numeric>
 
 using namespace mlir;
 
 /// Gets the first integer value from `attr`, assuming it is an integer array
 /// attribute.
+static uint64_t getFirstIntValue(ValueRange values) {
+  return values[0].getDefiningOp<arith::ConstantIndexOp>().value();
+}
+static uint64_t getFirstIntValue(ArrayRef<Attribute> attr) {
+  return cast<IntegerAttr>(attr[0]).getInt();
+}
 static uint64_t getFirstIntValue(ArrayAttr attr) {
   return (*attr.getAsValueRange<IntegerAttr>().begin()).getZExtValue();
 }
+static uint64_t getFirstIntValue(ArrayRef<OpFoldResult> foldResults) {
+  auto attr = foldResults[0].dyn_cast<Attribute>();
+  if (attr)
+    return getFirstIntValue(attr);
+
+  return getFirstIntValue(ValueRange{foldResults[0].get<Value>()});
+}
 
 /// Returns the number of bits for the given scalar/vector type.
 static int getNumBits(Type type) {
   // TODO: This does not take into account any memory layout or widening
   // constraints. E.g., a vector<3xi57> may report to occupy 3x57=171 bit, even
   // though in practice it will likely be stored as in a 4xi64 vector register.
   if (auto vectorType = dyn_cast<VectorType>(type))
     return vectorType.getNumElements() * vectorType.getElementTypeBitWidth();

 
 struct VectorExtractOpConvert final
     : public OpConversionPattern<vector::ExtractOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    // Only support extracting a scalar value now.
-    VectorType resultVectorType = dyn_cast<VectorType>(extractOp.getType());
-    if (resultVectorType && resultVectorType.getNumElements() > 1)
+    if (extractOp.hasDynamicPosition())
       return failure();
 
     Type dstType = getTypeConverter()->convertType(extractOp.getType());
     if (!dstType)
       return failure();
 
     if (isa<spirv::ScalarType>(adaptor.getVector().getType())) {
       rewriter.replaceOp(extractOp, adaptor.getVector());
       return success();
     }
 
-    int32_t id = extractOp.getPosition()[0];
+    int32_t id = getFirstIntValue(extractOp.getMixedPosition());
     rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
         extractOp, adaptor.getVector(), id);
     return success();
   }
 };
 
 struct VectorExtractStridedSliceOpConvert final
     : public OpConversionPattern<vector::ExtractStridedSliceOp> {
matchAndRewrite(vector::InsertOp insertOp, OpAdaptor adaptor,
 
     // Special case for inserting scalar values into size-1 vectors.
     if (insertOp.getSourceType().isIntOrFloat() &&
         insertOp.getDestVectorType().getNumElements() == 1) {
       rewriter.replaceOp(insertOp, adaptor.getSource());
       return success();
     }
 
-    int32_t id = insertOp.getPosition()[0];
+    int32_t id = getFirstIntValue(insertOp.getMixedPosition());
     rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
         insertOp, adaptor.getSource(), adaptor.getDest(), id);
     return success();
   }
 };
 
 struct VectorExtractElementOpConvert final
     : public OpConversionPattern<vector::ExtractElementOp> {

mlir/lib/Dialect/Arith/Transforms/IntNarrowing.cpp

struct ExtensionOverExtract final : NarrowingPattern<vector::ExtractOp> {
   LogicalResult matchAndRewrite(vector::ExtractOp op,
                                 PatternRewriter &rewriter) const override {
     FailureOr<ExtensionOp> ext =
         ExtensionOp::from(op.getVector().getDefiningOp());
     if (failed(ext))
       return failure();
 
     Value newExtract = rewriter.create<vector::ExtractOp>(
-        op.getLoc(), ext->getIn(), op.getPosition());
+        op.getLoc(), ext->getIn(), op.getMixedPosition());
     ext->recreateAndReplace(rewriter, op, newExtract);
     return success();
   }
 };
 
 struct ExtensionOverExtractElement final
     : NarrowingPattern<vector::ExtractElementOp> {
   using NarrowingPattern::NarrowingPattern;

 struct ExtensionOverInsert final
     : ExtensionOverInsertionPattern<vector::InsertOp> {
   using ExtensionOverInsertionPattern::ExtensionOverInsertionPattern;
 
   vector::InsertOp createInsertionOp(PatternRewriter &rewriter,
                                      vector::InsertOp origInsert,
                                      Value narrowValue,
                                      Value narrowDest) const override {
-    return rewriter.create<vector::InsertOp>(
-        origInsert.getLoc(), narrowValue, narrowDest, origInsert.getPosition());
+    return rewriter.create<vector::InsertOp>(origInsert.getLoc(), narrowValue,
+                                             narrowDest,
+                                             origInsert.getMixedPosition());
   }
 };
 
 struct ExtensionOverInsertElement final
     : ExtensionOverInsertionPattern<vector::InsertElementOp> {
   using ExtensionOverInsertionPattern::ExtensionOverInsertionPattern;
 
   vector::InsertElementOp createInsertionOp(PatternRewriter &rewriter,

mlir/lib/Dialect/NVGPU/Transforms/CreateAsyncGroups.cpp

if (auto maskOp =
     return TransferMask{maskOp, {}};
 
   // Case 2: Mask is the result of a vector.extract(vector.create_mask).
   if (auto extractOp =
           transferRead.getMask().getDefiningOp<vector::ExtractOp>())
     if (auto maskOp =
             extractOp.getVector().getDefiningOp<vector::CreateMaskOp>())
       return TransferMask{maskOp,
-                          SmallVector<int64_t>(extractOp.getPosition())};
+                          SmallVector<int64_t>(extractOp.getStaticPosition())};
 
   // All other cases: not supported.
   return failure();
 }
 
 /// Build an SSA value that represents the number of read elements.
 static Value buildNumReadElements(OpBuilder &b, Location loc,
                                   Operation *readOp) {

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

for (auto [posInDim, dimSize, offsetInDim] :
 
     // Carry the overflow to the next loop iteration.
     posInDim = offsetInDim;
   }
 
   return failure();
 }
 
+/// Returns the integer numbers in `values`. `values` are expected to be
+/// constant operations.
+SmallVector<int64_t> vector::getAsIntegers(ArrayRef<Value> values) {
+  SmallVector<int64_t> ints;
+  llvm::transform(values, std::back_inserter(ints), [](Value value) {
+    auto constOp = value.getDefiningOp<arith::ConstantIndexOp>();
+    assert(constOp && "Unexpected non-constant index");
+    return constOp.value();
+  });
+  return ints;
+}
+
+/// Returns the integer numbers in `foldResults`. `foldResults` are expected to
+/// be constant operations.
+SmallVector<int64_t> vector::getAsIntegers(ArrayRef<OpFoldResult> foldResults) {
+  SmallVector<int64_t> ints;
+  llvm::transform(
+      foldResults, std::back_inserter(ints), [](OpFoldResult foldResult) {
+        assert(foldResult.is<Attribute>() && "Unexpected non-constant index");
+        return cast<IntegerAttr>(foldResult.get<Attribute>()).getInt();
+      });
+  return ints;
+}
+
+/// Convert `foldResults` into Values. Integer attributes are converted to
+/// constant op.
+SmallVector<Value> vector::getAsValues(OpBuilder &builder, Location loc,
+                                       ArrayRef<OpFoldResult> foldResults) {
+  SmallVector<Value> values;
+  llvm::transform(foldResults, std::back_inserter(values),
+                  [&](OpFoldResult foldResult) {
+                    if (auto attr = foldResult.dyn_cast<Attribute>())
+                      return builder
+                          .create<arith::ConstantIndexOp>(
+                              loc, cast<IntegerAttr>(attr).getInt())
+                          .getResult();
+
+                    return foldResult.get<Value>();
+                  });
+  return values;
+}
+
 //===----------------------------------------------------------------------===//
 // CombiningKindAttr
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace vector {
 namespace detail {
 struct BitmaskEnumStorage : public AttributeStorage {
if (auto dstVecType = dyn_cast<VectorType>(reductionOp.getDestType())) {
                             dstVecType.getScalableDims());
         mask = rewriter.create<vector::ShapeCastOp>(loc, newMaskType, mask);
       }
       cast = rewriter.create<vector::ShapeCastOp>(
           loc, reductionOp.getDestType(), reductionOp.getSource());
     } else {
       // This means we are reducing all the dimensions, and all reduction
       // dimensions are of size 1. So a simple extraction would do.
-      SmallVector<int64_t> zeroAttr(shape.size(), 0);
+      SmallVector<int64_t> zeroIdx(shape.size(), 0);
       if (mask)
-        mask = rewriter.create<vector::ExtractOp>(loc, rewriter.getI1Type(),
-                                                  mask, zeroAttr);
-      cast = rewriter.create<vector::ExtractOp>(
-          loc, reductionOp.getDestType(), reductionOp.getSource(), zeroAttr);
+        mask = rewriter.create<vector::ExtractOp>(loc, mask, zeroIdx);
+      cast = rewriter.create<vector::ExtractOp>(loc, reductionOp.getSource(),
+                                                zeroIdx);
     }
 
     Value result = vector::makeArithReduction(
         rewriter, loc, reductionOp.getKind(), acc, cast, mask);
     rewriter.replaceOp(rootOp, result);
     return success();
   }
 };
LogicalResult matchAndRewrite(ReductionOp reductionOp,
 
     Location loc = reductionOp.getLoc();
     Value result;
     if (vectorType.getRank() == 0) {
       if (mask)
         mask = rewriter.create<ExtractElementOp>(loc, mask);
       result = rewriter.create<ExtractElementOp>(loc, reductionOp.getVector());
     } else {
-      if (mask) {
-        mask = rewriter.create<ExtractOp>(loc, rewriter.getI1Type(), mask, 0);
-      }
-      result = rewriter.create<ExtractOp>(loc, reductionOp.getType(),
-                                          reductionOp.getVector(), 0);
+      if (mask)
+        mask = rewriter.create<ExtractOp>(loc, mask, 0);
+      result = rewriter.create<ExtractOp>(loc, reductionOp.getVector(), 0);
     }
 
     if (Value acc = reductionOp.getAcc())
       result = vector::makeArithReduction(rewriter, loc, reductionOp.getKind(),
                                           result, acc, mask);
 
     rewriter.replaceOp(rootOp, result);
     return success();
OpFoldResult vector::ExtractElementOp::fold(FoldAdaptor adaptor) {
 
   return srcElements[posIdx];
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractOp
 //===----------------------------------------------------------------------===//
 
-// Convenience builder which assumes the values are constant indices.
 void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
-                              Value source, ValueRange position) {
-  SmallVector<int64_t> positionConstants = llvm::to_vector(llvm::map_range(
-      position, [](Value pos) { return getConstantIntValue(pos).value(); }));
-  build(builder, result, source, positionConstants);
+                              Value source, int64_t position) {
+  build(builder, result, source, ArrayRef<int64_t>{position});
+}
+
+void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
+                              Value source, OpFoldResult position) {
+  build(builder, result, source, ArrayRef<OpFoldResult>{position});
+}
+
+void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
+                              Value source, ArrayRef<int64_t> position) {
+  build(builder, result, source, /*dynamic_position=*/ArrayRef<Value>(),
+        builder.getDenseI64ArrayAttr(position));
+}
+
+void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
+                              Value source, ArrayRef<OpFoldResult> position) {
+  SmallVector<int64_t> staticPos;
+  SmallVector<Value> dynamicPos;
+  dispatchIndexOpFoldResults(position, dynamicPos, staticPos);
+  build(builder, result, source, dynamicPos,
+        builder.getDenseI64ArrayAttr(staticPos));
 }
 
 LogicalResult
 ExtractOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
                             ExtractOp::Adaptor adaptor,
                             SmallVectorImpl<Type> &inferredReturnTypes) {
   auto vectorType = llvm::cast<VectorType>(adaptor.getVector().getType());
-  if (static_cast<int64_t>(adaptor.getPosition().size()) ==
+  if (static_cast<int64_t>(adaptor.getStaticPosition().size()) ==
       vectorType.getRank()) {
     inferredReturnTypes.push_back(vectorType.getElementType());
   } else {
-    auto n =
-        std::min<size_t>(adaptor.getPosition().size(), vectorType.getRank());
+    auto n = std::min<size_t>(adaptor.getStaticPosition().size(),
+                              vectorType.getRank());
     inferredReturnTypes.push_back(VectorType::get(
         vectorType.getShape().drop_front(n), vectorType.getElementType(),
         vectorType.getScalableDims().drop_front(n)));
   }
   return success();
 }
 
 bool ExtractOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   // Allow extracting 1-element vectors instead of scalars.
   auto isCompatible = [](TypeRange l, TypeRange r) {
     auto vectorType = llvm::dyn_cast<VectorType>(l.front());
     return vectorType && vectorType.getShape().equals({1}) &&
            vectorType.getElementType() == r.front();
   };
   if (l.size() == 1 && r.size() == 1 &&
       (isCompatible(l, r) || isCompatible(r, l)))
     return true;
   return l == r;
 }
 
 LogicalResult vector::ExtractOp::verify() {
-  ArrayRef<int64_t> position = getPosition();
+  auto position = getMixedPosition();
   if (position.size() > static_cast<unsigned>(getSourceVectorType().getRank()))
     return emitOpError(
         "expected position attribute of rank no greater than vector rank");
-  for (const auto &en : llvm::enumerate(position)) {
-    if (en.value() < 0 ||
-        en.value() >= getSourceVectorType().getDimSize(en.index()))
+  for (auto [idx, pos] : llvm::enumerate(position)) {
+    if (pos.is<Attribute>()) {
+      int64_t constIdx = cast<IntegerAttr>(pos.get<Attribute>()).getInt();
+      if (constIdx < 0 || constIdx >= getSourceVectorType().getDimSize(idx)) {
         return emitOpError("expected position attribute #")
-             << (en.index() + 1)
-             << " to be a non-negative integer smaller than the corresponding "
+               << (idx + 1)
+               << " to be a non-negative integer smaller than the "
+                  "corresponding "
                   "vector dimension";
       }
+    }
+  }
   return success();
 }
 
 template <typename IntType>
 static SmallVector<IntType> extractVector(ArrayAttr arrayAttr) {
   return llvm::to_vector<4>(llvm::map_range(
       arrayAttr.getAsRange<IntegerAttr>(),
       [](IntegerAttr attr) { return static_cast<IntType>(attr.getInt()); }));
 }
 
 /// Fold the result of chains of ExtractOp in place by simply concatenating the
 /// positions.
 static LogicalResult foldExtractOpFromExtractChain(ExtractOp extractOp) {
   if (!extractOp.getVector().getDefiningOp<ExtractOp>())
     return failure();
 
-  SmallVector<int64_t, 4> globalPosition;
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return failure();
+
+  SmallVector<int64_t> globalPosition;
   ExtractOp currentOp = extractOp;
-  ArrayRef<int64_t> extrPos = currentOp.getPosition();
+  ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
   globalPosition.append(extrPos.rbegin(), extrPos.rend());
   while (ExtractOp nextOp = currentOp.getVector().getDefiningOp<ExtractOp>()) {
     currentOp = nextOp;
-    ArrayRef<int64_t> extrPos = currentOp.getPosition();
+    ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
     globalPosition.append(extrPos.rbegin(), extrPos.rend());
   }
-  extractOp.setOperand(currentOp.getVector());
+  extractOp.setOperand(0, currentOp.getVector());
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(extractOp.getContext());
   std::reverse(globalPosition.begin(), globalPosition.end());
-  extractOp.setPosition(globalPosition);
+  extractOp.setStaticPosition(globalPosition);
   return success();
 }
 
 namespace {
 /// Fold an ExtractOp that is fed by a chain of InsertOps and TransposeOps.
 /// Walk back a chain of InsertOp/TransposeOp until we hit a match.
 /// Compose TransposeOp permutations as we walk back.
 /// This helper class keeps an updated extraction position `extractPosition`
private:
   SmallVector<int64_t> sentinels;
   SmallVector<int64_t> extractPosition;
 };
 } // namespace
 
 ExtractFromInsertTransposeChainState::ExtractFromInsertTransposeChainState(
     ExtractOp e)
     : extractOp(e), vectorRank(extractOp.getSourceVectorType().getRank()),
-      extractedRank(extractOp.getPosition().size()) {
+      extractedRank(extractOp.getNumIndices()) {
   assert(vectorRank >= extractedRank && "extracted pos overflow");
   sentinels.reserve(vectorRank - extractedRank);
   for (int64_t i = 0, e = vectorRank - extractedRank; i < e; ++i)
     sentinels.push_back(-(i + 1));
-  extractPosition.assign(extractOp.getPosition().begin(),
-                         extractOp.getPosition().end());
+  extractPosition.assign(extractOp.getStaticPosition().begin(),
+                         extractOp.getStaticPosition().end());
   llvm::append_range(extractPosition, sentinels);
 }
 
 // Case 1. If we hit a transpose, just compose the map and iterate.
 // Invariant: insert + transpose do not change rank, we can always compose.
 LogicalResult ExtractFromInsertTransposeChainState::handleTransposeOp() {
   if (!nextTransposeOp)
     return failure();
   auto permutation = extractVector<unsigned>(nextTransposeOp.getTransp());
   AffineMap m = inversePermutation(
       AffineMap::getPermutationMap(permutation, extractOp.getContext()));
   extractPosition = applyPermutationMap(m, ArrayRef(extractPosition));
   return success();
 }
 
 // Case 2: the insert position matches extractPosition exactly, early return.
 LogicalResult
 ExtractFromInsertTransposeChainState::handleInsertOpWithMatchingPos(
     Value &res) {
-  ArrayRef<int64_t> insertedPos = nextInsertOp.getPosition();
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (nextInsertOp.hasDynamicPosition())
+    return failure();
+
+  ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
   if (insertedPos != llvm::ArrayRef(extractPosition).take_front(extractedRank))
     return failure();
   // Case 2.a. early-exit fold.
   res = nextInsertOp.getSource();
   // Case 2.b. if internal transposition is present, canFold will be false.
   return success(canFold());
 }
 
 /// Case 3: if inserted position is a prefix of extractPosition,
 /// extract a portion of the source of the insertion.
 /// This method updates the internal state.
 LogicalResult
 ExtractFromInsertTransposeChainState::handleInsertOpWithPrefixPos(Value &res) {
-  ArrayRef<int64_t> insertedPos = nextInsertOp.getPosition();
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (nextInsertOp.hasDynamicPosition())
+    return failure();
+
+  ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
   if (!isContainedWithin(insertedPos, extractPosition))
     return failure();
   // Set leading dims to zero.
   std::fill_n(extractPosition.begin(), insertedPos.size(), 0);
   // Drop extra leading dims.
   extractPosition.erase(extractPosition.begin(),
                         extractPosition.begin() + insertedPos.size());
   extractedRank = extractPosition.size() - sentinels.size();
   // Case 3.a. early-exit fold (break and delegate to post-while path).
   res = nextInsertOp.getSource();
   // Case 3.b. if internal transposition is present, canFold will be false.
   return success();
 }
 
 /// Try to fold in place to extract(source, extractPosition) and return the
 /// folded result. Return null if folding is not possible (e.g. due to an
 /// internal tranposition in the result).
 Value ExtractFromInsertTransposeChainState::tryToFoldExtractOpInPlace(
     Value source) {
   // If we can't fold (either internal transposition, or nothing to fold), bail.
   bool nothingToFold = (source == extractOp.getVector());
   if (nothingToFold || !canFold())
     return Value();
+
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
   // Otherwise, fold by updating the op inplace and return its result.
   OpBuilder b(extractOp.getContext());
-  extractOp.setPosition(ArrayRef(extractPosition).take_front(extractedRank));
+  extractOp.setStaticPosition(
+      ArrayRef(extractPosition).take_front(extractedRank));
   extractOp.getVectorMutable().assign(source);
   return extractOp.getResult();
 }
 
 /// Iterate over producing insert and transpose ops until we find a fold.
 Value ExtractFromInsertTransposeChainState::fold() {
   Value valueToExtractFrom = extractOp.getVector();
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
   updateStateForNextIteration(valueToExtractFrom);
   while (nextInsertOp || nextTransposeOp) {
     // Case 1. If we hit a transpose, just compose the map and iterate.
     // Invariant: insert + transpose do not change rank, we can always compose.
     if (succeeded(handleTransposeOp())) {
       valueToExtractFrom = nextTransposeOp.getVector();
       updateStateForNextIteration(valueToExtractFrom);
       continue;
     }
 
     Value result;
     // Case 2: the position match exactly.
     if (succeeded(handleInsertOpWithMatchingPos(result)))
       return result;
 
     // Case 3: if the inserted position is a prefix of extractPosition, we can
     // just extract a portion of the source of the insert.
     if (succeeded(handleInsertOpWithPrefixPos(result)))
       return tryToFoldExtractOpInPlace(result);
 
     // Case 4: extractPositionRef intersects insertedPosRef on non-sentinel
     // values. This is a more difficult case and we bail.
-    ArrayRef<int64_t> insertedPos = nextInsertOp.getPosition();
+    ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
     if (isContainedWithin(extractPosition, insertedPos) ||
         intersectsWhereNonNegative(extractPosition, insertedPos))
       return Value();
 
     // Case 5: No intersection, we forward the extract to insertOp.dest().
     valueToExtractFrom = nextInsertOp.getDest();
     updateStateForNextIteration(valueToExtractFrom);
   }
static bool hasZeroDimVectors(Operation *op) {
   };
 
   return llvm::any_of(op->getOperandTypes(), hasZeroDimVectorType) ||
          llvm::any_of(op->getResultTypes(), hasZeroDimVectorType);
 }
 
 /// Fold extractOp with scalar result coming from BroadcastOp or SplatOp.
 static Value foldExtractFromBroadcast(ExtractOp extractOp) {
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
   Operation *defOp = extractOp.getVector().getDefiningOp();
   if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
     return Value();
 
   // 0-D vectors not supported.
   assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
   if (hasZeroDimVectors(defOp))
     return Value();
static Value foldExtractFromBroadcast(ExtractOp extractOp) {
   auto broadcastOp = cast<vector::BroadcastOp>(defOp);
   int64_t broadcastDstRank = broadcastOp.getResultVectorType().getRank();
 
   // Detect all the positions that come from "dim-1" broadcasting.
   // These dimensions correspond to "dim-1" broadcasted dims; set the mathching
   // extract position to `0` when extracting from the source operand.
   llvm::SetVector<int64_t> broadcastedUnitDims =
       broadcastOp.computeBroadcastedUnitDims();
-  SmallVector<int64_t> extractPos(extractOp.getPosition());
+  SmallVector<int64_t> extractPos(extractOp.getStaticPosition());
   int64_t broadcastRankDiff = broadcastDstRank - broadcastSrcRank;
   for (int64_t i = broadcastRankDiff, e = extractPos.size(); i < e; ++i)
     if (broadcastedUnitDims.contains(i))
       extractPos[i] = 0;
   // `rankDiff` leading dimensions correspond to new broadcasted dims, drop the
   // matching extract position when extracting from the source operand.
   int64_t rankDiff = broadcastSrcRank - extractResultRank;
   extractPos.erase(extractPos.begin(),
                    std::next(extractPos.begin(), extractPos.size() - rankDiff));
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(extractOp.getContext());
-  extractOp.setOperand(source);
-  extractOp.setPosition(extractPos);
+  extractOp.setOperand(0, source);
+  extractOp.setStaticPosition(extractPos);
   return extractOp.getResult();
 }
 
 // Fold extractOp with source coming from ShapeCast op.
 static Value foldExtractFromShapeCast(ExtractOp extractOp) {
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
   auto shapeCastOp = extractOp.getVector().getDefiningOp<vector::ShapeCastOp>();
   if (!shapeCastOp)
     return Value();
 
   // 0-D vectors not supported.
   assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
   if (hasZeroDimVectors(shapeCastOp))
     return Value();
for (int64_t i = 0; i < destinationRank; i++) {
       // TODO: This case could be support in a canonicalization pattern.
       if (getDimReverse(shapeCastOp.getSourceVectorType(), i) !=
           getDimReverse(destinationType, i))
         return Value();
     }
   }
   // Extract the strides associated with the extract op vector source. Then use
   // this to calculate a linearized position for the extract.
-  SmallVector<int64_t> extractedPos(extractOp.getPosition());
+  SmallVector<int64_t> extractedPos(extractOp.getStaticPosition());
   std::reverse(extractedPos.begin(), extractedPos.end());
   SmallVector<int64_t, 4> strides;
   int64_t stride = 1;
   for (int64_t i = 0, e = extractedPos.size(); i < e; i++) {
     strides.push_back(stride);
     stride *=
         getDimReverse(extractOp.getSourceVectorType(), i + destinationRank);
   }
for (int64_t i = 0; i < numDimension; i++) {
     newStrides.push_back(stride);
     stride *=
         getDimReverse(shapeCastOp.getSourceVectorType(), i + destinationRank);
   }
   std::reverse(newStrides.begin(), newStrides.end());
   SmallVector<int64_t, 4> newPosition = delinearize(position, newStrides);
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(extractOp.getContext());
-  extractOp.setPosition(newPosition);
-  extractOp.setOperand(shapeCastOp.getSource());
+  extractOp.setStaticPosition(newPosition);
+  extractOp.setOperand(0, shapeCastOp.getSource());
   return extractOp.getResult();
 }
 
 /// Fold an ExtractOp from ExtractStridedSliceOp.
 static Value foldExtractFromExtractStrided(ExtractOp extractOp) {
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
   auto extractStridedSliceOp =
       extractOp.getVector().getDefiningOp<vector::ExtractStridedSliceOp>();
   if (!extractStridedSliceOp)
     return Value();
 
   // 0-D vectors not supported.
   assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
   if (hasZeroDimVectors(extractStridedSliceOp))
static Value foldExtractFromExtractStrided(ExtractOp extractOp) {
   unsigned destinationRank = 0;
   if (auto vecType = llvm::dyn_cast<VectorType>(extractOp.getType()))
     destinationRank = vecType.getRank();
   // The dimensions of the result need to be untouched by the
   // extractStridedSlice op.
   if (destinationRank > extractStridedSliceOp.getSourceVectorType().getRank() -
                             sliceOffsets.size())
     return Value();
-  SmallVector<int64_t> extractedPos(extractOp.getPosition());
+
+  SmallVector<int64_t> extractedPos(extractOp.getStaticPosition());
   assert(extractedPos.size() >= sliceOffsets.size());
   for (size_t i = 0, e = sliceOffsets.size(); i < e; i++)
     extractedPos[i] = extractedPos[i] + sliceOffsets[i];
   extractOp.getVectorMutable().assign(extractStridedSliceOp.getVector());
+
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(extractOp.getContext());
-  extractOp.setPosition(extractedPos);
+  extractOp.setStaticPosition(extractedPos);
   return extractOp.getResult();
 }
 
 /// Fold extract_op fed from a chain of insertStridedSlice ops.
 static Value foldExtractStridedOpFromInsertChain(ExtractOp extractOp) {
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
   int64_t destinationRank =
       llvm::isa<VectorType>(extractOp.getType())
           ? llvm::cast<VectorType>(extractOp.getType()).getRank()
           : 0;
   auto insertOp = extractOp.getVector().getDefiningOp<InsertStridedSliceOp>();
   if (!insertOp)
     return Value();
 
   // 0-D vectors not supported.
   assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
   if (hasZeroDimVectors(insertOp))
     return Value();
 
   while (insertOp) {
     int64_t insertRankDiff = insertOp.getDestVectorType().getRank() -
                              insertOp.getSourceVectorType().getRank();
     if (destinationRank > insertOp.getSourceVectorType().getRank())
       return Value();
     auto insertOffsets = extractVector<int64_t>(insertOp.getOffsets());
-    ArrayRef<int64_t> extractOffsets = extractOp.getPosition();
+    ArrayRef<int64_t> extractOffsets = extractOp.getStaticPosition();
 
     if (llvm::any_of(insertOp.getStrides(), [](Attribute attr) {
           return llvm::cast<IntegerAttr>(attr).getInt() != 1;
         }))
       return Value();
     bool disjoint = false;
     SmallVector<int64_t, 4> offsetDiffs;
     for (unsigned dim = 0, e = extractOffsets.size(); dim < e; ++dim) {
if (!disjoint) {
         if (insertOp.getSourceVectorType().getDimSize(i + srcRankDiff) !=
             insertOp.getDestVectorType().getDimSize(i + srcRankDiff +
                                                     insertRankDiff))
           return Value();
       }
       extractOp.getVectorMutable().assign(insertOp.getSource());
       // OpBuilder is only used as a helper to build an I64ArrayAttr.
       OpBuilder b(extractOp.getContext());
-      extractOp.setPosition(offsetDiffs);
+      extractOp.setStaticPosition(offsetDiffs);
       return extractOp.getResult();
     }
     // If the chunk extracted is disjoint from the chunk inserted, keep
     // looking in the insert chain.
     insertOp = insertOp.getDest().getDefiningOp<InsertStridedSliceOp>();
   }
   return Value();
 }
 
 OpFoldResult ExtractOp::fold(FoldAdaptor) {
-  if (getPosition().empty())
+  if (getNumIndices() == 0)
     return getVector();
   if (succeeded(foldExtractOpFromExtractChain(*this)))
     return getResult();
   if (auto res = ExtractFromInsertTransposeChainState(*this).fold())
     return res;
   if (auto res = foldExtractFromBroadcast(*this))
     return res;
   if (auto res = foldExtractFromShapeCast(*this))

 // Pattern to rewrite a ExtractOp(non-splat ConstantOp)[...] -> ConstantOp.
 class ExtractOpNonSplatConstantFolder final
     : public OpRewritePattern<ExtractOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractOp extractOp,
                                 PatternRewriter &rewriter) const override {
+    // TODO: Canonicalization for dynamic position not implemented yet.
+    if (extractOp.hasDynamicPosition())
+      return failure();
+
     // Return if 'ExtractOp' operand is not defined by a compatible vector
     // ConstantOp.
     Value sourceVector = extractOp.getVector();
     Attribute vectorCst;
     if (!matchPattern(sourceVector, m_Constant(&vectorCst)))
       return failure();
 
     auto vecTy = llvm::cast<VectorType>(sourceVector.getType());
     if (vecTy.isScalable())
       return failure();
 
     // The splat case is handled by `ExtractOpSplatConstantFolder`.
     auto dense = llvm::dyn_cast<DenseElementsAttr>(vectorCst);
     if (!dense || dense.isSplat())
       return failure();
 
     // Calculate the linearized position of the continuous chunk of elements to
     // extract.
     llvm::SmallVector<int64_t> completePositions(vecTy.getRank(), 0);
-    copy(extractOp.getPosition(), completePositions.begin());
+    copy(extractOp.getStaticPosition(), completePositions.begin());
     int64_t elemBeginPosition =
         linearize(completePositions, computeStrides(vecTy.getShape()));
     auto denseValuesBegin = dense.value_begin<TypedAttr>() + elemBeginPosition;
 
     TypedAttr newAttr;
     if (auto resVecTy = llvm::dyn_cast<VectorType>(extractOp.getType())) {
       SmallVector<Attribute> elementValues(
           denseValuesBegin, denseValuesBegin + resVecTy.getNumElements());
OpFoldResult vector::InsertElementOp::fold(FoldAdaptor adaptor) {
 
   return DenseElementsAttr::get(getDestVectorType(), results);
 }
 
 //===----------------------------------------------------------------------===//
 // InsertOp
 //===----------------------------------------------------------------------===//
 
-// Convenience builder which assumes the values are constant indices.
-void InsertOp::build(OpBuilder &builder, OperationState &result, Value source,
-                     Value dest, ValueRange position) {
-  SmallVector<int64_t, 4> positionConstants =
-      llvm::to_vector<4>(llvm::map_range(position, [](Value pos) {
-        return getConstantIntValue(pos).value();
-      }));
-  build(builder, result, source, dest, positionConstants);
+void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
+                             Value source, Value dest, int64_t position) {
+  build(builder, result, source, dest, ArrayRef<int64_t>{position});
+}
+
+void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
+                             Value source, Value dest, OpFoldResult position) {
+  build(builder, result, source, dest, ArrayRef<OpFoldResult>{position});
+}
+
+void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
+                             Value source, Value dest,
+                             ArrayRef<int64_t> position) {
+  SmallVector<OpFoldResult> posVals;
+  posVals.reserve(position.size());
+  llvm::transform(position, std::back_inserter(posVals),
+                  [&](int64_t pos) { return builder.getI64IntegerAttr(pos); });
+  build(builder, result, source, dest, posVals);
+}
+
+void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
+                             Value source, Value dest,
+                             ArrayRef<OpFoldResult> position) {
+  SmallVector<int64_t> staticPos;
+  SmallVector<Value> dynamicPos;
+  dispatchIndexOpFoldResults(position, dynamicPos, staticPos);
+  build(builder, result, source, dest, dynamicPos,
+        builder.getDenseI64ArrayAttr(staticPos));
 }
 
 LogicalResult InsertOp::verify() {
-  ArrayRef<int64_t> position = getPosition();
+  SmallVector<OpFoldResult> position = getMixedPosition();
   auto destVectorType = getDestVectorType();
   if (position.size() > static_cast<unsigned>(destVectorType.getRank()))
     return emitOpError(
         "expected position attribute of rank no greater than dest vector rank");
   auto srcVectorType = llvm::dyn_cast<VectorType>(getSourceType());
   if (srcVectorType &&
       (static_cast<unsigned>(srcVectorType.getRank()) + position.size() !=
        static_cast<unsigned>(destVectorType.getRank())))
     return emitOpError("expected position attribute rank + source rank to "
                        "match dest vector rank");
   if (!srcVectorType &&
       (position.size() != static_cast<unsigned>(destVectorType.getRank())))
     return emitOpError(
         "expected position attribute rank to match the dest vector rank");
-  for (const auto &en : llvm::enumerate(position)) {
-    int64_t attr = en.value();
-    if (attr < 0 || attr >= destVectorType.getDimSize(en.index()))
+  for (auto [idx, pos] : llvm::enumerate(position)) {
+    if (auto attr = pos.dyn_cast<Attribute>()) {
+      int64_t constIdx = cast<IntegerAttr>(attr).getInt();
+      if (constIdx < 0 || constIdx >= destVectorType.getDimSize(idx)) {
         return emitOpError("expected position attribute #")
-             << (en.index() + 1)
-             << " to be a non-negative integer smaller than the corresponding "
+               << (idx + 1)
+               << " to be a non-negative integer smaller than the "
+                  "corresponding "
                   "dest vector dimension";
       }
+    }
+  }
   return success();
 }
 
 namespace {
 
 // If insertOp is only inserting unit dimensions it can be transformed to a
 // broadcast.
 class InsertToBroadcast final : public OpRewritePattern<InsertOp> {
public:
   using OpRewritePattern::OpRewritePattern;
 
   // Do not create constants with more than `vectorSizeFoldThreashold` elements,
   // unless the source vector constant has a single use.
   static constexpr int64_t vectorSizeFoldThreshold = 256;
 
   LogicalResult matchAndRewrite(InsertOp op,
                                 PatternRewriter &rewriter) const override {
+    // TODO: Canonicalization for dynamic position not implemented yet.
+    if (op.hasDynamicPosition())
+      return failure();
+
     // Return if 'InsertOp' operand is not defined by a compatible vector
     // ConstantOp.
     TypedValue<VectorType> destVector = op.getDest();
     Attribute vectorDestCst;
     if (!matchPattern(destVector, m_Constant(&vectorDestCst)))
       return failure();
 
     VectorType destTy = destVector.getType();
LogicalResult matchAndRewrite(InsertOp op,
     Value sourceValue = op.getSource();
     Attribute sourceCst;
     if (!matchPattern(sourceValue, m_Constant(&sourceCst)))
       return failure();
 
     // Calculate the linearized position of the continuous chunk of elements to
     // insert.
     llvm::SmallVector<int64_t> completePositions(destTy.getRank(), 0);
-    copy(op.getPosition(), completePositions.begin());
+    copy(op.getStaticPosition(), completePositions.begin());
     int64_t insertBeginPosition =
         linearize(completePositions, computeStrides(destTy.getShape()));
 
     SmallVector<Attribute> insertedValues;
     if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(sourceCst))
       llvm::append_range(insertedValues, denseSource.getValues<Attribute>());
     else
       insertedValues.push_back(sourceCst);
void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<InsertToBroadcast, BroadcastFolder, InsertSplatToSplat,
               InsertOpConstantFolder>(context);
 }
 
 // Eliminates insert operations that produce values identical to their source
 // value. This happens when the source and destination vectors have identical
 // sizes.
 OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) {
-  if (getPosition().empty())
+  if (getNumIndices() == 0)
     return getSource();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // InsertStridedSliceOp
 //===----------------------------------------------------------------------===//
 

mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp

 
 // Helper method to possibly drop a dimension in a load.
 // TODO
 static Value reshapeLoad(Location loc, Value val, VectorType type,
                          int64_t index, int64_t pos,
                          PatternRewriter &rewriter) {
   if (index == -1)
     return val;
-  Type lowType = type.getRank() > 1 ? VectorType::Builder(type).dropDim(0)
-                                    : type.getElementType();
+
   // At extraction dimension?
   if (index == 0)
-    return rewriter.create<vector::ExtractOp>(loc, lowType, val, pos);
+    return rewriter.create<vector::ExtractOp>(loc, val, pos);
+
   // Unroll leading dimensions.
+  Type lowType = VectorType::Builder(type).dropDim(0);
   VectorType vType = cast<VectorType>(lowType);
   VectorType resType = VectorType::Builder(type).dropDim(index);
   Value result = rewriter.create<arith::ConstantOp>(
       loc, resType, rewriter.getZeroAttr(resType));
   for (int64_t d = 0, e = resType.getDimSize(0); d < e; d++) {
-    Value ext = rewriter.create<vector::ExtractOp>(loc, vType, val, d);
+    Value ext = rewriter.create<vector::ExtractOp>(loc, val, d);
     Value load = reshapeLoad(loc, ext, vType, index - 1, pos, rewriter);
-    result = rewriter.create<vector::InsertOp>(loc, resType, load, result, d);
+    result = rewriter.create<vector::InsertOp>(loc, load, result, d);
   }
   return result;
 }
 
 // Helper method to possibly drop a dimension in a store.
 // TODO
 static Value reshapeStore(Location loc, Value val, Value result,
                           VectorType type, int64_t index, int64_t pos,
                           PatternRewriter &rewriter) {
   // Unmodified?
   if (index == -1)
     return val;
   // At insertion dimension?
   if (index == 0)
-    return rewriter.create<vector::InsertOp>(loc, type, val, result, pos);
+    return rewriter.create<vector::InsertOp>(loc, val, result, pos);
+
   // Unroll leading dimensions.
-  VectorType lowType = VectorType::Builder(type).dropDim(0);
-  Type insType = lowType.getRank() > 1 ? VectorType::Builder(lowType).dropDim(0)
-                                       : lowType.getElementType();
+  Type lowType = VectorType::Builder(type).dropDim(0);
+  auto vType = cast<VectorType>(lowType);
   for (int64_t d = 0, e = type.getDimSize(0); d < e; d++) {
-    Value ext = rewriter.create<vector::ExtractOp>(loc, lowType, result, d);
-    Value ins = rewriter.create<vector::ExtractOp>(loc, insType, val, d);
-    Value sto = reshapeStore(loc, ins, ext, lowType, index - 1, pos, rewriter);
-    result = rewriter.create<vector::InsertOp>(loc, type, sto, result, d);
+    Value ext = rewriter.create<vector::ExtractOp>(loc, result, d);
+    Value ins = rewriter.create<vector::ExtractOp>(loc, val, d);
+    Value sto = reshapeStore(loc, ins, ext, vType, index - 1, pos, rewriter);
+    result = rewriter.create<vector::InsertOp>(loc, sto, result, d);
   }
   return result;
 }
 
 /// Helper to create arithmetic operation associated with a kind of contraction.
 static std::optional<Value>
 createContractArithOp(Location loc, Value x, Value y, Value acc,
                       vector::CombiningKind kind, PatternRewriter &rewriter,
for (int64_t d = 0, e = resType.getDimSize(0); d < e; ++d) {
       Value extrMask;
       if (mask)
         extrMask = rewriter.create<vector::ExtractOp>(loc, mask, d);
 
       std::optional<Value> m = createContractArithOp(
           loc, a, op.getRhs(), r, kind, rewriter, isInt, extrMask);
       if (!m.has_value())
         return failure();
-      result = rewriter.create<vector::InsertOp>(loc, resType, *m, result, d);
+      result = rewriter.create<vector::InsertOp>(loc, *m, result, d);
     }
 
     rewriter.replaceOp(rootOp, result);
     return success();
   }
 };
 
 /// Progressively lower a `vector.contract %a, %b, %c` with row-major matmul

mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp

LogicalResult matchAndRewrite(vector::CreateMaskOp op,
     Value result = rewriter.create<arith::ConstantOp>(
         loc, dstType, rewriter.getZeroAttr(dstType));
     for (int64_t d = 0; d < dim; d++) {
       Value bnd =
           rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(d));
       Value val = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt,
                                                  bnd, idx);
       Value sel = rewriter.create<arith::SelectOp>(loc, val, trueVal, falseVal);
-      result = rewriter.create<vector::InsertOp>(loc, dstType, sel, result, d);
+      result = rewriter.create<vector::InsertOp>(loc, sel, result, d);
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
 /// Progressive lowering of ConstantMaskOp.
 /// One:
LogicalResult matchAndRewrite(vector::ConstantMaskOp op,
     SmallVector<int64_t> newDimSizes;
     for (int64_t r = 1; r < rank; r++)
       newDimSizes.push_back(cast<IntegerAttr>(dimSizes[r]).getInt());
     Value trueVal = rewriter.create<vector::ConstantMaskOp>(
         loc, lowType, rewriter.getI64ArrayAttr(newDimSizes));
     Value result = rewriter.create<arith::ConstantOp>(
         loc, dstType, rewriter.getZeroAttr(dstType));
     for (int64_t d = 0; d < trueDim; d++)
-      result =
-          rewriter.create<vector::InsertOp>(loc, dstType, trueVal, result, d);
+      result = rewriter.create<vector::InsertOp>(loc, trueVal, result, d);
+
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 } // namespace
 
 void mlir::vector::populateVectorMaskOpLoweringPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
     VectorType extractSrcType = extractOp.getSourceVectorType();
     Location loc = extractOp.getLoc();
 
     // "vector.extract %v[] : vector<f32>" is an invalid op.
     assert(extractSrcType.getRank() > 0 &&
            "vector.extract does not support rank 0 sources");
 
     // "vector.extract %v[] : vector<...xf32>" can be canonicalized to %v.
-    if (extractOp.getPosition().empty())
+    if (extractOp.getNumIndices() == 0)
       return failure();
 
     // Rewrite vector.extract with 1d source to vector.extractelement.
     if (extractSrcType.getRank() == 1) {
-      assert(extractOp.getPosition().size() == 1 && "expected 1 index");
-      int64_t pos = extractOp.getPosition()[0];
+      assert(!extractOp.hasDynamicPosition() && "Unsupported dynamic index");
+      assert(extractOp.getNumIndices() == 1 && "expected 1 index");
+      int64_t pos = extractOp.getStaticPosition()[0];
       rewriter.setInsertionPoint(extractOp);
       rewriter.replaceOpWithNewOp<vector::ExtractElementOp>(
           extractOp, extractOp.getVector(),
           rewriter.create<arith::ConstantIndexOp>(loc, pos));
       return success();
     }
 
     // All following cases are 2d or higher dimensional source vectors.
 
     if (warpOp.getResult(operandNumber).getType() == operand->get().getType()) {
       // There is no distribution, this is a broadcast. Simply move the extract
       // out of the warp op.
       // TODO: This could be optimized. E.g., in case of a scalar result, let
       // one lane extract and shuffle the result to all other lanes (same as
       // the 1d case).
       SmallVector<size_t> newRetIndices;
       WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
           rewriter, warpOp, {extractOp.getVector()},
           {extractOp.getSourceVectorType()}, newRetIndices);
       rewriter.setInsertionPointAfter(newWarpOp);
       Value distributedVec = newWarpOp->getResult(newRetIndices[0]);
       // Extract from distributed vector.
       Value newExtract = rewriter.create<vector::ExtractOp>(
-          loc, distributedVec, extractOp.getPosition());
+          loc, distributedVec, extractOp.getMixedPosition());
       rewriter.replaceAllUsesWith(newWarpOp->getResult(operandNumber),
                                   newExtract);
       return success();
     }
 
     // Find the distributed dimension. There should be exactly one.
     auto distributedType =
         cast<VectorType>(warpOp.getResult(operandNumber).getType());
LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
     }
     assert(distributedDim != -1 && "could not find distributed dimension");
     (void)distributedDim;
 
     // Yield source vector from warp op.
     SmallVector<int64_t> newDistributedShape(extractSrcType.getShape().begin(),
                                              extractSrcType.getShape().end());
     for (int i = 0; i < distributedType.getRank(); ++i)
-      newDistributedShape[i + extractOp.getPosition().size()] =
+      newDistributedShape[i + extractOp.getNumIndices()] =
           distributedType.getDimSize(i);
     auto newDistributedType =
         VectorType::get(newDistributedShape, distributedType.getElementType());
     SmallVector<size_t> newRetIndices;
     WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
         rewriter, warpOp, {extractOp.getVector()}, {newDistributedType},
         newRetIndices);
     rewriter.setInsertionPointAfter(newWarpOp);
     Value distributedVec = newWarpOp->getResult(newRetIndices[0]);
     // Extract from distributed vector.
     Value newExtract = rewriter.create<vector::ExtractOp>(
-        loc, distributedVec, extractOp.getPosition());
+        loc, distributedVec, extractOp.getMixedPosition());
     rewriter.replaceAllUsesWith(newWarpOp->getResult(operandNumber),
                                 newExtract);
     return success();
   }
 };
 
 /// Pattern to move out vector.extractelement of 0-D tensors. Those don't
 /// need to be distributed and can just be propagated outside of the region.
OpOperand *operand = getWarpResult(
         warpOp, [](Operation *op) { return isa<vector::InsertOp>(op); });
     if (!operand)
       return failure();
     unsigned int operandNumber = operand->getOperandNumber();
     auto insertOp = operand->get().getDefiningOp<vector::InsertOp>();
     Location loc = insertOp.getLoc();
 
     // "vector.insert %v, %v[] : ..." can be canonicalized to %v.
-    if (insertOp.getPosition().empty())
+    if (insertOp.getNumIndices() == 0)
       return failure();
 
     // Rewrite vector.insert with 1d dest to vector.insertelement.
     if (insertOp.getDestVectorType().getRank() == 1) {
-      assert(insertOp.getPosition().size() == 1 && "expected 1 index");
-      int64_t pos = insertOp.getPosition()[0];
+      assert(!insertOp.hasDynamicPosition() && "Unsupported dynamic position");
+      assert(insertOp.getNumIndices() == 1 && "expected 1 index");
+      int64_t pos = insertOp.getStaticPosition()[0];
       rewriter.setInsertionPoint(insertOp);
       rewriter.replaceOpWithNewOp<vector::InsertElementOp>(
           insertOp, insertOp.getSource(), insertOp.getDest(),
           rewriter.create<arith::ConstantIndexOp>(loc, pos));
       return success();
     }
 
     if (warpOp.getResult(operandNumber).getType() == operand->get().getType()) {
       // There is no distribution, this is a broadcast. Simply move the insert
       // out of the warp op.
       SmallVector<size_t> newRetIndices;
       WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
           rewriter, warpOp, {insertOp.getSource(), insertOp.getDest()},
           {insertOp.getSourceType(), insertOp.getDestVectorType()},
           newRetIndices);
       rewriter.setInsertionPointAfter(newWarpOp);
       Value distributedSrc = newWarpOp->getResult(newRetIndices[0]);
       Value distributedDest = newWarpOp->getResult(newRetIndices[1]);
       Value newResult = rewriter.create<vector::InsertOp>(
-          loc, distributedSrc, distributedDest, insertOp.getPosition());
+          loc, distributedSrc, distributedDest, insertOp.getMixedPosition());
       rewriter.replaceAllUsesWith(newWarpOp->getResult(operandNumber),
                                   newResult);
       return success();
     }
 
     // Find the distributed dimension. There should be exactly one.
     auto distrDestType =
         cast<VectorType>(warpOp.getResult(operandNumber).getType());
LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
     SmallVector<int64_t> distrSrcShape(srcVecType.getShape().begin(),
                                        srcVecType.getShape().end());
     // E.g.: vector.insert %s, %d [2] : vector<96xf32> into vector<128x96xf32>
     // Case 1: distrDestDim = 1 (dim of size 96). In that case, each lane will
     //         insert a smaller vector<3xf32>.
     // Case 2: distrDestDim = 0 (dim of size 128) => distrSrcDim = -1. In that
     //         case, one lane will insert the source vector<96xf32>. The other
     //         lanes will not do anything.
-    int64_t distrSrcDim = distrDestDim - insertOp.getPosition().size();
+    int64_t distrSrcDim = distrDestDim - insertOp.getNumIndices();
     if (distrSrcDim >= 0)
       distrSrcShape[distrSrcDim] = distrDestType.getDimSize(distrDestDim);
     auto distrSrcType =
         VectorType::get(distrSrcShape, distrDestType.getElementType());
 
     // Yield source and dest vectors from warp op.
     SmallVector<size_t> newRetIndices;
     WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
         rewriter, warpOp, {insertOp.getSource(), insertOp.getDest()},
         {distrSrcType, distrDestType}, newRetIndices);
     rewriter.setInsertionPointAfter(newWarpOp);
     Value distributedSrc = newWarpOp->getResult(newRetIndices[0]);
     Value distributedDest = newWarpOp->getResult(newRetIndices[1]);
 
     // Insert into the distributed vector.
     Value newResult;
     if (distrSrcDim >= 0) {
       // Every lane inserts a small piece.
       newResult = rewriter.create<vector::InsertOp>(
-          loc, distributedSrc, distributedDest, insertOp.getPosition());
+          loc, distributedSrc, distributedDest, insertOp.getMixedPosition());
     } else {
       // One lane inserts the entire source vector.
       int64_t elementsPerLane = distrDestType.getDimSize(distrDestDim);
-      SmallVector<int64_t> newPos(insertOp.getPosition());
+      SmallVector<OpFoldResult> pos = insertOp.getMixedPosition();
+      SmallVector<int64_t> newPos = getAsIntegers(pos);
       // tid of inserting lane: pos / elementsPerLane
       Value insertingLane = rewriter.create<arith::ConstantIndexOp>(
           loc, newPos[distrDestDim] / elementsPerLane);
       Value isInsertingLane = rewriter.create<arith::CmpIOp>(
           loc, arith::CmpIPredicate::eq, newWarpOp.getLaneid(), insertingLane);
       // Insert position: pos % elementsPerLane
       newPos[distrDestDim] %= elementsPerLane;
       auto insertingBuilder = [&](OpBuilder &builder, Location loc) {

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

 //===- VectorDropLeadUnitDim.cpp - Conversion within the Vector dialect ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/TypeUtilities.h"
 
LogicalResult matchAndRewrite(vector::InsertOp insertOp,
     }
     Value newDstVector = rewriter.create<vector::ExtractOp>(
         loc, insertOp.getDest(), splatZero(dstDropCount));
 
     // New position rank needs to be computed in two steps: (1) if destination
     // type has leading unit dims, we also trim the position array accordingly,
     // then (2) if source type also has leading unit dims, we need to append
     // zeroes to the position array accordingly.
-    unsigned oldPosRank = insertOp.getPosition().size();
+    unsigned oldPosRank = insertOp.getNumIndices();
     unsigned newPosRank = std::max<int64_t>(0, oldPosRank - dstDropCount);
-    SmallVector<int64_t> newPositions =
-        llvm::to_vector(insertOp.getPosition().take_back(newPosRank));
-    newPositions.resize(newDstType.getRank() - newSrcRank, 0);
+    SmallVector<OpFoldResult> oldPosition = insertOp.getMixedPosition();
+    SmallVector<OpFoldResult> newPosition =
+        llvm::to_vector(ArrayRef(oldPosition).take_back(newPosRank));
+    newPosition.resize(newDstType.getRank() - newSrcRank,
+                       rewriter.getI64IntegerAttr(0));
 
     auto newInsertOp = rewriter.create<vector::InsertOp>(
-        loc, newDstType, newSrcVector, newDstVector, newPositions);
+        loc, newSrcVector, newDstVector, newPosition);
 
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
                                                      newInsertOp);
 
     return success();
   }
 };
 

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

using RewriteScalarExtractOfTransferReadBase::
       RewriteScalarExtractOfTransferReadBase;
 
   void rewrite(vector::ExtractOp extractOp,
                PatternRewriter &rewriter) const override {
     // Construct scalar load.
     auto xferOp = extractOp.getVector().getDefiningOp<vector::TransferReadOp>();
     SmallVector<Value> newIndices(xferOp.getIndices().begin(),
                                   xferOp.getIndices().end());
-    for (const auto &it : llvm::enumerate(extractOp.getPosition())) {
-      int64_t offset = it.value();
-      int64_t idx =
-          newIndices.size() - extractOp.getPosition().size() + it.index();
+    for (auto [i, pos] : llvm::enumerate(extractOp.getMixedPosition())) {
+      assert(pos.is<Attribute>() && "Unespected non-constant index");
+      int64_t offset = cast<IntegerAttr>(pos.get<Attribute>()).getInt();
+      int64_t idx = newIndices.size() - extractOp.getNumIndices() + i;
       OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
           rewriter, extractOp.getLoc(),
           rewriter.getAffineSymbolExpr(0) + offset, {newIndices[idx]});
       if (ofr.is<Value>()) {
         newIndices[idx] = ofr.get<Value>();
       } else {
         newIndices[idx] = rewriter.create<arith::ConstantIndexOp>(
             extractOp.getLoc(), *getConstantIntValue(ofr));

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
     // Only support casting to a larger number of elements or now.
     // E.g., vector<4xf32> -> vector<8xf16>.
     if (castSrcType.getNumElements() > castDstType.getNumElements())
       return failure();
 
     unsigned expandRatio =
         castDstType.getNumElements() / castSrcType.getNumElements();
 
-    uint64_t index = extractOp.getPosition()[0];
+    auto getFirstIntValue = [](ArrayRef<OpFoldResult> values) -> uint64_t {
+      assert(values[0].is<Attribute>() && "Unexpected non-constant index");
+      return cast<IntegerAttr>(values[0].get<Attribute>()).getInt();
+    };
+
+    uint64_t index = getFirstIntValue(extractOp.getMixedPosition());
 
     // Get the single scalar (as a vector) in the source value that packs the
     // desired scalar. E.g. extract vector<1xf32> from vector<4xf32>
-    VectorType oneScalarType =
-        VectorType::get({1}, castSrcType.getElementType());
+    Location loc = extractOp.getLoc();
     Value packedValue = rewriter.create<vector::ExtractOp>(
-        extractOp.getLoc(), oneScalarType, castOp.getSource(),
-        index / expandRatio);
+        loc, castOp.getSource(), index / expandRatio);
+    Type packedVecType = VectorType::get(/*shape=*/{1}, packedValue.getType());
+    Value zero = rewriter.create<arith::ConstantOp>(
+        loc, packedVecType, rewriter.getZeroAttr(packedVecType));
+    packedValue = rewriter.create<vector::InsertOp>(loc, packedValue, zero,
+                                                    /*position=*/0);
 
     // Cast it to a vector with the desired scalar's type.
     // E.g. f32 -> vector<2xf16>
     VectorType packedType =
         VectorType::get({expandRatio}, castDstType.getElementType());
-    Value castedValue = rewriter.create<vector::BitCastOp>(
-        extractOp.getLoc(), packedType, packedValue);
+    Value castedValue =
+        rewriter.create<vector::BitCastOp>(loc, packedType, packedValue);
 
     // Finally extract the desired scalar.
-    rewriter.replaceOpWithNewOp<vector::ExtractOp>(
-        extractOp, extractOp.getType(), castedValue, index % expandRatio);
-
+    rewriter.replaceOpWithNewOp<vector::ExtractOp>(extractOp, castedValue,
+                                                   index % expandRatio);
     return success();
   }
 };
 
 // Shuffles vector.bitcast op after vector.extract_strided_slice op.
 //
 // This transforms IR like:
 //    %cast = vector.bitcast %arg0: vector<4xf32> to vector<8xf16>

mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir

 
 // -----
 
 // CHECK-LABEL: @extract
 //  CHECK-SAME: %[[ARG:.+]]: vector<2xf32>
 //       CHECK:   spirv.CompositeExtract %[[ARG]][0 : i32] : vector<2xf32>
 //       CHECK:   spirv.CompositeExtract %[[ARG]][1 : i32] : vector<2xf32>
 func.func @extract(%arg0 : vector<2xf32>) -> (vector<1xf32>, f32) {
-  %0 = "vector.extract"(%arg0) <{position = array<i64: 0>}> : (vector<2xf32>) -> vector<1xf32>
-  %1 = "vector.extract"(%arg0) <{position = array<i64: 1>}> : (vector<2xf32>) -> f32
+  %0 = "vector.extract"(%arg0) <{static_position = array<i64: 0>}> : (vector<2xf32>) -> vector<1xf32>
+  %1 = "vector.extract"(%arg0) <{static_position = array<i64: 1>}> : (vector<2xf32>) -> f32
   return %0, %1: vector<1xf32>, f32
 }
 
 // -----
 
 // CHECK-LABEL: @extract_size1_vector
 //  CHECK-SAME: %[[ARG0:.+]]: vector<1xf32>
 //       CHECK:   %[[R:.+]] = builtin.unrealized_conversion_cast %[[ARG0]]

mlir/test/Dialect/Vector/invalid.mlir

func.func @extract_position_rank_overflow(%arg0: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute of rank no greater than vector rank}}
   %1 = vector.extract %arg0[0, 0, 0, 0] : vector<4x8x16xf32>
 }
 
 // -----
 
 func.func @extract_position_rank_overflow_generic(%arg0: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute of rank no greater than vector rank}}
-  %1 = "vector.extract" (%arg0) <{position = array<i64: 0, 0, 0, 0>}> : (vector<4x8x16xf32>) -> (vector<16xf32>)
+  %1 = "vector.extract" (%arg0) <{static_position = array<i64: 0, 0, 0, 0>}> : (vector<4x8x16xf32>) -> (vector<16xf32>)
 }
 
 // -----
 
 func.func @extract_position_overflow(%arg0: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute #2 to be a non-negative integer smaller than the corresponding vector dimension}}
   %1 = vector.extract %arg0[0, 43, 0] : vector<4x8x16xf32>
 }

mlir/test/Dialect/Vector/vector-transforms.mlir

func.func @contraction4x4_ikj_xfer_read_tensor(%arg0 : tensor<4x2xf32>,
   return %r : tensor<4x4xf32>
 }
 
 // CHECK-LABEL: func @bubble_down_bitcast_in_extract
 //  CHECK-SAME: %[[SRC:.+]]: vector<4xf32>
 func.func @bubble_down_bitcast_in_extract(%src: vector<4xf32>) -> (f16, f16) {
   %0 = vector.bitcast %src : vector<4xf32> to vector<8xf16>
   // CHECK: %[[EXTRACT1:.+]] = vector.extract %[[SRC]][1] : vector<4xf32>
-  // CHECK:    %[[CAST1:.+]] = vector.bitcast %[[EXTRACT1]] : vector<1xf32> to vector<2xf16>
+  // CHECK:  %[[INSERT1:.+]] = vector.insert %[[EXTRACT1]], %{{.+}} [0] : f32 into vector<1xf32>
+  // CHECK:    %[[CAST1:.+]] = vector.bitcast %[[INSERT1]] : vector<1xf32> to vector<2xf16>
   // CHECK: %[[EXTRACT2:.+]] = vector.extract %[[CAST1]][1] : vector<2xf16>
   %1 = vector.extract %0[3] : vector<8xf16>
   // CHECK: %[[EXTRACT3:.+]] = vector.extract %[[SRC]][2] : vector<4xf32>
-  // CHECK:    %[[CAST2:.+]] = vector.bitcast %[[EXTRACT3]] : vector<1xf32> to vector<2xf16>
+  // CHECK:  %[[INSERT3:.+]] = vector.insert %[[EXTRACT3]], %{{.+}} [0] : f32 into vector<1xf32>
+  // CHECK:    %[[CAST2:.+]] = vector.bitcast %[[INSERT3]] : vector<1xf32> to vector<2xf16>
   // CHECK: %[[EXTRACT4:.+]] = vector.extract %[[CAST2]][0] : vector<2xf16>
   %2 = vector.extract %0[4] : vector<8xf16>
   // CHECK: return %[[EXTRACT2]], %[[EXTRACT4]]
   return %1, %2: f16, f16
 }
 
 // CHECK-LABEL: func @bubble_down_bitcast_in_strided_slice_extract
 //  CHECK-SAME: %[[SRC:.+]]: vector<4xf32>