[MLIR][SPIRVToLLVM] Implementation of spv.BitFieldSExtract and spv.BitFieldUExtract patterns

[georgemitenkov]

Summary

This patch adds conversion patterns for spv.BitFieldSExtract and spv.BitFieldUExtract. As in the patch for spv.BitFieldInsert, offset and count have to be broadcasted in vector case and casted to match the type of the base.

[rriddle]

It doesn't seem like you are actually testing that these patterns replace the op.

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

238

nit: Drop trivial braces.

265

Where are you replacing op?

319

Same here.

[georgemitenkov]

• [MLIR][SPIRVToLLVM] Addressed comments on style and op replacement
• [MLIR][SPIRVToLLVM] Refactored the code to follow clang style

Removed trivial braces in conditional statement. Used replaceOpWithNewOp to replace the op.

[antiagainst]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

219

The logic to broadcast and bitwidth-cast is the same for all there bitfield ops. Could we have a utility function to deduplicate them?

236

Nit: . to end.

mlir/test/Conversion/SPIRVToLLVM/bitwise-ops-to-llvm.mlir

205

Same as the previous patch. I think we can drop the bitwidth case permutation for vector cases.

[georgemitenkov]

• [MLIR][SPIRVToLLVM] Addressed style comments on bit manipulation patterns and removed unnecessary tests

Added a full stop in the comment. Reduced the number of vector tests to one. I suppose that for scalars it can also be the case? The functuinality is the same accross all 3 bitfiled ops. So we can have 3 scalar cases for spv.BitFieldInsert and 1 for extraction ops.
Broadcasting and truncating/extending will be moved to utility once BitFieldInsertPattern is merged.

[antiagainst]

In D82640#2122564, @georgemitenkov wrote:

• [MLIR][SPIRVToLLVM] Addressed style comments on bit manipulation patterns and removed unnecessary tests

Added a full stop in the comment. Reduced the number of vector tests to one. I suppose that for scalars it can also be the case? The functuinality is the same accross all 3 bitfiled ops. So we can have 3 scalar cases for spv.BitFieldInsert and 1 for extraction ops.
Broadcasting and truncating/extending will be moved to utility once BitFieldInsertPattern is merged.

It is merged now. Please rebase and address issues. :)

[georgemitenkov]

• [MLIR][SPIRVToLLVM] Extracted common logic for broadcasting and extension/truncation for bit ops
• [MLIR][SPIRVToLLVM] Refactored the code to follow clang style

Rebased against master. Extracted common functionality for bit manipulation ops to utilility functions. Now, both offset and count are processed via processCountOrOffset() function, that would call optionallyBroadcast() and optionallyTruncateOrExtend().
Also, processing offset and count is more logical now: offset first and then count (before the ordering was different for broadcasting and extension/truncation).

[antiagainst]

Nice! I believe you can land it by yourself now. So I won't spoil the fun here and will let you do it by yourself. :) Please make sure to follow https://llvm.org/docs/Phabricator.html#committing-a-change.

[georgemitenkov]

I see that the revision hasn't been landed due to changes picked up when rebasing against master (I think this is because I was committing via arc patch D*). Do I reopen the review in this case?

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

if (srcType.isa<VectorType>())
     return rewriter.create<LLVM::ConstantOp>(
         loc, dstType,
         SplatElementsAttr::get(srcType.cast<ShapedType>(),
                                minusOneIntegerAttribute(srcType, rewriter)));
   return rewriter.create<LLVM::ConstantOp>(
       loc, dstType, minusOneIntegerAttribute(srcType, rewriter));
 }
 
-/// This is a utility function for bit manipulations ops (`BitFieldInsert`)
-/// and operates on their `Count` or `Offset` values. It casts the given
-/// value to match the target type.
+/// This is a utility function for bit manipulations ops (`BitFieldInsert`,
+/// `BitFieldSExtract` and `BitFieldUExtract`), and operates on their `Count`
+/// or `Offset` values. It casts the given value to match the target type.
 static Value optionallyCast(Location loc, Value value, Type dstType,
                             ConversionPatternRewriter &rewriter) {
   auto llvmType = dstType.cast<LLVM::LLVMType>();
   unsigned targetBitWidth =
       llvmType.isVectorTy()
           ? llvmType.getVectorElementType()
                 .getUnderlyingType()
                 ->getIntegerBitWidth()
Value insertShiftedByOffset = rewriter.create<LLVM::ShlOp>(
         loc, dstType, op.insert(), optionallyCastedOffset);
     rewriter.create<LLVM::OrOp>(loc, dstType, baseAndMask,
                                 insertShiftedByOffset);
     rewriter.eraseOp(op);
     return success();
   }
 };
 
