[mlir][arith] Add canonicalization patterns for 'mul*i_extended'

[kuhar]

Summary

• Add a fold for mulsi_extended(x, 1)
• Add folds to demote wide integer multiplication to mul*i_extended when the result is shifted and truncated: trunci(shrui(mul(*ext(x), *ext(y)), c)) -> mul*i_extended(x, y)

[kuhar]

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

125

FYI, this is my first time writing non-trivial DDR, I'd appreciate feedback if you see how to simplify these or make them more idiomatic.

[Mogball]

Can you update the description to explain the pattern? It's hard to tell at-a-glance what I'm looking for

[Mogball]

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

125

This is as good as it gets

275

can you make sure these fit under 80 characters?

[kuhar]

Fix formatting in ArithCanonicalization.td. Rebase.

[jpienaar]

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

289

Do you care about matching arith constantop explicitly? Else use constantlikematcher

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

80

I would probably use getElementTypeOrSame followed by this rather than getElementTypeBitWidth.

86

With

(getScalarOrElementWidth($0) - getScalarOrElementWidth($1)) == "
        "getIntOrSplatIntValue($2)

it seems this will happily allow for known and unsupported cases with this chosen sentinel. And one could end up with $0 having width 1, $1 unsupported (so -1) and a shift amount of 2 and the constraint is happy.

101

-1 is a valid integer so using it as a sentinel here seems problematic given generic name of function.

[kuhar]

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

289

These patterns don't but this seems to be the convention used in this file.

[kuhar]

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

86

I think this can't happen in valid IR because of the other type constraints

[kuhar]

Address post-commit comments

[kuhar]

Clean up includes

[jpienaar]

Thanks!

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

