[mlir][index] Add shl, shrs, and shru ops

[Mogball]

Summary

This patch adds the left shift, signed right shift, and unsigned right
shift operations to the index dialects with folders and LLVM lowerings.

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

let description = [{
     ```mlir
     // c = max(a, b)
     %c = index.maxu %a, %b
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// ShlOp
+//===----------------------------------------------------------------------===//
+
+def Index_ShlOp : IndexBinaryOp<"shl"> {
+  let summary = "index shift left";
+  let description = [{
+    The `index.shl` operation shifts an index value to the left by a variable
+    amount. The low order bits are filled with zeroes. The RHS operand is always
+    treated as unsigned. If the RHS operand is equal to or greater than the
+    index bitwidth, the operation is undefined.
+
+    Example:
+
+    ```mlir
+    // c = a << b
+    %c = index.shl %a, %b
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// ShrSOp
+//===----------------------------------------------------------------------===//
+
+def Index_ShrSOp : IndexBinaryOp<"shrs"> {
+  let summary = "signed index shift right";
+  let description = [{
+    The `index.shrs` operation shifts an index value to the right by a variable
+    amount. The LHS operand is treated as signed. The high order bits are filled
+    with copies of the most significant bit. If the RHS operand is equal to or
+    greater than the index bitwidth, the operation is undefined.
+
+    Example:
+
+    ```mlir
+    // c = a >> b
+    %c = index.shrs %a, %b
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// ShrUOp
+//===----------------------------------------------------------------------===//
+
+def Index_ShrUOp : IndexBinaryOp<"shru"> {
+  let summary = "unsigned index shift right";
+  let description = [{
+    The `index.shru` operation shifts an index value to the right by a variable
+    amount. The LHS operand is treated as unsigned. The high order bits are
+    filled with zeroes. If the RHS operand is equal to or greater than the index
+    bitwidth, the operation is undefined.
+
+    Example:
+
+    ```mlir
+    // c = a >> b
+    %c = index.shru %a, %b
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // CastSOp
 //===----------------------------------------------------------------------===//
 
 def Index_CastSOp : IndexOp<"casts",
                            [DeclareOpInterfaceMethods<CastOpInterface>]> {
   let summary = "index signed cast";
   let description = [{
     The `index.casts` operation enables conversions between values of index type

mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp

 using ConvertIndexRemS =
     mlir::OneToOneConvertToLLVMPattern<RemSOp, LLVM::SRemOp>;
 using ConvertIndexRemU =
     mlir::OneToOneConvertToLLVMPattern<RemUOp, LLVM::URemOp>;
 using ConvertIndexMaxS =
     mlir::OneToOneConvertToLLVMPattern<MaxSOp, LLVM::SMaxOp>;
 using ConvertIndexMaxU =
     mlir::OneToOneConvertToLLVMPattern<MaxUOp, LLVM::UMaxOp>;
+using ConvertIndexShl = mlir::OneToOneConvertToLLVMPattern<ShlOp, LLVM::ShlOp>;
+using ConvertIndexShrS =
+    mlir::OneToOneConvertToLLVMPattern<ShrSOp, LLVM::AShrOp>;
+using ConvertIndexShrU =
+    mlir::OneToOneConvertToLLVMPattern<ShrUOp, LLVM::LShrOp>;
 using ConvertIndexBoolConstant =
     mlir::OneToOneConvertToLLVMPattern<BoolConstantOp, LLVM::ConstantOp>;
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pattern Population
 //===----------------------------------------------------------------------===//
 
 void index::populateIndexToLLVMConversionPatterns(
     LLVMTypeConverter &typeConverter, RewritePatternSet &patterns) {
   patterns.insert<
       // clang-format off
       ConvertIndexAdd,
       ConvertIndexSub,
       ConvertIndexMul,
       ConvertIndexDivS,
       ConvertIndexDivU,
       ConvertIndexRemS,
       ConvertIndexRemU,
       ConvertIndexMaxS,
       ConvertIndexMaxU,
+      ConvertIndexShl,
+      ConvertIndexShrS,
+      ConvertIndexShrU,
       ConvertIndexCeilDivS,
       ConvertIndexCeilDivU,
       ConvertIndexFloorDivS,
       ConvertIndexCastS,
       ConvertIndexCastU,
       ConvertIndexCmp,
       ConvertIndexSizeOf,
       ConvertIndexConstant,

mlir/lib/Dialect/Index/IR/IndexOps.cpp

 /// ```
 /// trunc(f(a, b)) = f(trunc(a), trunc(b))
 /// ```
 ///
 /// For all values of `a` and `b`. The function accepts a lambda that computes
 /// the integer result, which in turn must satisfy the above property.
 static OpFoldResult foldBinaryOpUnchecked(
     ArrayRef<Attribute> operands,
-    function_ref<APInt(const APInt &, const APInt &)> calculate) {
+    function_ref<Optional<APInt>(const APInt &, const APInt &)> calculate) {
   assert(operands.size() == 2 && "binary operation expected 2 operands");
   auto lhs = dyn_cast_if_present<IntegerAttr>(operands[0]);
   auto rhs = dyn_cast_if_present<IntegerAttr>(operands[1]);
   if (!lhs || !rhs)
     return {};
 
-  APInt result = calculate(lhs.getValue(), rhs.getValue());
-  assert(result.trunc(32) ==
+  Optional<APInt> result = calculate(lhs.getValue(), rhs.getValue());
+  if (!result)
+    return {};
+  assert(result->trunc(32) ==
          calculate(lhs.getValue().trunc(32), rhs.getValue().trunc(32)));
-  return IntegerAttr::get(IndexType::get(lhs.getContext()), std::move(result));
+  return IntegerAttr::get(IndexType::get(lhs.getContext()), std::move(*result));
 }
 
 /// Fold an index operation only if the truncated 64-bit result matches the
 /// 32-bit result for operations that don't satisfy the above property. These
 /// are operations where the upper bits of the operands can affect the lower
 /// bits of the results.
 ///
 /// The function accepts a lambda that computes the integer result in both

 
 OpFoldResult MaxUOp::fold(ArrayRef<Attribute> operands) {
   return foldBinaryOpChecked(operands, [](const APInt &lhs, const APInt &rhs) {
     return lhs.ugt(rhs) ? lhs : rhs;
   });
 }
 
 //===----------------------------------------------------------------------===//
+// ShlOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ShlOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpUnchecked(
+      operands, [](const APInt &lhs, const APInt &rhs) -> Optional<APInt> {
+        // We cannot fold if the RHS is greater than or equal to 32 because
+        // this would be UB in 32-bit systems but not on 64-bit systems. RHS is
+        // already treated as unsigned.
+        if (rhs.uge(32))
+          return {};
+        return lhs << rhs;
+      });
+}
+
+//===----------------------------------------------------------------------===//
+// ShrSOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ShrSOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(
+      operands, [](const APInt &lhs, const APInt &rhs) -> Optional<APInt> {
+        // Don't fold if RHS is greater than or equal to 32.
+        if (rhs.uge(32))
+          return {};
+        return lhs.ashr(rhs);
+      });
+}
+
+//===----------------------------------------------------------------------===//
+// ShrUOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ShrUOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(
+      operands, [](const APInt &lhs, const APInt &rhs) -> Optional<APInt> {
+        // Don't fold if RHS is greater than or equal to 32.
+        if (rhs.uge(32))
+          return {};
+        return lhs.lshr(rhs);
+      });
+}
+
+//===----------------------------------------------------------------------===//
 // CastSOp
 //===----------------------------------------------------------------------===//
 
 bool CastSOp::areCastCompatible(TypeRange lhsTypes, TypeRange rhsTypes) {
   return lhsTypes.front().isa<IndexType>() != rhsTypes.front().isa<IndexType>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir

func.func @trivial_ops(%a: index, %b: index) {
   // CHECK: llvm.srem
   %5 = index.rems %a, %b
   // CHECK: llvm.urem
   %6 = index.remu %a, %b
   // CHECK: llvm.intr.smax
   %7 = index.maxs %a, %b
   // CHECK: llvm.intr.umax
   %8 = index.maxu %a, %b
+  // CHECK: llvm.shl
+  %9 = index.shl %a, %b
+  // CHECK: llvm.ashr
+  %10 = index.shrs %a, %b
+  // CHECK: llvm.lshr
+  %11 = index.shru %a, %b
   // CHECK: llvm.mlir.constant(true
-  %9 = index.bool.constant true
+  %12 = index.bool.constant true
   return
 }
 
 // CHECK-LABEL: @ceildivs
 // CHECK-SAME: %[[NI:.*]]: index, %[[MI:.*]]: index
 func.func @ceildivs(%n: index, %m: index) -> index {
   // CHECK: %[[N:.*]] = builtin.unrealized_conversion_cast %[[NI]]
   // CHECK: %[[M:.*]] = builtin.unrealized_conversion_cast %[[MI]]

mlir/test/Dialect/Index/index-canonicalize.mlir

func.func @maxu() -> index {
   %lhs = index.constant -1
   %rhs = index.constant 1
   // CHECK: %[[A:.*]] = index.constant -1
   %0 = index.maxu %lhs, %rhs
   // CHECK: return %[[A]]
   return %0 : index
 }
 
+// CHECK-LABEL: @shl
+func.func @shl() -> index {
+  %lhs = index.constant 128
+  %rhs = index.constant 2
+  // CHECK: %[[A:.*]] = index.constant 512
+  %0 = index.shl %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
+// CHECK-LABEL: @shl_32
+func.func @shl_32() -> index {
+  %lhs = index.constant 1
+  %rhs = index.constant 32
+  // CHECK: index.shl
+  %0 = index.shl %lhs, %rhs
+  return %0 : index
+}
+
+// CHECK-LABEL: @shl_edge
+func.func @shl_edge() -> index {
+  %lhs = index.constant 4000000000
+  %rhs = index.constant 31
+  // CHECK: %[[A:.*]] = index.constant 858{{[0-9]+}}
+  %0 = index.shl %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
+// CHECK-LABEL: @shrs
+func.func @shrs() -> index {
+  %lhs = index.constant 128
+  %rhs = index.constant 2
+  // CHECK: %[[A:.*]] = index.constant 32
+  %0 = index.shrs %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
+// CHECK-LABEL: @shrs_32
+func.func @shrs_32() -> index {
+  %lhs = index.constant 4000000000000
+  %rhs = index.constant 32
+  // CHECK: index.shrs
+  %0 = index.shrs %lhs, %rhs
+  return %0 : index
+}
+
+// CHECK-LABEL: @shrs_nofold
+func.func @shrs_nofold() -> index {
+  %lhs = index.constant 0x100000000
+  %rhs = index.constant 1
+  // CHECK: index.shrs
+  %0 = index.shrs %lhs, %rhs
+  return %0 : index
+}
+
+// CHECK-LABEL: @shrs_edge
+func.func @shrs_edge() -> index {
+  %lhs = index.constant 0x10000000000
+  %rhs = index.constant 3
+  // CHECK: %[[A:.*]] = index.constant 137{{[0-9]+}}
+  %0 = index.shrs %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
+// CHECK-LABEL: @shru
+func.func @shru() -> index {
+  %lhs = index.constant 128
+  %rhs = index.constant 2
+  // CHECK: %[[A:.*]] = index.constant 32
+  %0 = index.shru %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
+// CHECK-LABEL: @shru_32
+func.func @shru_32() -> index {
+  %lhs = index.constant 4000000000000
+  %rhs = index.constant 32
+  // CHECK: index.shru
+  %0 = index.shru %lhs, %rhs
+  return %0 : index
+}
+
+// CHECK-LABEL: @shru_nofold
+func.func @shru_nofold() -> index {
+  %lhs = index.constant 0x100000000
+  %rhs = index.constant 1
+  // CHECK: index.shru
+  %0 = index.shru %lhs, %rhs
+  return %0 : index
+}
+
+// CHECK-LABEL: @shru_edge
+func.func @shru_edge() -> index {
+  %lhs = index.constant 0x10000000000
+  %rhs = index.constant 3
+  // CHECK: %[[A:.*]] = index.constant 137{{[0-9]+}}
+  %0 = index.shru %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
 // CHECK-LABEL: @cmp
 func.func @cmp() -> (i1, i1, i1, i1) {
   %a = index.constant 0
   %b = index.constant -1
   %c = index.constant -2
   %d = index.constant 4
 
   %0 = index.cmp slt(%a, %b)

mlir/test/Dialect/Index/index-ops.mlir

func.func @binary_ops(%a: index, %b: index) {
   // CHECK-NEXT: index.rems %[[A]], %[[B]]
   %8 = index.rems %a, %b
   // CHECK-NEXT: index.remu %[[A]], %[[B]]
   %9 = index.remu %a, %b
   // CHECK-NEXT: index.maxs %[[A]], %[[B]]
   %10 = index.maxs %a, %b
   // CHECK-NEXT: index.maxu %[[A]], %[[B]]
   %11 = index.maxu %a, %b
+  // CHECK-NEXT: index.shl %[[A]], %[[B]]
+  %12 = index.shl %a, %b
+  // CHECK-NEXT: index.shrs %[[A]], %[[B]]
+  %13 = index.shrs %a, %b
+  // CHECK-NEXT: index.shru %[[A]], %[[B]]
+  %14 = index.shru %a, %b
   return
 }
 
 // CHECK-LABEL: @cmp_op
 // CHECK-SAME: %[[A:.*]]: index, %[[B:.*]]: index
 func.func @cmp_op(%a: index, %b: index) {
   // CHECK-NEXT: index.cmp eq(%[[A]], %[[B]])
   %0 = index.cmp eq(%a, %b)