[mlir][arith] Support wide int shrui emulation

[kuhar]

Summary

Tested by checking all 16-bit LHS and all valid RHS when emulating i16 with i8 operations.

[Mogball]

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

469

[kuhar]

Dropped const from a local variable.

[kuhar]

Fixed a typo.

[kuhar]

Removed an unused value & variable.

[kuhar]

Rebased. Fixed documentation comments.

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

 } // namespace mlir::arith
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Common Helper Functions
 //===----------------------------------------------------------------------===//
 
-// Returns N bottom and N top bits from `value`, where N = `newBitWidth`.
-// Treats `value` as a 2*N bits-wide integer.
-// The bottom bits are returned in the first pair element, while the top bits in
-// the second one.
+/// Returns N bottom and N top bits from `value`, where N = `newBitWidth`.
+/// Treats `value` as a 2*N bits-wide integer.
+/// The bottom bits are returned in the first pair element, while the top bits
+/// in the second one.
 static std::pair<APInt, APInt> getHalves(const APInt &value,
                                          unsigned newBitWidth) {
   APInt low = value.extractBits(newBitWidth, 0);
   APInt high = value.extractBits(newBitWidth, newBitWidth);
   return {std::move(low), std::move(high)};
 }
 
-// Returns the type with the last (innermost) dimention reduced to x1.
-// Scalarizes 1D vector inputs to match how we extract/insert vector values,
-// e.g.:
-//   - vector<3x2xi16> --> vector<3x1xi16>
-//   - vector<2xi16>   --> i16
+/// Returns the type with the last (innermost) dimention reduced to x1.
+/// Scalarizes 1D vector inputs to match how we extract/insert vector values,
+/// e.g.:
+///   - vector<3x2xi16> --> vector<3x1xi16>
+///   - vector<2xi16>   --> i16
 static Type reduceInnermostDim(VectorType type) {
   if (type.getShape().size() == 1)
     return type.getElementType();
 
   auto newShape = to_vector(type.getShape());
   newShape.back() = 1;
   return VectorType::get(newShape, type.getElementType());
 }
 
-// Returns a constant of integer of vector type filled with (repeated) `value`.
+/// Returns a constant of integer of vector type filled with (repeated) `value`.
 static Value createScalarOrSplatConstant(ConversionPatternRewriter &rewriter,
                                          Location loc, Type type,
                                          const APInt &value) {
   Attribute attr;
   if (auto intTy = type.dyn_cast<IntegerType>()) {
     attr = rewriter.getIntegerAttr(type, value);
   } else {
     auto vecTy = type.cast<VectorType>();
     attr = SplatElementsAttr::get(vecTy, value);
   }
 
   return rewriter.create<arith::ConstantOp>(loc, attr);
 }
 
-// Returns a constant of integer of vector type filled with (repeated) `value`.
+/// Returns a constant of integer of vector type filled with (repeated) `value`.
 static Value createScalarOrSplatConstant(ConversionPatternRewriter &rewriter,
                                          Location loc, Type type,
                                          int64_t value) {
   unsigned elementBitWidth = 0;
   if (auto intTy = type.dyn_cast<IntegerType>())
     elementBitWidth = intTy.getWidth();
   else
     elementBitWidth = type.cast<VectorType>().getElementTypeBitWidth();
 
   return createScalarOrSplatConstant(rewriter, loc, type,
                                      APInt(elementBitWidth, value));
 }
 
