[mlir][arith] Add narrowing pattern to commute extension over insertion

[kuhar]

Summary

This enabled more optimization opportunities by moving
zero/sign-extension closer to the use.

[kuhar]

Add missing testcase

[antiagainst]

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

297

Worth to put some comments here to explain what this is for and why we require it > inserted/dest bitwidth. Something like "Get the immediate destination/result type bitwidth for insert. It's the the intermediate bitwidth intercepting op chains involving extension ops; so we need to make sure it's larger than the required bitwidths needed to represent both the inserted value and the destination value."

mlir/test/Dialect/Arith/int-narrowing.mlir

306

This would actually fold with your https://reviews.llvm.org/D149286?

[kuhar]

Add comments

mlir/test/Dialect/Arith/int-narrowing.mlir

306

It won't fold after createOrFold but will disappear after canonicalization.

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

FailureOr<unsigned> bestBitwidth =
         getNarrowestCompatibleBitwidth(bitsRequired);
     if (failed(bestBitwidth))
       return failure();
 
     Type elemTy = getElementTypeOrSelf(origTy);
     if (!isa<IntegerType>(elemTy))
       return failure();
 
-    auto newElemTy = IntegerType::get(origTy.getContext(), bitsRequired);
+    auto newElemTy = IntegerType::get(origTy.getContext(), *bestBitwidth);
     if (newElemTy == elemTy)
       return failure();
 
     if (origTy == elemTy)
       return newElemTy;
 
     if (auto shapedTy = dyn_cast<ShapedType>(origTy))
       if (auto elemTy = dyn_cast<IntegerType>(shapedTy.getElementType()))
FailureOr<unsigned> calculateBitsRequired(Type type) {
   if (auto intTy = dyn_cast<IntegerType>(getElementTypeOrSelf(type)))
     return intTy.getWidth();
 
   return failure();
 }
 
 enum class ExtensionKind { Sign, Zero };
 
+ExtensionKind getExtensionKind(Operation *op) {
+  assert(op);
+  assert((isa<arith::ExtSIOp, arith::ExtUIOp>(op)) && "Not an extension op");
+  return isa<arith::ExtSIOp>(op) ? ExtensionKind::Sign : ExtensionKind::Zero;
+}
+
+/// Returns the integer bitwidth required to represent `value`.
+unsigned calculateBitsRequired(const APInt &value,
+                               ExtensionKind lookThroughExtension) {
+  // For unsigned values, we only need the active bits. As a special case, zero
+  // requires one bit.
+  if (lookThroughExtension == ExtensionKind::Zero)
+    return std::max(value.getActiveBits(), 1u);
+
+  // If a signed value is nonnegative, we need one extra bit for the sign.
+  if (value.isNonNegative())
+    return value.getActiveBits() + 1;
+
+  // For the signed min, we need all the bits.
+  if (value.isMinSignedValue())
+    return value.getBitWidth();
+
+  // For negative values, we need all the non-sign bits and one extra bit for
+  // the sign.
+  return value.getBitWidth() - value.getNumSignBits() + 1;
+}
+
 /// Returns the integer bitwidth required to represent `value`.
 /// Looks through either sign- or zero-extension as specified by
 /// `lookThroughExtension`.
 FailureOr<unsigned> calculateBitsRequired(Value value,
                                           ExtensionKind lookThroughExtension) {
+  // Handle constants.
+  if (TypedAttr attr; matchPattern(value, m_Constant(&attr))) {
+    if (auto intAttr = dyn_cast<IntegerAttr>(attr))
+      return calculateBitsRequired(intAttr.getValue(), lookThroughExtension);
+
+    if (auto elemsAttr = dyn_cast<DenseElementsAttr>(attr)) {
+      if (elemsAttr.getElementType().isIntOrIndex()) {
+        if (elemsAttr.isSplat())
+          return calculateBitsRequired(elemsAttr.getSplatValue<APInt>(),
+                                       lookThroughExtension);
+
+        unsigned maxBits = 1;
+        for (const APInt &elemValue : elemsAttr.getValues<APInt>())
+          maxBits = std::max(
+              maxBits, calculateBitsRequired(elemValue, lookThroughExtension));
+        return maxBits;
+      }
+    }
+  }
+
   if (lookThroughExtension == ExtensionKind::Sign) {
     if (auto sext = value.getDefiningOp<arith::ExtSIOp>())
       return calculateBitsRequired(sext.getIn().getType());
   } else if (lookThroughExtension == ExtensionKind::Zero) {
     if (auto zext = value.getDefiningOp<arith::ExtUIOp>())
       return calculateBitsRequired(zext.getIn().getType());
   }
 

 };
 using SIToFPPattern = IToFPPattern<arith::SIToFPOp, ExtensionKind::Sign>;
 using UIToFPPattern = IToFPPattern<arith::UIToFPOp, ExtensionKind::Zero>;
 
 //===----------------------------------------------------------------------===//
 // Patterns to Commute Extension Ops
 //===----------------------------------------------------------------------===//
 