+class BitFieldSExtractPattern
+    : public SPIRVToLLVMConversion<spirv::BitFieldSExtractOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::BitFieldSExtractOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::BitFieldSExtractOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto srcType = op.getType();
+    auto dstType = this->typeConverter.convertType(srcType);
+    if (!dstType)
+      return failure();
+    Location loc = op.getLoc();
+
+    // Broadcast `Offset` and `Count` to match the type of `Base`.
+    // If `Base` is of a vector type, construct a vector that has:
+    //  - same number of elements as `Base`
+    //  - each element has the type that is the same as the type of `Offset` or
+    // `Count`
+    //  - each element has the same value as `Offset` or `Count`
+    Value offset;
+    Value count;
+    if (auto vectorType = srcType.dyn_cast<VectorType>()) {
+      int64_t numElements = static_cast<int64_t>(vectorType.getNumElements());
+      broadcast(loc, op.offset(), offset, numElements, typeConverter, rewriter);
+      broadcast(loc, op.count(), count, numElements, typeConverter, rewriter);
+    } else {
+      offset = op.offset();
+      count = op.count();
+    }
+
+    // Need to zero extend `Offset` and `Count` if their bit width is different
+    // from the result bit width.
+    Value optionallyCastedCount = optionallyCast(loc, count, dstType, rewriter);
+    Value optionallyCastedOffset =
+        optionallyCast(loc, offset, dstType, rewriter);
+
+    // Create a constant that holds the size of the `Base`.
+    IntegerType integerType;
+    if (auto vecType = srcType.dyn_cast<VectorType>())
+      integerType = vecType.getElementType().cast<IntegerType>();
+    else
+      integerType = srcType.cast<IntegerType>();
+
+    auto baseSize = rewriter.getIntegerAttr(integerType, getBitWidth(srcType));
+    Value size =
+        srcType.isa<VectorType>()
+            ? rewriter.create<LLVM::ConstantOp>(
+                  loc, dstType,
+                  SplatElementsAttr::get(srcType.cast<ShapedType>(), baseSize))
+            : rewriter.create<LLVM::ConstantOp>(loc, dstType, baseSize);
+
+    // Shift `Base` left by [sizeof(Base) - (Count + Offset)], so that the bit
+    // at Offset + Count - 1 is the most significant bit now.
+    Value countPlusOffset = rewriter.create<LLVM::AddOp>(
+        loc, dstType, optionallyCastedCount, optionallyCastedOffset);
+    Value amountToShiftLeft =
+        rewriter.create<LLVM::SubOp>(loc, dstType, size, countPlusOffset);
+    Value baseShiftedLeft = rewriter.create<LLVM::ShlOp>(
+        loc, dstType, op.base(), amountToShiftLeft);
+
+    // Shift the result right, filling the bits with the sign bit.
+    Value amountToShiftRight = rewriter.create<LLVM::AddOp>(
+        loc, dstType, optionallyCastedOffset, amountToShiftLeft);
+    rewriter.replaceOpWithNewOp<LLVM::AShrOp>(op, dstType, baseShiftedLeft,
+                                              amountToShiftRight);
+    return success();
+  }
+};
+
+class BitFieldUExtractPattern
+    : public SPIRVToLLVMConversion<spirv::BitFieldUExtractOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::BitFieldUExtractOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::BitFieldUExtractOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto srcType = op.getType();
+    auto dstType = this->typeConverter.convertType(srcType);
+    if (!dstType)
+      return failure();
+    Location loc = op.getLoc();
+
+    // Broadcast `Offset` and `Count` to match the type of `Base`.
+    // If `Base` is of a vector type, construct a vector that has:
+    //  - same number of elements as `Base`
+    //  - each element has the type that is the same as the type of `Offset` or
+    // `Count`
+    //  - each element has the same value as `Offset` or `Count`
+    Value offset;
+    Value count;
+    if (auto vectorType = srcType.dyn_cast<VectorType>()) {
+      int64_t numElements = static_cast<int64_t>(vectorType.getNumElements());
+      broadcast(loc, op.offset(), offset, numElements, typeConverter, rewriter);
+      broadcast(loc, op.count(), count, numElements, typeConverter, rewriter);
+    } else {
+      offset = op.offset();
+      count = op.count();
+    }
+
+    // Create a mask with all bits set of the same type as `srcType`
+    Value minusOne = createConstantAllBitsSet(loc, srcType, dstType, rewriter);
+
+    // Need to cast `Offset` and `Count` if their bit width is different from
+    // the result bit width.
+    Value optionallyCastedCount = optionallyCast(loc, count, dstType, rewriter);
+    Value optionallyCastedOffset =
+        optionallyCast(loc, offset, dstType, rewriter);
+
+    // Create a mask with bits set at [0, Count - 1].
+    Value maskShiftedByCount = rewriter.create<LLVM::ShlOp>(
+        loc, dstType, minusOne, optionallyCastedCount);
+    Value mask = rewriter.create<LLVM::XOrOp>(loc, dstType, maskShiftedByCount,
+                                              minusOne);
+
+    // Shift `Base` by `Offset` and apply the mask on it.
+    Value shiftedBase = rewriter.create<LLVM::LShrOp>(loc, dstType, op.base(),
+                                                      optionallyCastedOffset);
+    rewriter.replaceOpWithNewOp<LLVM::AndOp>(op, dstType, shiftedBase, mask);
+    return success();
+  }
+};
+
 /// Converts SPIR-V operations that have straightforward LLVM equivalent
 /// into LLVM dialect operations.
 template <typename SPIRVOp, typename LLVMOp>
 class DirectConversionPattern : public SPIRVToLLVMConversion<SPIRVOp> {
 public:
   using SPIRVToLLVMConversion<SPIRVOp>::SPIRVToLLVMConversion;
 
   LogicalResult
patterns.insert<
       DirectConversionPattern<spirv::FRemOp, LLVM::FRemOp>,
       DirectConversionPattern<spirv::FSubOp, LLVM::FSubOp>,
       DirectConversionPattern<spirv::SDivOp, LLVM::SDivOp>,
       DirectConversionPattern<spirv::SRemOp, LLVM::SRemOp>,
       DirectConversionPattern<spirv::UDivOp, LLVM::UDivOp>,
       DirectConversionPattern<spirv::UModOp, LLVM::URemOp>,
 
       // Bitwise ops
-      BitFieldInsertPattern,
+      BitFieldInsertPattern, BitFieldUExtractPattern, BitFieldSExtractPattern,
       DirectConversionPattern<spirv::BitwiseAndOp, LLVM::AndOp>,
       DirectConversionPattern<spirv::BitwiseOrOp, LLVM::OrOp>,
       DirectConversionPattern<spirv::BitwiseXorOp, LLVM::XOrOp>,
       NotPattern<spirv::NotOp>,
 
       // Cast ops
       DirectConversionPattern<spirv::ConvertFToSOp, LLVM::FPToSIOp>,
       DirectConversionPattern<spirv::ConvertFToUOp, LLVM::FPToUIOp>,

mlir/test/Conversion/SPIRVToLLVM/bitwise-ops-to-llvm.mlir

func @bitfield_insert_vector_greater_bit_width(%base: vector<2xi32>, %insert: vector<2xi32>, %offset: i64, %count: i64) {
     // CHECK: %[[NEW_BASE:.*]] = llvm.and %[[BASE]], %[[MASK]] : !llvm<"<2 x i32>">
     // CHECK: %[[SHIFTED_INSERT:.*]] = llvm.shl %[[INSERT]], %[[TRUNC_OFFSET]] : !llvm<"<2 x i32>">
     // CHECK: %{{.*}} = llvm.or %[[NEW_BASE]], %[[SHIFTED_INSERT]] : !llvm<"<2 x i32>">
     %0 = spv.BitFieldInsert %base, %insert, %offset, %count : vector<2xi32>, i64, i64
     return
 }
 
 //===----------------------------------------------------------------------===//
+// spv.BitFieldSExtract
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @bitfield_sextract_scalar_same_bit_width
+// CHECK-SAME: %[[BASE:.*]]: !llvm.i64, %[[OFFSET:.*]]: !llvm.i64, %[[COUNT:.*]]: !llvm.i64
+func @bitfield_sextract_scalar_same_bit_width(%base: i64, %offset: i64, %count: i64) {
+    // CHECK: %[[SIZE:.]] = llvm.mlir.constant(64 : i64) : !llvm.i64
+    // CHECK: %[[T0:.*]] = llvm.add %[[COUNT]], %[[OFFSET]] : !llvm.i64
+    // CHECK: %[[T1:.*]] = llvm.sub %[[SIZE]], %[[T0]] : !llvm.i64
+    // CHECK: %[[SHIFTED_LEFT:.*]] = llvm.shl %[[BASE]], %[[T1]] : !llvm.i64
+    // CHECK: %[[T2:.*]] = llvm.add %[[OFFSET]], %[[T1]] : !llvm.i64
+    // CHECK: %{{.*}} = llvm.ashr %[[SHIFTED_LEFT]], %[[T2]] : !llvm.i64
+    %0 = spv.BitFieldSExtract %base, %offset, %count : i64, i64, i64
+    return
+}
+
+// CHECK-LABEL: func @bitfield_sextract_scalar_smaller_bit_width
+// CHECK-SAME: %[[BASE:.*]]: !llvm.i32, %[[OFFSET:.*]]: !llvm.i8, %[[COUNT:.*]]: !llvm.i8
+func @bitfield_sextract_scalar_smaller_bit_width(%base: i32, %offset: i8, %count: i8) {
+    // CHECK: %[[EXT_COUNT:.*]] = llvm.zext %[[COUNT]] : !llvm.i8 to !llvm.i32
+    // CHECK: %[[EXT_OFFSET:.*]] = llvm.zext %[[OFFSET]] : !llvm.i8 to !llvm.i32
+    // CHECK: %[[SIZE:.]] = llvm.mlir.constant(32 : i32) : !llvm.i32
+    // CHECK: %[[T0:.*]] = llvm.add %[[EXT_COUNT]], %[[EXT_OFFSET]] : !llvm.i32
+    // CHECK: %[[T1:.*]] = llvm.sub %[[SIZE]], %[[T0]] : !llvm.i32
+    // CHECK: %[[SHIFTED_LEFT:.*]] = llvm.shl %[[BASE]], %[[T1]] : !llvm.i32
+    // CHECK: %[[T2:.*]] = llvm.add %[[EXT_OFFSET]], %[[T1]] : !llvm.i32
+    // CHECK: %{{.*}} = llvm.ashr %[[SHIFTED_LEFT]], %[[T2]] : !llvm.i32
+    %0 = spv.BitFieldSExtract %base, %offset, %count : i32, i8, i8
+    return
+}
+
+// CHECK-LABEL: func @bitfield_sextract_scalar_greater_bit_width
+// CHECK-SAME: %[[BASE:.*]]: !llvm.i32, %[[OFFSET:.*]]: !llvm.i64, %[[COUNT:.*]]: !llvm.i64
+func @bitfield_sextract_scalar_greater_bit_width(%base: i32, %offset: i64, %count: i64) {
+    // CHECK: %[[TRUNC_COUNT:.*]] = llvm.trunc %[[COUNT]] : !llvm.i64 to !llvm.i32
+    // CHECK: %[[TRUNC_OFFSET:.*]] = llvm.trunc %[[OFFSET]] : !llvm.i64 to !llvm.i32
+    // CHECK: %[[SIZE:.]] = llvm.mlir.constant(32 : i32) : !llvm.i32
+    // CHECK: %[[T0:.*]] = llvm.add %[[TRUNC_COUNT]], %[[TRUNC_OFFSET]] : !llvm.i32
+    // CHECK: %[[T1:.*]] = llvm.sub %[[SIZE]], %[[T0]] : !llvm.i32
+    // CHECK: %[[SHIFTED_LEFT:.*]] = llvm.shl %[[BASE]], %[[T1]] : !llvm.i32
+    // CHECK: %[[T2:.*]] = llvm.add %[[TRUNC_OFFSET]], %[[T1]] : !llvm.i32
+    // CHECK: %{{.*}} = llvm.ashr %[[SHIFTED_LEFT]], %[[T2]] : !llvm.i32
+    %0 = spv.BitFieldSExtract %base, %offset, %count : i32, i64, i64
+    return
+}
+
+// CHECK-LABEL: func @bitfield_sextract_vector
+// CHECK-SAME: %[[BASE:.*]]: !llvm<"<2 x i32>">, %[[OFFSET:.*]]: !llvm.i32, %[[COUNT:.*]]: !llvm.i32
+func @bitfield_sextract_vector(%base: vector<2xi32>, %offset: i32, %count: i32) {
+    // CHECK: %[[OFFSET_V0:.*]] = llvm.mlir.undef : !llvm<"<2 x i32>">
+    // CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i32) : !llvm.i32
+    // CHECK: %[[OFFSET_V1:.*]] = llvm.insertelement %[[OFFSET]], %[[OFFSET_V0]][%[[ZERO]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+    // CHECK: %[[OFFSET_V2:.*]] = llvm.insertelement  %[[OFFSET]], %[[OFFSET_V1]][%[[ONE]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[COUNT_V0:.*]] = llvm.mlir.undef : !llvm<"<2 x i32>">
+    // CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i32) : !llvm.i32
+    // CHECK: %[[COUNT_V1:.*]] = llvm.insertelement %[[COUNT]], %[[COUNT_V0]][%[[ZERO]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+    // CHECK: %[[COUNT_V2:.*]] = llvm.insertelement %[[COUNT]], %[[COUNT_V1]][%[[ONE]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(dense<32> : vector<2xi32>) : !llvm<"<2 x i32>">
+    // CHECK: %[[T0:.*]] = llvm.add %[[COUNT_V2]], %[[OFFSET_V2]] : !llvm<"<2 x i32>">
+    // CHECK: %[[T1:.*]] = llvm.sub %[[SIZE]], %[[T0]] : !llvm<"<2 x i32>">
+    // CHECK: %[[SHIFTED_LEFT:.*]] = llvm.shl %[[BASE]], %[[T1]] : !llvm<"<2 x i32>">
+    // CHECK: %[[T2:.*]] = llvm.add %[[OFFSET_V2]], %[[T1]] : !llvm<"<2 x i32>">
+    // CHECK: %{{.*}} = llvm.ashr %[[SHIFTED_LEFT]], %[[T2]] : !llvm<"<2 x i32>">
+    %0 = spv.BitFieldSExtract %base, %offset, %count : vector<2xi32>, i32, i32
+    return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.BitFieldUExtract
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @bitfield_uextract_scalar_same_bit_width
+// CHECK-SAME: %[[BASE:.*]]: !llvm.i32, %[[OFFSET:.*]]: !llvm.i32, %[[COUNT:.*]]: !llvm.i32
+func @bitfield_uextract_scalar_same_bit_width(%base: i32, %offset: i32, %count: i32) {
+    // CHECK: %[[MINUS_ONE:.*]] = llvm.mlir.constant(-1 : i32) : !llvm.i32
+    // CHECK: %[[T0:.*]] = llvm.shl %[[MINUS_ONE]], %[[COUNT]] : !llvm.i32
+    // CHECK: %[[MASK:.*]] = llvm.xor %[[T0]], %[[MINUS_ONE]] : !llvm.i32
+    // CHECK: %[[SHIFTED_BASE:.*]] = llvm.lshr %[[BASE]], %[[OFFSET]] : !llvm.i32
+    // CHECK: %{{.*}} = llvm.and %[[SHIFTED_BASE]], %[[MASK]] : !llvm.i32
+    %0 = spv.BitFieldUExtract %base, %offset, %count : i32, i32, i32
+    return
+}
+
+// CHECK-LABEL: func @bitfield_uextract_scalar_smaller_bit_width
+// CHECK-SAME: %[[BASE:.*]]: !llvm.i32, %[[OFFSET:.*]]: !llvm.i16, %[[COUNT:.*]]: !llvm.i8
+func @bitfield_uextract_scalar_smaller_bit_width(%base: i32, %offset: i16, %count: i8) {
+    // CHECK: %[[MINUS_ONE:.*]] = llvm.mlir.constant(-1 : i32) : !llvm.i32
+    // CHECK: %[[EXT_COUNT:.*]] = llvm.zext %[[COUNT]] : !llvm.i8 to !llvm.i32
+    // CHECK: %[[EXT_OFFSET:.*]] = llvm.zext %[[OFFSET]] : !llvm.i16 to !llvm.i32
+    // CHECK: %[[T0:.*]] = llvm.shl %[[MINUS_ONE]], %[[EXT_COUNT]] : !llvm.i32
+    // CHECK: %[[MASK:.*]] = llvm.xor %[[T0]], %[[MINUS_ONE]] : !llvm.i32
+    // CHECK: %[[SHIFTED_BASE:.*]] = llvm.lshr %[[BASE]], %[[EXT_OFFSET]] : !llvm.i32
+    // CHECK: %{{.*}} = llvm.and %[[SHIFTED_BASE]], %[[MASK]] : !llvm.i32
+    %0 = spv.BitFieldUExtract %base, %offset, %count : i32, i16, i8
+    return
+}
+
+// CHECK-LABEL: func @bitfield_uextract_scalar_greater_bit_width
+// CHECK-SAME: %[[BASE:.*]]: !llvm.i8, %[[OFFSET:.*]]: !llvm.i16, %[[COUNT:.*]]: !llvm.i8
+func @bitfield_uextract_scalar_greater_bit_width(%base: i8, %offset: i16, %count: i8) {
+    // CHECK: %[[MINUS_ONE:.*]] = llvm.mlir.constant(-1 : i8) : !llvm.i8
+    // CHECK: %[[TRUNC_OFFSET:.*]] = llvm.trunc %[[OFFSET]] : !llvm.i16 to !llvm.i8
+    // CHECK: %[[T0:.*]] = llvm.shl %[[MINUS_ONE]], %[[COUNT]] : !llvm.i8
+    // CHECK: %[[MASK:.*]] = llvm.xor %[[T0]], %[[MINUS_ONE]] : !llvm.i8
+    // CHECK: %[[SHIFTED_BASE:.*]] = llvm.lshr %[[BASE]], %[[TRUNC_OFFSET]] : !llvm.i8
+    // CHECK: %{{.*}} = llvm.and %[[SHIFTED_BASE]], %[[MASK]] : !llvm.i8
+    %0 = spv.BitFieldUExtract %base, %offset, %count : i8, i16, i8
+    return
+}
+
+// CHECK-LABEL: func @bitfield_uextract_vector
+// CHECK-SAME: %[[BASE:.*]]: !llvm<"<2 x i32>">, %[[OFFSET:.*]]: !llvm.i32, %[[COUNT:.*]]: !llvm.i32
+func @bitfield_uextract_vector(%base: vector<2xi32>, %offset: i32, %count: i32) {
+    // CHECK: %[[OFFSET_V0:.*]] = llvm.mlir.undef : !llvm<"<2 x i32>">
+    // CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i32) : !llvm.i32
+    // CHECK: %[[OFFSET_V1:.*]] = llvm.insertelement %[[OFFSET]], %[[OFFSET_V0]][%[[ZERO]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+    // CHECK: %[[OFFSET_V2:.*]] = llvm.insertelement  %[[OFFSET]], %[[OFFSET_V1]][%[[ONE]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[COUNT_V0:.*]] = llvm.mlir.undef : !llvm<"<2 x i32>">
+    // CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i32) : !llvm.i32
+    // CHECK: %[[COUNT_V1:.*]] = llvm.insertelement %[[COUNT]], %[[COUNT_V0]][%[[ZERO]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+    // CHECK: %[[COUNT_V2:.*]] = llvm.insertelement %[[COUNT]], %[[COUNT_V1]][%[[ONE]] : !llvm.i32] : !llvm<"<2 x i32>">
+    // CHECK: %[[MINUS_ONE:.*]] = llvm.mlir.constant(dense<-1> : vector<2xi32>) : !llvm<"<2 x i32>">
+    // CHECK: %[[T0:.*]] = llvm.shl %[[MINUS_ONE]], %[[COUNT_V2]] : !llvm<"<2 x i32>">
+    // CHECK: %[[MASK:.*]] = llvm.xor %[[T0]], %[[MINUS_ONE]] : !llvm<"<2 x i32>">
+    // CHECK: %[[SHIFTED_BASE:.*]] = llvm.lshr %[[BASE]], %[[OFFSET_V2]] : !llvm<"<2 x i32>">
+    // CHECK: %{{.*}} = llvm.and %[[SHIFTED_BASE]], %[[MASK]] : !llvm<"<2 x i32>">
+    %0 = spv.BitFieldUExtract %base, %offset, %count : vector<2xi32>, i32, i32
+    return
+}
+
+//===----------------------------------------------------------------------===//
 // spv.BitwiseAnd
 //===----------------------------------------------------------------------===//
 
 func @bitwise_and_scalar(%arg0: i32, %arg1: i32) {
   // CHECK: %{{.*}} = llvm.and %{{.*}}, %{{.*}} : !llvm.i32
   %0 = spv.BitwiseAnd %arg0, %arg1 : i32
   return
 }