-// Extracts the `input` vector slice with elements at the last dimension offset
-// by `lastOffset`. Returns a value of vector type with the last dimension
-// reduced to x1 or fully scalarized, e.g.:
-//   - vector<3x2xi16> --> vector<3x1xi16>
-//   - vector<2xi16>   --> i16
+/// Extracts the `input` vector slice with elements at the last dimension offset
+/// by `lastOffset`. Returns a value of vector type with the last dimension
+/// reduced to x1 or fully scalarized, e.g.:
+///   - vector<3x2xi16> --> vector<3x1xi16>
+///   - vector<2xi16>   --> i16
 static Value extractLastDimSlice(ConversionPatternRewriter &rewriter,
                                  Location loc, Value input,
                                  int64_t lastOffset) {
   ArrayRef<int64_t> shape = input.getType().cast<VectorType>().getShape();
   assert(lastOffset < shape.back() && "Offset out of bounds");
 
   // Scalarize the result in case of 1D vectors.
   if (shape.size() == 1)
     return rewriter.create<vector::ExtractOp>(loc, input, lastOffset);
 
   SmallVector<int64_t> offsets(shape.size(), 0);
   offsets.back() = lastOffset;
   auto sizes = llvm::to_vector(shape);
   sizes.back() = 1;
   SmallVector<int64_t> strides(shape.size(), 1);
 
   return rewriter.create<vector::ExtractStridedSliceOp>(loc, input, offsets,
                                                         sizes, strides);
 }
 
-// Extracts two vector slices from the `input` whose type is `vector<...x2T>`,
-// with the first element at offset 0 and the second element at offset 1.
+/// Extracts two vector slices from the `input` whose type is `vector<...x2T>`,
+/// with the first element at offset 0 and the second element at offset 1.
 static std::pair<Value, Value>
 extractLastDimHalves(ConversionPatternRewriter &rewriter, Location loc,
                      Value input) {
   return {extractLastDimSlice(rewriter, loc, input, 0),
           extractLastDimSlice(rewriter, loc, input, 1)};
 }
 
 // Performs a vector shape cast to drop the trailing x1 dimension. If the
 // `input` is a scalar, this is a noop.
 static Value dropTrailingX1Dim(ConversionPatternRewriter &rewriter,
                                Location loc, Value input) {
   auto vecTy = input.getType().dyn_cast<VectorType>();
   if (!vecTy)
     return input;
 
   // Shape cast to drop the last x1 dimention.
   ArrayRef<int64_t> shape = vecTy.getShape();
   assert(shape.size() >= 2 && "Expected vector with at list two dims");
   assert(shape.back() == 1 && "Expected the last vector dim to be x1");
 
   auto newVecTy = VectorType::get(shape.drop_back(), vecTy.getElementType());
   return rewriter.create<vector::ShapeCastOp>(loc, newVecTy, input);
 }
 
-// Performs a vector shape cast to append an x1 dimension. If the
-// `input` is a scalar, this is a noop.
+/// Performs a vector shape cast to append an x1 dimension. If the
+/// `input` is a scalar, this is a noop.
 static Value appendX1Dim(ConversionPatternRewriter &rewriter, Location loc,
                          Value input) {
   auto vecTy = input.getType().dyn_cast<VectorType>();
   if (!vecTy)
     return input;
 
   // Add a trailing x1 dim.
   auto newShape = llvm::to_vector(vecTy.getShape());
   newShape.push_back(1);
   auto newTy = VectorType::get(newShape, vecTy.getElementType());
   return rewriter.create<vector::ShapeCastOp>(loc, newTy, input);
 }
 
-// Inserts the `source` vector slice into the `dest` vector at offset
-// `lastOffset` in the last dimension. `source` can be a scalar when `dest` is a
-// 1D vector.
+/// Inserts the `source` vector slice into the `dest` vector at offset
+/// `lastOffset` in the last dimension. `source` can be a scalar when `dest` is
+/// a 1D vector.
 static Value insertLastDimSlice(ConversionPatternRewriter &rewriter,
                                 Location loc, Value source, Value dest,
                                 int64_t lastOffset) {
   ArrayRef<int64_t> shape = dest.getType().cast<VectorType>().getShape();
   assert(lastOffset < shape.back() && "Offset out of bounds");
 
   // Handle scalar source.
   if (source.getType().isa<IntegerType>())
     return rewriter.create<vector::InsertOp>(loc, source, dest, lastOffset);
 
   SmallVector<int64_t> offsets(shape.size(), 0);
   offsets.back() = lastOffset;
   SmallVector<int64_t> strides(shape.size(), 1);
   return rewriter.create<vector::InsertStridedSliceOp>(loc, source, dest,
                                                        offsets, strides);
 }
 