```mlir
       %2 = arith.trunci %1 : i5 to i4 // %2 is 0b0101
       %3 = arith.trunci %1 : i5 to i3 // %3 is 0b101
 
       %5 = arith.trunci %0 : vector<2 x i32> to vector<2 x i16>
     ```
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TruncFOp
 //===----------------------------------------------------------------------===//
 
 def Arith_TruncFOp : Arith_FToFCastOp<"truncf"> {

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

 
 // mulsi_extended(x, y) -> [muli(x, y), x], when the `high` result is unused.
 // Since the `high` result it not used, any replacement value will do.
 def MulSIExtendedToMulI :
     Pattern<(Arith_MulSIExtendedOp:$res $x, $y),
         [(Arith_MulIOp $x, $y), (replaceWithValue $x)],
         [(Constraint<CPred<"$0.getUses().empty()">> $res__1)]>;
 
+
+def IsScalarOrSplatOne :
+    Constraint<And<[
+      CPred<"succeeded(getIntOrSplatIntValue($0))">,
+      CPred<"getIntOrSplatIntValue($0).value() == 1">]>>;
+
+// mulsi_extended(x, 1) -> [x, extsi(cmpi slt, x, 0)]
+def MulSIExtendedRHSOne :
+    Pattern<(Arith_MulSIExtendedOp $x, (ConstantLikeMatcher AnyAttr:$c1)),
+            [(replaceWithValue $x),
+             (Arith_ExtSIOp(Arith_CmpIOp
+                              (NativeCodeCall<"arith::CmpIPredicate::slt">),
+                              $x,
+                              (Arith_ConstantOp (GetZeroAttr $x))))],
+            [(IsScalarOrSplatOne $c1)]>;
+
 //===----------------------------------------------------------------------===//
 // MulUIExtendedOp
 //===----------------------------------------------------------------------===//
 
 // mului_extended(x, y) -> [muli(x, y), x], when the `high` result is unused.
 // Since the `high` result it not used, any replacement value will do.
 def MulUIExtendedToMulI :
     Pattern<(Arith_MulUIExtendedOp:$res $x, $y),

 
 // or extsi(x), extsi(y) -> extsi(or(x,y))
 def OrOfExtSI :
     Pat<(Arith_OrIOp (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)),
         (Arith_ExtSIOp (Arith_OrIOp $x, $y)),
         [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
 
 //===----------------------------------------------------------------------===//
+// TruncIOp
+//===----------------------------------------------------------------------===//
+
+def ValuesWithSameType :
+    Constraint<
+      CPred<"llvm::all_equal({$0.getType(), $1.getType(), $2.getType()})">>;
+
+def ValueWiderThan :
+    Constraint<And<[
+      CPred<"getScalarOrElementWidth($0) > getScalarOrElementWidth($1)">,
+      CPred<"getScalarOrElementWidth($1) > 0">]>>;
+
+def TruncationMatchesShiftAmount :
+    Constraint<And<[
+      CPred<"succeeded(getIntOrSplatIntValue($2))">,
+      CPred<"(getScalarOrElementWidth($0) - getScalarOrElementWidth($1)) == "
+              "getIntOrSplatIntValue($2).value()">]>>;
+
+// trunci(shrsi(x, c)) -> trunci(shrui(x, c))
+def TruncIShrSIToTrunciShrUI :
+    Pat<(Arith_TruncIOp:$tr
+          (Arith_ShRSIOp $x, (ConstantLikeMatcher AnyAttr:$c0))),
+        (Arith_TruncIOp (Arith_ShRUIOp $x, (Arith_ConstantOp $c0))),
+        [(TruncationMatchesShiftAmount $x, $tr, $c0)]>;
+
+// trunci(shrui(mul(sext(x), sext(y)), c)) -> mulsi_extended(x, y)
+def TruncIShrUIMulIToMulSIExtended :
+    Pat<(Arith_TruncIOp:$tr (Arith_ShRUIOp
+                              (Arith_MulIOp:$mul
+                                (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)),
+                              (ConstantLikeMatcher AnyAttr:$c0))),
+        (Arith_MulSIExtendedOp:$res__1 $x, $y),
+      [(ValuesWithSameType $tr, $x, $y),
+       (ValueWiderThan $mul, $x),
+       (TruncationMatchesShiftAmount $mul, $x, $c0)]>;
+
+// trunci(shrui(mul(zext(x), zext(y)), c)) -> mului_extended(x, y)
+def TruncIShrUIMulIToMulUIExtended :
+    Pat<(Arith_TruncIOp:$tr (Arith_ShRUIOp
+                              (Arith_MulIOp:$mul
+                                (Arith_ExtUIOp $x), (Arith_ExtUIOp $y)),
+                              (ConstantLikeMatcher AnyAttr:$c0))),
+        (Arith_MulUIExtendedOp:$res__1 $x, $y),
+      [(ValuesWithSameType $tr, $x, $y),
+       (ValueWiderThan $mul, $x),
+       (TruncationMatchesShiftAmount $mul, $x, $c0)]>;
+
+//===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 // mulf(negf(x), negf(y)) -> mulf(x,y)
 // (retain fastmath flags of original mulf)
 def MulFOfNegF :
     Pat<(Arith_MulFOp (Arith_NegFOp $x, $_), (Arith_NegFOp $y, $_), $fmf),
         (Arith_MulFOp $x, $y, $fmf),

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

 //===- ArithOps.cpp - MLIR Arith dialect ops implementation -----===//
 //
 // 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 <cassert>
+#include <cstdint>
 #include <utility>
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Support/LogicalResult.h"
 
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::arith;
 
arith::CmpIPredicate arith::invertPredicate(arith::CmpIPredicate pred) {
   llvm_unreachable("unknown cmpi predicate kind");
 }
 
 static arith::CmpIPredicateAttr invertPredicate(arith::CmpIPredicateAttr pred) {
   return arith::CmpIPredicateAttr::get(pred.getContext(),
                                        invertPredicate(pred.getValue()));
 }
 
+static int64_t getScalarOrElementWidth(Type type) {
+  Type elemTy = getElementTypeOrSelf(type);
+  if (elemTy.isIntOrFloat())
+    return elemTy.getIntOrFloatBitWidth();
+
+  return -1;
+}
+
+static int64_t getScalarOrElementWidth(Value value) {
+  return getScalarOrElementWidth(value.getType());
+}
+
+static FailureOr<APInt> getIntOrSplatIntValue(Attribute attr) {
+  if (auto intAttr = attr.dyn_cast<IntegerAttr>())
+    return intAttr.getValue();
+
+  if (auto splatAttr = attr.dyn_cast<SplatElementsAttr>())
+    if (splatAttr.getElementType().isa<IntegerType>())
+      return splatAttr.getSplatValue<APInt>();
+
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd canonicalization patterns
 //===----------------------------------------------------------------------===//
 
 namespace {
 #include "ArithCanonicalization.inc"
 } // namespace
 
if (Attribute lowAttr = constFoldBinaryOp<IntegerAttr>(
     return success();
   }
 
   return failure();
 }
 
 void arith::MulSIExtendedOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
-  patterns.add<MulSIExtendedToMulI>(context);
+  patterns.add<MulSIExtendedToMulI, MulSIExtendedRHSOne>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // MulUIExtendedOp
 //===----------------------------------------------------------------------===//
 
 Optional<SmallVector<int64_t, 4>> arith::MulUIExtendedOp::getShapeForUnroll() {
   if (auto vt = getType(0).dyn_cast<VectorType>())
return constFoldCastOp<IntegerAttr, IntegerAttr>(
         return a.trunc(bitWidth);
       });
 }
 
 bool arith::TruncIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::less, IntegerType>(inputs, outputs);
 }
 
+void arith::TruncIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                  MLIRContext *context) {
+  patterns.add<TruncIShrSIToTrunciShrUI, TruncIShrUIMulIToMulSIExtended,
+               TruncIShrUIMulIToMulUIExtended>(context);
+}
+
 LogicalResult arith::TruncIOp::verify() {
   return verifyTruncateOp<IntegerType>(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // TruncFOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Arith/canonicalize.mlir

 //  CHECK-NEXT:   %[[zero:.+]] = arith.constant 0 : i32
 //  CHECK-NEXT:   return %[[zero]], %[[zero]]
 func.func @mulsiExtendedZeroLhs(%arg0: i32) -> (i32, i32) {
   %zero = arith.constant 0 : i32
   %low, %high = arith.mulsi_extended %zero, %arg0: i32
   return %low, %high : i32, i32
 }
 
+// CHECK-LABEL: @mulsiExtendedOneRhs
+//  CHECK-SAME:   (%[[ARG:.+]]: i32) -> (i32, i32)
+//  CHECK-NEXT:   %[[C0:.+]]  = arith.constant 0 : i32
+//  CHECK-NEXT:   %[[CMP:.+]] = arith.cmpi slt, %[[ARG]], %[[C0]] : i32
+//  CHECK-NEXT:   %[[EXT:.+]] = arith.extsi %[[CMP]] : i1 to i32
+//  CHECK-NEXT:   return %[[ARG]], %[[EXT]] : i32, i32
+func.func @mulsiExtendedOneRhs(%arg0: i32) -> (i32, i32) {
+  %one = arith.constant 1 : i32
+  %low, %high = arith.mulsi_extended %arg0, %one: i32
+  return %low, %high : i32, i32
+}
+
+// CHECK-LABEL: @mulsiExtendedOneRhsSplat
+//  CHECK-SAME:   (%[[ARG:.+]]: vector<3xi32>) -> (vector<3xi32>, vector<3xi32>)
+//  CHECK-NEXT:   %[[C0:.+]]  = arith.constant dense<0> : vector<3xi32>
+//  CHECK-NEXT:   %[[CMP:.+]] = arith.cmpi slt, %[[ARG]], %[[C0]] : vector<3xi32>
+//  CHECK-NEXT:   %[[EXT:.+]] = arith.extsi %[[CMP]] : vector<3xi1> to vector<3xi32>
+//  CHECK-NEXT:   return %[[ARG]], %[[EXT]] : vector<3xi32>, vector<3xi32>
+func.func @mulsiExtendedOneRhsSplat(%arg0: vector<3xi32>) -> (vector<3xi32>, vector<3xi32>) {
+  %one = arith.constant dense<1> : vector<3xi32>
+  %low, %high = arith.mulsi_extended %arg0, %one: vector<3xi32>
+  return %low, %high : vector<3xi32>, vector<3xi32>
+}
+
 // CHECK-LABEL: @mulsiExtendedUnusedHigh
 //  CHECK-SAME:   (%[[ARG:.+]]: i32) -> i32
 //  CHECK-NEXT:   %[[RES:.+]] = arith.muli %[[ARG]], %[[ARG]] : i32
 //  CHECK-NEXT:   return %[[RES]]
 func.func @mulsiExtendedUnusedHigh(%arg0: i32) -> i32 {
   %low, %high = arith.mulsi_extended %arg0, %arg0: i32
   return %low : i32
 }

 //  CHECK-SAME:   (%[[A:.*]]: i32, %[[B:.*]]: i32)
 //       CHECK:   %[[RES:.*]] = arith.andi %[[A]], %[[B]] : i32
 //       CHECK:   return %[[RES]]
 func.func @andand3(%a : i32, %b : i32) -> i32 {
   %c = arith.andi %a, %b : i32
   %res = arith.andi %c, %b : i32
   return %res : i32
 }
+
+// -----
+
+// CHECK-LABEL: @truncIShrSIToTrunciShrUI
+//  CHECK-SAME:   (%[[A:.+]]: i64)
+//  CHECK-NEXT:   %[[C32:.+]] = arith.constant 32 : i64
+//  CHECK-NEXT:   %[[SHR:.+]] = arith.shrui %[[A]], %[[C32]] : i64
+//  CHECK-NEXT:   %[[TRU:.+]] = arith.trunci %[[SHR]] : i64 to i32
+//  CHECK-NEXT:   return %[[TRU]] : i32
+func.func @truncIShrSIToTrunciShrUI(%a: i64) -> i32 {
+  %c32 = arith.constant 32: i64
+  %sh = arith.shrsi %a, %c32 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @truncIShrSIToTrunciShrUIBadShiftAmt1
+//       CHECK:   arith.shrsi
+func.func @truncIShrSIToTrunciShrUIBadShiftAmt1(%a: i64) -> i32 {
+  %c33 = arith.constant 33: i64
+  %sh = arith.shrsi %a, %c33 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @truncIShrSIToTrunciShrUIBadShiftAmt2
+//  CHECK:        arith.shrsi
+func.func @truncIShrSIToTrunciShrUIBadShiftAmt2(%a: i64) -> i32 {
+  %c31 = arith.constant 31: i64
+  %sh = arith.shrsi %a, %c31 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @wideMulToMulSIExtended
+//  CHECK-SAME:   (%[[A:.+]]: i32, %[[B:.+]]: i32)
+//  CHECK-NEXT:   %[[LOW:.+]], %[[HIGH:.+]] = arith.mulsi_extended %[[A]], %[[B]] : i32
+//  CHECK-NEXT:   return %[[HIGH]] : i32
+func.func @wideMulToMulSIExtended(%a: i32, %b: i32) -> i32 {
+  %x = arith.extsi %a: i32 to i64
+  %y = arith.extsi %b: i32 to i64
+  %m = arith.muli %x, %y: i64
+  %c32 = arith.constant 32: i64
+  %sh = arith.shrui %m, %c32 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @wideMulToMulSIExtendedVector
+//  CHECK-SAME:   (%[[A:.+]]: vector<3xi32>, %[[B:.+]]: vector<3xi32>)
+//  CHECK-NEXT:   %[[LOW:.+]], %[[HIGH:.+]] = arith.mulsi_extended %[[A]], %[[B]] : vector<3xi32>
+//  CHECK-NEXT:   return %[[HIGH]] : vector<3xi32>
+func.func @wideMulToMulSIExtendedVector(%a: vector<3xi32>, %b: vector<3xi32>) -> vector<3xi32> {
+  %x = arith.extsi %a: vector<3xi32> to vector<3xi64>
+  %y = arith.extsi %b: vector<3xi32> to vector<3xi64>
+  %m = arith.muli %x, %y: vector<3xi64>
+  %c32 = arith.constant dense<32>: vector<3xi64>
+  %sh = arith.shrui %m, %c32 : vector<3xi64>
+  %hi = arith.trunci %sh: vector<3xi64> to vector<3xi32>
+  return %hi : vector<3xi32>
+}
+
+// CHECK-LABEL: @wideMulToMulUIExtended
+//  CHECK-SAME:   (%[[A:.+]]: i32, %[[B:.+]]: i32)
+//  CHECK-NEXT:   %[[LOW:.+]], %[[HIGH:.+]] = arith.mului_extended %[[A]], %[[B]] : i32
+//  CHECK-NEXT:   return %[[HIGH]] : i32
+func.func @wideMulToMulUIExtended(%a: i32, %b: i32) -> i32 {
+  %x = arith.extui %a: i32 to i64
+  %y = arith.extui %b: i32 to i64
+  %m = arith.muli %x, %y: i64
+  %c32 = arith.constant 32: i64
+  %sh = arith.shrui %m, %c32 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @wideMulToMulUIExtendedVector
+//  CHECK-SAME:   (%[[A:.+]]: vector<3xi32>, %[[B:.+]]: vector<3xi32>)
+//  CHECK-NEXT:   %[[LOW:.+]], %[[HIGH:.+]] = arith.mului_extended %[[A]], %[[B]] : vector<3xi32>
+//  CHECK-NEXT:   return %[[HIGH]] : vector<3xi32>
+func.func @wideMulToMulUIExtendedVector(%a: vector<3xi32>, %b: vector<3xi32>) -> vector<3xi32> {
+  %x = arith.extui %a: vector<3xi32> to vector<3xi64>
+  %y = arith.extui %b: vector<3xi32> to vector<3xi64>
+  %m = arith.muli %x, %y: vector<3xi64>
+  %c32 = arith.constant dense<32>: vector<3xi64>
+  %sh = arith.shrui %m, %c32 : vector<3xi64>
+  %hi = arith.trunci %sh: vector<3xi64> to vector<3xi32>
+  return %hi : vector<3xi32>
+}
+
+// CHECK-LABEL: @wideMulToMulIExtendedMixedExt
+//       CHECK:   arith.muli
+//       CHECK:   arith.shrui
+//       CHECK:   arith.trunci
+func.func @wideMulToMulIExtendedMixedExt(%a: i32, %b: i32) -> i32 {
+  %x = arith.extsi %a: i32 to i64
+  %y = arith.extui %b: i32 to i64
+  %m = arith.muli %x, %y: i64
+  %c32 = arith.constant 32: i64
+  %sh = arith.shrui %m, %c32 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @wideMulToMulSIExtendedBadExt
+//       CHECK:   arith.muli
+//       CHECK:   arith.shrui
+//       CHECK:   arith.trunci
+func.func @wideMulToMulSIExtendedBadExt(%a: i16, %b: i16) -> i32 {
+  %x = arith.extsi %a: i16 to i64
+  %y = arith.extsi %b: i16 to i64
+  %m = arith.muli %x, %y: i64
+  %c32 = arith.constant 32: i64
+  %sh = arith.shrui %m, %c32 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @wideMulToMulSIExtendedBadShift1
+//       CHECK:   arith.muli
+//       CHECK:   arith.shrui
+//       CHECK:   arith.trunci
+func.func @wideMulToMulSIExtendedBadShift1(%a: i32, %b: i32) -> i32 {
+  %x = arith.extsi %a: i32 to i64
+  %y = arith.extsi %b: i32 to i64
+  %m = arith.muli %x, %y: i64
+  %c33 = arith.constant 33: i64
+  %sh = arith.shrui %m, %c33 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}
+
+// CHECK-LABEL: @wideMulToMulSIExtendedBadShift2
+//       CHECK:   arith.muli
+//       CHECK:   arith.shrui
+//       CHECK:   arith.trunci
+func.func @wideMulToMulSIExtendedBadShift2(%a: i32, %b: i32) -> i32 {
+  %x = arith.extsi %a: i32 to i64
+  %y = arith.extsi %b: i32 to i64
+  %m = arith.muli %x, %y: i64
+  %c31 = arith.constant 31: i64
+  %sh = arith.shrui %m, %c31 : i64
+  %hi = arith.trunci %sh: i64 to i32
+  return %hi : i32
+}