-struct ExtensionOverExtract final : OpRewritePattern<vector::ExtractOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct ExtensionOverExtract final : NarrowingPattern<vector::ExtractOp> {
+  using NarrowingPattern::NarrowingPattern;
 
   LogicalResult matchAndRewrite(vector::ExtractOp op,
                                 PatternRewriter &rewriter) const override {
     Operation *def = op.getVector().getDefiningOp();
     if (!def)
       return failure();
 
     return TypeSwitch<Operation *, LogicalResult>(def)
         .Case<arith::ExtSIOp, arith::ExtUIOp>([&](auto extOp) {
           Value newExtract = rewriter.create<vector::ExtractOp>(
               op.getLoc(), extOp.getIn(), op.getPosition());
           rewriter.replaceOpWithNewOp<decltype(extOp)>(op, op.getType(),
                                                        newExtract);
           return success();
         })
         .Default(failure());
   }
 };
 
 struct ExtensionOverExtractElement final
-    : OpRewritePattern<vector::ExtractElementOp> {
-  using OpRewritePattern::OpRewritePattern;
+    : NarrowingPattern<vector::ExtractElementOp> {
+  using NarrowingPattern::NarrowingPattern;
 
   LogicalResult matchAndRewrite(vector::ExtractElementOp op,
                                 PatternRewriter &rewriter) const override {
     Operation *def = op.getVector().getDefiningOp();
     if (!def)
       return failure();
 
     return TypeSwitch<Operation *, LogicalResult>(def)
         .Case<arith::ExtSIOp, arith::ExtUIOp>([&](auto extOp) {
           Value newExtract = rewriter.create<vector::ExtractElementOp>(
               op.getLoc(), extOp.getIn(), op.getPosition());
           rewriter.replaceOpWithNewOp<decltype(extOp)>(op, op.getType(),
                                                        newExtract);
           return success();
         })
         .Default(failure());
   }
 };
 
 struct ExtensionOverExtractStridedSlice final
-    : OpRewritePattern<vector::ExtractStridedSliceOp> {
-  using OpRewritePattern::OpRewritePattern;
+    : NarrowingPattern<vector::ExtractStridedSliceOp> {
+  using NarrowingPattern::NarrowingPattern;
 
   LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp op,
                                 PatternRewriter &rewriter) const override {
     Operation *def = op.getVector().getDefiningOp();
     if (!def)
       return failure();
 
     return TypeSwitch<Operation *, LogicalResult>(def)
         .Case<arith::ExtSIOp, arith::ExtUIOp>([&](auto extOp) {
           VectorType origTy = op.getType();
           Type inElemTy =
               cast<VectorType>(extOp.getIn().getType()).getElementType();
           VectorType extractTy = origTy.cloneWith(origTy.getShape(), inElemTy);
           Value newExtract = rewriter.create<vector::ExtractStridedSliceOp>(
               op.getLoc(), extractTy, extOp.getIn(), op.getOffsets(),
               op.getSizes(), op.getStrides());
           rewriter.replaceOpWithNewOp<decltype(extOp)>(op, op.getType(),
                                                        newExtract);
           return success();
         })
         .Default(failure());
   }
 };
 
+struct ExtensionOverInsert final : NarrowingPattern<vector::InsertOp> {
+  using NarrowingPattern::NarrowingPattern;
+
+  LogicalResult matchAndRewrite(vector::InsertOp op,
+                                PatternRewriter &rewriter) const override {
+    Operation *def = op.getSource().getDefiningOp();
+    if (!def)
+      return failure();
+
+    return TypeSwitch<Operation *, LogicalResult>(def)
+        .Case<arith::ExtSIOp, arith::ExtUIOp>([&](auto extOp) {
+          FailureOr<vector::InsertOp> newInsert =
+              createNarrowInsert(op, rewriter, extOp);
+          if (failed(newInsert))
+            return failure();
+          rewriter.replaceOpWithNewOp<decltype(extOp)>(op, op.getType(),
+                                                       *newInsert);
+          return success();
+        })
+        .Default(failure());
+  }
+
+  FailureOr<vector::InsertOp> createNarrowInsert(vector::InsertOp op,
+                                                 PatternRewriter &rewriter,
+                                                 Operation *insValue) const {
+    assert((isa<arith::ExtSIOp, arith::ExtUIOp>(insValue)));
+
+    FailureOr<unsigned> origBitsRequired = calculateBitsRequired(op.getType());
+    if (failed(origBitsRequired))
+      return failure();
+
+    ExtensionKind kind = getExtensionKind(insValue);
+    FailureOr<unsigned> destBitsRequired =
+        calculateBitsRequired(op.getDest(), kind);
+    if (failed(destBitsRequired) || *destBitsRequired >= *origBitsRequired)
+      return failure();
+
+    FailureOr<unsigned> insertedBitsRequired =
+        calculateBitsRequired(insValue->getOperands().front(), kind);
+    if (failed(insertedBitsRequired) ||
+        *insertedBitsRequired >= *origBitsRequired)
+      return failure();
+
+    // Find a narrower element type that satisfies the bitwidth requirements of
+    // both the source and the destination values.
+    unsigned newInsertionBits =
+        std::max(*destBitsRequired, *insertedBitsRequired);
+    FailureOr<Type> newVecTy = getNarrowType(newInsertionBits, op.getType());
+    if (failed(newVecTy) || *newVecTy == op.getType())
+      return failure();
+
+    FailureOr<Type> newInsertedValueTy =
+        getNarrowType(newInsertionBits, insValue->getResultTypes().front());
+    if (failed(newInsertedValueTy))
+      return failure();
+
+    Location loc = op.getLoc();
+    Value narrowValue = rewriter.createOrFold<arith::TruncIOp>(
+        loc, *newInsertedValueTy, insValue->getResult(0));
+    Value narrowDest =
+        rewriter.createOrFold<arith::TruncIOp>(loc, *newVecTy, op.getDest());
+    return rewriter.create<vector::InsertOp>(loc, narrowValue, narrowDest,
+                                             op.getPosition());
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Pass Definitions
 //===----------------------------------------------------------------------===//
 
 struct ArithIntNarrowingPass final
     : impl::ArithIntNarrowingBase<ArithIntNarrowingPass> {
   using ArithIntNarrowingBase::ArithIntNarrowingBase;
 

 // Public API
 //===----------------------------------------------------------------------===//
 
 void populateArithIntNarrowingPatterns(
     RewritePatternSet &patterns, const ArithIntNarrowingOptions &options) {
   // Add commute patterns with a higher benefit. This is to expose more
   // optimization opportunities to narrowing patterns.
   patterns.add<ExtensionOverExtract, ExtensionOverExtractElement,
-               ExtensionOverExtractStridedSlice>(patterns.getContext(),
-                                                 PatternBenefit(2));
+               ExtensionOverExtractStridedSlice, ExtensionOverInsert>(
+      patterns.getContext(), options, PatternBenefit(2));
 
   patterns.add<SIToFPPattern, UIToFPPattern>(patterns.getContext(), options);
 }
 
 } // namespace mlir::arith

mlir/test/Dialect/Arith/int-narrowing.mlir

 // CHECK-NEXT:    %[[RET:.+]]  = arith.extui %[[EXTR]] : vector<1x2xi16> to vector<1x2xi32>
 // CHECK-NEXT:    return %[[RET]] : vector<1x2xi32>
 func.func @extui_over_extract_strided_slice_2d(%a: vector<2x3xi16>) -> vector<1x2xi32> {
   %b = arith.extui %a : vector<2x3xi16> to vector<2x3xi32>
   %c = vector.extract_strided_slice %b
    {offsets = [1, 1], sizes = [1, 2], strides = [1, 1]} : vector<2x3xi32> to vector<1x2xi32>
   return %c : vector<1x2xi32>
 }
+
+// CHECK-LABEL: func.func @extsi_over_insert_3xi16
+// CHECK-SAME:    (%[[ARG0:.+]]: vector<3xi16>, %[[ARG1:.+]]: i16)
+// CHECK-NEXT:    %[[INS:.+]] = vector.insert %[[ARG1]], %[[ARG0]] [1] : i16 into vector<3xi16>
+// CHECK-NEXT:    %[[RET:.+]] = arith.extsi %[[INS]] : vector<3xi16> to vector<3xi32>
+// CHECK-NEXT:    return %[[RET]] : vector<3xi32>
+func.func @extsi_over_insert_3xi16(%a: vector<3xi16>, %b: i16) -> vector<3xi32> {
+  %c = arith.extsi %a : vector<3xi16> to vector<3xi32>
+  %d = arith.extsi %b : i16 to i32
+  %e = vector.insert %d, %c [1] : i32 into vector<3xi32>
+  return %e : vector<3xi32>
+}
+
+// CHECK-LABEL: func.func @extsi_over_insert_3xi16_cst_0
+// CHECK-SAME:    (%[[ARG:.+]]: i16)
+// CHECK-NEXT:    %[[CST:.+]] = arith.constant dense<0> : vector<3xi16>
+// CHECK-NEXT:    %[[INS:.+]] = vector.insert %[[ARG]], %[[CST]] [1] : i16 into vector<3xi16>
+// CHECK-NEXT:    %[[RET:.+]] = arith.extsi %[[INS]] : vector<3xi16> to vector<3xi32>
+// CHECK-NEXT:    return %[[RET]] : vector<3xi32>
+func.func @extsi_over_insert_3xi16_cst_0(%a: i16) -> vector<3xi32> {
+  %cst = arith.constant dense<0> : vector<3xi32>
+  %d = arith.extsi %a : i16 to i32
+  %e = vector.insert %d, %cst [1] : i32 into vector<3xi32>
+  return %e : vector<3xi32>
+}
+
+// CHECK-LABEL: func.func @extsi_over_insert_3xi8_cst
+// CHECK-SAME:    (%[[ARG:.+]]: i8)
+// CHECK-NEXT:    %[[CST:.+]] = arith.constant dense<[-1, 127, -128]> : vector<3xi8>
+// CHECK-NEXT:    %[[INS:.+]] = vector.insert %[[ARG]], %[[CST]] [1] : i8 into vector<3xi8>
+// CHECK-NEXT:    %[[RET:.+]] = arith.extsi %[[INS]] : vector<3xi8> to vector<3xi32>
+// CHECK-NEXT:    return %[[RET]] : vector<3xi32>
+func.func @extsi_over_insert_3xi8_cst(%a: i8) -> vector<3xi32> {
+  %cst = arith.constant dense<[-1, 127, -128]> : vector<3xi32>
+  %d = arith.extsi %a : i8 to i32
+  %e = vector.insert %d, %cst [1] : i32 into vector<3xi32>
+  return %e : vector<3xi32>
+}
+
+// CHECK-LABEL: func.func @extui_over_insert_3xi8_cst
+// CHECK-SAME:    (%[[ARG:.+]]: i8)
+// CHECK-NEXT:    %[[CST:.+]] = arith.constant dense<[1, 127, -1]> : vector<3xi8>
+// CHECK-NEXT:    %[[INS:.+]] = vector.insert %[[ARG]], %[[CST]] [1] : i8 into vector<3xi8>
+// CHECK-NEXT:    %[[RET:.+]] = arith.extui %[[INS]] : vector<3xi8> to vector<3xi32>
+// CHECK-NEXT:    return %[[RET]] : vector<3xi32>
+func.func @extui_over_insert_3xi8_cst(%a: i8) -> vector<3xi32> {
+  %cst = arith.constant dense<[1, 127, 255]> : vector<3xi32>
+  %d = arith.extui %a : i8 to i32
+  %e = vector.insert %d, %cst [1] : i32 into vector<3xi32>
+  return %e : vector<3xi32>
+}
+
+// CHECK-LABEL: func.func @extsi_over_insert_3xi16_cst_i16
+// CHECK-SAME:    (%[[ARG:.+]]: i8)
+// CHECK-NEXT:    %[[CST:.+]]  = arith.constant dense<[-1, 128, 0]> : vector<3xi16>
+// CHECK-NEXT:    %[[SRCE:.+]] = arith.extsi %[[ARG]] : i8 to i32
+// CHECK-NEXT:    %[[SRCT:.+]] = arith.trunci %[[SRCE]] : i32 to i16
+// CHECK-NEXT:    %[[INS:.+]]  = vector.insert %[[SRCT]], %[[CST]] [1] : i16 into vector<3xi16>
+// CHECK-NEXT:    %[[RET:.+]]  = arith.extsi %[[INS]] : vector<3xi16> to vector<3xi32>
+// CHECK-NEXT:    return %[[RET]] : vector<3xi32>
+func.func @extsi_over_insert_3xi16_cst_i16(%a: i8) -> vector<3xi32> {
+  %cst = arith.constant dense<[-1, 128, 0]> : vector<3xi32>
+  %d = arith.extsi %a : i8 to i32
+  %e = vector.insert %d, %cst [1] : i32 into vector<3xi32>
+  return %e : vector<3xi32>
+}
+
+// CHECK-LABEL: func.func @extui_over_insert_3xi16_cst_i16
+// CHECK-SAME:    (%[[ARG:.+]]: i8)
+// CHECK-NEXT:    %[[CST:.+]]  = arith.constant dense<[1, 256, 0]> : vector<3xi16>
+// CHECK-NEXT:    %[[SRCE:.+]] = arith.extui %[[ARG]] : i8 to i32
+// CHECK-NEXT:    %[[SRCT:.+]] = arith.trunci %[[SRCE]] : i32 to i16
+// CHECK-NEXT:    %[[INS:.+]]  = vector.insert %[[SRCT]], %[[CST]] [1] : i16 into vector<3xi16>
+// CHECK-NEXT:    %[[RET:.+]]  = arith.extui %[[INS]] : vector<3xi16> to vector<3xi32>
+// CHECK-NEXT:    return %[[RET]] : vector<3xi32>
+func.func @extui_over_insert_3xi16_cst_i16(%a: i8) -> vector<3xi32> {
+  %cst = arith.constant dense<[1, 256, 0]> : vector<3xi32>
+  %d = arith.extui %a : i8 to i32
+  %e = vector.insert %d, %cst [1] : i32 into vector<3xi32>
+  return %e : vector<3xi32>
+}