-// Constructs a new vector of type `resultType` by creating a series of
-// insertions of `resultComponents`, each at the next offset of the last vector
-// dimension.
-// When all `resultComponents` are scalars, the result type is `vector<NxT>`;
-// when `resultComponents` are `vector<...x1xT>`s, the result type is
-// `vector<...xNxT>`, where `N` is the number of `resultComponenets`.
+/// Constructs a new vector of type `resultType` by creating a series of
+/// insertions of `resultComponents`, each at the next offset of the last vector
+/// dimension.
+/// When all `resultComponents` are scalars, the result type is `vector<NxT>`;
+/// when `resultComponents` are `vector<...x1xT>`s, the result type is
+/// `vector<...xNxT>`, where `N` is the number of `resultComponenets`.
 static Value constructResultVector(ConversionPatternRewriter &rewriter,
                                    Location loc, VectorType resultType,
                                    ValueRange resultComponents) {
   llvm::ArrayRef<int64_t> resultShape = resultType.getShape();
   (void)resultShape;
   assert(!resultShape.empty() && "Result expected to have dimentions");
   assert(resultShape.back() == static_cast<int64_t>(resultComponents.size()) &&
          "Wrong number of result components");
matchAndRewrite(arith::ExtUIOp op, OpAdaptor adaptor,
     Value zeroCst = createScalarOrSplatConstant(rewriter, loc, newTy, 0);
     Value newRes = insertLastDimSlice(rewriter, loc, extended, zeroCst, 0);
     rewriter.replaceOp(op, newRes);
     return success();
   }
 };
 
 //===----------------------------------------------------------------------===//
+// ConvertShRUI
+//===----------------------------------------------------------------------===//
+
+struct ConvertShRUI final : OpConversionPattern<arith::ShRUIOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::ShRUIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+
+    Type oldTy = op.getType();
+    auto newTy = getTypeConverter()->convertType(oldTy).cast<VectorType>();
+    Type newOperandTy = reduceInnermostDim(newTy);
+    unsigned newBitWidth = newTy.getElementTypeBitWidth();
+
+    auto [lhsElem0, lhsElem1] =
+        extractLastDimHalves(rewriter, loc, adaptor.getLhs());
+    Value rhsElem0 = extractLastDimSlice(rewriter, loc, adaptor.getRhs(), 0);
+
+    // Assume that the shift amount is < 2 * newBitWidth. Calculate the low and
+    // high halves of the results separately:
+    //   1. low := a or b or c, where:
+    //     a) Bits from LHS.low, shifted by the RHS.
+    //     b) Bits from LHS.high, shifted left. These matter when
+    //        RHS < newBitWidth, e.g.:
+    //         [hhhh][0000] shrui 3 --> [000h][hhh0]
+    //                                          ^
+    //                                          |
+    //                                 [hhhh] shli (4 - 1)
+    //     c) Bits from LHS.high, shifted right. These come into play when
+    //        RHS > newBitWidth, e.g.:
+    //         [hhhh][0000] shrui 7 --> [0000][000h]
+    //                                          ^
+    //                                          |
+    //                                 [hhhh] shrui (7 - 4)
+    //
+    //   2. high := LHS.high shrui RHS
+    //
+    // Because shifts by values >= newBitWidth are undefined, we ignore the high
+    // half of RHS, and introduce 'bounds checks' to account for
+    // RHS.low > newBitWidth.
+    //
+    // TODO: Explore possible optimizations.
+    Value zeroCst = createScalarOrSplatConstant(rewriter, loc, newOperandTy, 0);
+    Value elemBitWidth =
+        createScalarOrSplatConstant(rewriter, loc, newOperandTy, newBitWidth);
+
+    Value illegalElemShift = rewriter.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::uge, rhsElem0, elemBitWidth);
+
+    Value shiftedElem0 =
+        rewriter.create<arith::ShRUIOp>(loc, lhsElem0, rhsElem0);
+    Value resElem0Low = rewriter.create<arith::SelectOp>(loc, illegalElemShift,
+                                                         zeroCst, shiftedElem0);
+    Value shiftedElem1 =
+        rewriter.create<arith::ShRUIOp>(loc, lhsElem1, rhsElem0);
+    Value resElem1 = rewriter.create<arith::SelectOp>(loc, illegalElemShift,
+                                                      zeroCst, shiftedElem1);
+
+    Value cappedShiftAmount = rewriter.create<arith::SelectOp>(
+        loc, illegalElemShift, elemBitWidth, rhsElem0);
+    Value leftShiftAmount =
+        rewriter.create<arith::SubIOp>(loc, elemBitWidth, cappedShiftAmount);
+    Value shiftedLeft =
+        rewriter.create<arith::ShLIOp>(loc, lhsElem1, leftShiftAmount);
+    Value overshotShiftAmount =
+        rewriter.create<arith::SubIOp>(loc, rhsElem0, elemBitWidth);
+    Value shiftedRight =
+        rewriter.create<arith::ShRUIOp>(loc, lhsElem1, overshotShiftAmount);
+
+    Value resElem0High = rewriter.create<arith::SelectOp>(
+        loc, illegalElemShift, shiftedRight, shiftedLeft);
+    Value resElem0 =
+        rewriter.create<arith::OrIOp>(loc, resElem0Low, resElem0High);
+
+    Value resultVec =
+        constructResultVector(rewriter, loc, newTy, {resElem0, resElem1});
+    rewriter.replaceOp(op, resultVec);
+    return success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
 // ConvertTruncI
 //===----------------------------------------------------------------------===//
 
 struct ConvertTruncI final : OpConversionPattern<arith::TruncIOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::TruncIOp op, OpAdaptor adaptor,
void arith::populateWideIntEmulationPatterns(
   populateCallOpTypeConversionPattern(patterns, typeConverter);
   populateReturnOpTypeConversionPattern(patterns, typeConverter);
 
   // Populate `arith.*` conversion patterns.
   patterns.add<
       // Misc ops.
       ConvertConstant, ConvertVectorPrint,
       // Binary ops.
-      ConvertAddI, ConvertMulI,
+      ConvertAddI, ConvertMulI, ConvertShRUI,
       // Extension and truncation ops.
       ConvertExtSI, ConvertExtUI, ConvertTruncI>(typeConverter,
                                                  patterns.getContext());
 }

mlir/test/Dialect/Arithmetic/emulate-wide-int.mlir

 
 // CHECK-LABEL: func.func @muli_vector
 // CHECK-SAME:    ({{%.+}}: vector<3x2xi32>, {{%.+}}: vector<3x2xi32>) -> vector<3x2xi32>
 // CHECK:       return {{%.+}} : vector<3x2xi32>
 func.func @muli_vector(%a : vector<3xi64>, %b : vector<3xi64>) -> vector<3xi64> {
     %m = arith.muli %a, %b : vector<3xi64>
     return %m : vector<3xi64>
 }
+
+// CHECK-LABEL: func.func @shrui_scalar
+// CHECK-SAME:    ([[ARG0:%.+]]: vector<2xi32>, [[ARG1:%.+]]: vector<2xi32>) -> vector<2xi32>
+// CHECK-NEXT:    [[LOW0:%.+]]     = vector.extract [[ARG0]][0] : vector<2xi32>
+// CHECK-NEXT:    [[HIGH0:%.+]]    = vector.extract [[ARG0]][1] : vector<2xi32>
+// CHECK-NEXT:    [[LOW1:%.+]]     = vector.extract [[ARG1]][0] : vector<2xi32>
+// CHECK-NEXT:    [[CST0:%.+]]     = arith.constant 0 : i32
+// CHECK-NEXT:    [[CST32:%.+]]    = arith.constant 32 : i32
+// CHECK-DAG:     [[OOB:%.+]]      = arith.cmpi uge, [[LOW1]], [[CST32]] : i32
+// CHECK-DAG:     [[SHLOW0:%.+]]   = arith.shrui [[LOW0]], [[LOW1]] : i32
+// CHECK-NEXT:    [[RES0LOW:%.+]]  = arith.select [[OOB]], [[CST0]], [[SHLOW0]] : i32
+// CHECK-NEXT:    [[SHRHIGH0:%.+]] = arith.shrui [[HIGH0]], [[LOW1]] : i32
+// CHECK-NEXT:    [[RESLOW1:%.+]]  = arith.select [[OOB]], [[CST0]], [[SHRHIGH0]] : i32
+// CHECK-NEXT:    [[SHAMT:%.+]]    = arith.select [[OOB]], [[CST32]], [[LOW1]] : i32
+// CHECK-NEXT:    [[LSHAMT:%.+]]   = arith.subi [[CST32]], [[SHAMT]] : i32
+// CHECK-NEXT:    [[SHLHIGH0:%.+]] = arith.shli [[HIGH0]], [[LSHAMT]] : i32
+// CHECK-NEXT:    [[RSHAMT:%.+]]   = arith.subi [[LOW1]], [[CST32]] : i32
+// CHECK-NEXT:    [[SHRHIGH0:%.+]] = arith.shrui [[HIGH0]], [[RSHAMT]] : i32
+// CHECK-NEXT:    [[RES0HIGH:%.+]] = arith.select [[OOB]], [[SHRHIGH0]], [[SHLHIGH0]] : i32
+// CHECK-NEXT:    [[RES0:%.+]]     = arith.ori [[RES0LOW]], [[RES0HIGH]] : i32
+// CHECK-NEXT:    [[VZ:%.+]]       = arith.constant dense<0> : vector<2xi32>
+// CHECK-NEXT:    [[INS0:%.+]]     = vector.insert [[RES0]], [[VZ]] [0] : i32 into vector<2xi32>
+// CHECK-NEXT:    [[INS1:%.+]]     = vector.insert [[RESLOW1]], [[INS0]] [1] : i32 into vector<2xi32>
+// CHECK-NEXT:    return [[INS1]] : vector<2xi32>
+func.func @shrui_scalar(%a : i64, %b : i64) -> i64 {
+    %c = arith.shrui %a, %b : i64
+    return %c : i64
+}
+
+// CHECK-LABEL: func.func @shrui_scalar_cst_2
+// CHECK-SAME:    ({{%.+}}: vector<2xi32>) -> vector<2xi32>
+// CHECK:       return {{%.+}} : vector<2xi32>
+func.func @shrui_scalar_cst_2(%a : i64) -> i64 {
+    %b = arith.constant 2 : i64
+    %c = arith.shrui %a, %b : i64
+    return %c : i64
+}
+
+// CHECK-LABEL: func.func @shrui_scalar_cst_36
+// CHECK-SAME:    ({{%.+}}: vector<2xi32>) -> vector<2xi32>
+// CHECK:       return {{%.+}} : vector<2xi32>
+func.func @shrui_scalar_cst_36(%a : i64) -> i64 {
+    %b = arith.constant 36 : i64
+    %c = arith.shrui %a, %b : i64
+    return %c : i64
+}
+
+// CHECK-LABEL: func.func @shrui_vector
+// CHECK-SAME:    ({{%.+}}: vector<3x2xi32>, {{%.+}}: vector<3x2xi32>) -> vector<3x2xi32>
+// CHECK:       return {{%.+}} : vector<3x2xi32>
+func.func @shrui_vector(%a : vector<3xi64>, %b : vector<3xi64>) -> vector<3xi64> {
+    %m = arith.shrui %a, %b : vector<3xi64>
+    return %m : vector<3xi64>
+}