Add integer bit-shift operations to the standard dialect.

Rename the 'shlis' operation in the standard dialect to 'shift_left'. Add tests
for this operation (these have been missing so far) and add a lowering to the
'shl' operation in the LLVM dialect.

Add also 'shift_right_signed' (lowered to LLVM's 'ashr') and 'shift_right_unsigned'
(lowered to 'lshr').

The original plan was to name these operations 'shift.left', 'shift.right.signed'
and 'shift.right.unsigned'. This works if the operations are prefixed with 'std.'
in MLIR assembly. Unfortunately during import the short form is ambigous with
operations from a hypothetical 'shift' dialect. The best solution seems to omit
dots in standard operations for now.

Closes tensorflow/mlir#226

PiperOrigin-RevId: 286803388

7 files changed, 51 insertions, 32 deletions

diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
index 9208ce8ab6d..3f613c6bfb5 100644
--- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
+++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
@@ -94,7 +94,7 @@ public:
     auto rhs = visit(expr.getRHS());
     assert(lhs && rhs && "unexpected affine expr lowering failure");
 
-    Value *remainder = builder.create<RemISOp>(loc, lhs, rhs);
+    Value *remainder = builder.create<SignedRemIOp>(loc, lhs, rhs);
     Value *zeroCst = builder.create<ConstantIndexOp>(loc, 0);
     Value *isRemainderNegative =
         builder.create<CmpIOp>(loc, CmpIPredicate::slt, remainder, zeroCst);
@@ -138,7 +138,7 @@ public:
     Value *negatedDecremented = builder.create<SubIOp>(loc, noneCst, lhs);
     Value *dividend =
         builder.create<SelectOp>(loc, negative, negatedDecremented, lhs);
-    Value *quotient = builder.create<DivISOp>(loc, dividend, rhs);
+    Value *quotient = builder.create<SignedDivIOp>(loc, dividend, rhs);
     Value *correctedQuotient = builder.create<SubIOp>(loc, noneCst, quotient);
     Value *result =
         builder.create<SelectOp>(loc, negative, correctedQuotient, quotient);
@@ -178,7 +178,7 @@ public:
     Value *decremented = builder.create<SubIOp>(loc, lhs, oneCst);
     Value *dividend =
         builder.create<SelectOp>(loc, nonPositive, negated, decremented);
-    Value *quotient = builder.create<DivISOp>(loc, dividend, rhs);
+    Value *quotient = builder.create<SignedDivIOp>(loc, dividend, rhs);
     Value *negatedQuotient = builder.create<SubIOp>(loc, zeroCst, quotient);
     Value *incrementedQuotient = builder.create<AddIOp>(loc, quotient, oneCst);
     Value *result = builder.create<SelectOp>(loc, nonPositive, negatedQuotient,
diff --git a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp
index c269dc5c45a..d663ae105f2 100644
--- a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp
+++ b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp
@@ -254,7 +254,7 @@ Optional<OpTy> LoopToGpuConverter::collectBounds(OpTy forOp,
         builder.create<SubIOp>(currentLoop.getLoc(), upperBound, lowerBound);
     Value *step = getOrCreateStep(currentLoop, builder);
     if (!isConstantOne(step))
-      range = builder.create<DivISOp>(currentLoop.getLoc(), range, step);
+      range = builder.create<SignedDivIOp>(currentLoop.getLoc(), range, step);
     dims.push_back(range);
 
     lbs.push_back(lowerBound);
diff --git a/mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp b/mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
index ea8501b9a7e..fdc90851b64 100644
--- a/mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
+++ b/mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
@@ -814,16 +814,20 @@ struct SubIOpLowering : public BinaryOpLLVMOpLowering<SubIOp, LLVM::SubOp> {
 struct MulIOpLowering : public BinaryOpLLVMOpLowering<MulIOp, LLVM::MulOp> {
   using Super::Super;
 };
-struct DivISOpLowering : public BinaryOpLLVMOpLowering<DivISOp, LLVM::SDivOp> {
+struct SignedDivIOpLowering
+    : public BinaryOpLLVMOpLowering<SignedDivIOp, LLVM::SDivOp> {
   using Super::Super;
 };
-struct DivIUOpLowering : public BinaryOpLLVMOpLowering<DivIUOp, LLVM::UDivOp> {
+struct UnsignedDivIOpLowering
+    : public BinaryOpLLVMOpLowering<UnsignedDivIOp, LLVM::UDivOp> {
   using Super::Super;
 };
-struct RemISOpLowering : public BinaryOpLLVMOpLowering<RemISOp, LLVM::SRemOp> {
+struct SignedRemIOpLowering
+    : public BinaryOpLLVMOpLowering<SignedRemIOp, LLVM::SRemOp> {
   using Super::Super;
 };
-struct RemIUOpLowering : public BinaryOpLLVMOpLowering<RemIUOp, LLVM::URemOp> {
+struct UnsignedRemIOpLowering
+    : public BinaryOpLLVMOpLowering<UnsignedRemIOp, LLVM::URemOp> {
   using Super::Super;
 };
 struct AndOpLowering : public BinaryOpLLVMOpLowering<AndOp, LLVM::AndOp> {
@@ -862,6 +866,18 @@ struct ConstLLVMOpLowering
     : public OneToOneLLVMOpLowering<ConstantOp, LLVM::ConstantOp> {
   using Super::Super;
 };
+struct ShiftLeftOpLowering
+    : public OneToOneLLVMOpLowering<ShiftLeftOp, LLVM::ShlOp> {
+  using Super::Super;
+};
+struct SignedShiftRightOpLowering
+    : public OneToOneLLVMOpLowering<SignedShiftRightOp, LLVM::AShrOp> {
+  using Super::Super;
+};
+struct UnsignedShiftRightOpLowering
+    : public OneToOneLLVMOpLowering<UnsignedShiftRightOp, LLVM::LShrOp> {
+  using Super::Super;
+};
 
 // Check if the MemRefType `type` is supported by the lowering. We currently
 // only support memrefs with identity maps.
@@ -2082,8 +2098,6 @@ void mlir::populateStdToLLVMNonMemoryConversionPatterns(
       CosOpLowering,
       ConstLLVMOpLowering,
       DivFOpLowering,
-      DivISOpLowering,
-      DivIUOpLowering,
       ExpOpLowering,
       LogOpLowering,
       Log10OpLowering,
@@ -2097,18 +2111,23 @@ void mlir::populateStdToLLVMNonMemoryConversionPatterns(
       OrOpLowering,
       PrefetchOpLowering,
       RemFOpLowering,
-      RemISOpLowering,
-      RemIUOpLowering,
       ReturnOpLowering,
       SIToFPLowering,
       SelectOpLowering,
+      ShiftLeftOpLowering,
       SignExtendIOpLowering,
+      SignedDivIOpLowering,
+      SignedRemIOpLowering,
+      SignedShiftRightOpLowering,
       SplatOpLowering,
       SplatNdOpLowering,
       SubFOpLowering,
       SubIOpLowering,
       TanhOpLowering,
       TruncateIOpLowering,
+      UnsignedDivIOpLowering,
+      UnsignedRemIOpLowering,
+      UnsignedShiftRightOpLowering,
       XOrOpLowering,
       ZeroExtendIOpLowering>(*converter.getDialect(), converter);
   // clang-format on
diff --git a/mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp b/mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp
index e87bd4ef861..a14271efbb6 100644
--- a/mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp
+++ b/mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp
@@ -316,8 +316,8 @@ void populateStandardToSPIRVPatterns(MLIRContext *context,
   patterns.insert<ConstantIndexOpConversion, CmpIOpConversion,
                   IntegerOpConversion<AddIOp, spirv::IAddOp>,
                   IntegerOpConversion<MulIOp, spirv::IMulOp>,
-                  IntegerOpConversion<DivISOp, spirv::SDivOp>,
-                  IntegerOpConversion<RemISOp, spirv::SModOp>,
+                  IntegerOpConversion<SignedDivIOp, spirv::SDivOp>,
+                  IntegerOpConversion<SignedRemIOp, spirv::SModOp>,
                   IntegerOpConversion<SubIOp, spirv::ISubOp>, LoadOpConversion,
                   ReturnOpConversion, SelectOpConversion, StoreOpConversion>(
       context, typeConverter);
diff --git a/mlir/lib/Dialect/StandardOps/Ops.cpp b/mlir/lib/Dialect/StandardOps/Ops.cpp
index d0fd1855f96..4116f6f14ae 100644
--- a/mlir/lib/Dialect/StandardOps/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/Ops.cpp
@@ -1320,10 +1320,10 @@ OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
 }
 
 //===----------------------------------------------------------------------===//
-// DivISOp
+// SignedDivIOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult DivISOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult SignedDivIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary operation takes two operands");
 
   // Don't fold if it would overflow or if it requires a division by zero.
@@ -1339,10 +1339,10 @@ OpFoldResult DivISOp::fold(ArrayRef<Attribute> operands) {
 }
 
 //===----------------------------------------------------------------------===//
-// DivIUOp
+// UnsignedDivIOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult DivIUOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult UnsignedDivIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary operation takes two operands");
 
   // Don't fold if it would require a division by zero.
@@ -1885,11 +1885,11 @@ OpFoldResult RankOp::fold(ArrayRef<Attribute> operands) {
 }
 
 //===----------------------------------------------------------------------===//
-// RemISOp
+// SignedRemIOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult RemISOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "remis takes two operands");
+OpFoldResult SignedRemIOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "remi_signed takes two operands");
 
   auto rhs = operands.back().dyn_cast_or_null<IntegerAttr>();
   if (!rhs)
@@ -1911,11 +1911,11 @@ OpFoldResult RemISOp::fold(ArrayRef<Attribute> operands) {
 }
 
 //===----------------------------------------------------------------------===//
-// RemIUOp
+// UnsignedRemIOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult RemIUOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "remiu takes two operands");
+OpFoldResult UnsignedRemIOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "remi_unsigned takes two operands");
 
   auto rhs = operands.back().dyn_cast_or_null<IntegerAttr>();
   if (!rhs)
diff --git a/mlir/lib/EDSC/Builders.cpp b/mlir/lib/EDSC/Builders.cpp
index 2956066a035..47e2dfed55e 100644
--- a/mlir/lib/EDSC/Builders.cpp
+++ b/mlir/lib/EDSC/Builders.cpp
@@ -390,14 +390,14 @@ ValueHandle mlir::edsc::op::operator*(ValueHandle lhs, ValueHandle rhs) {
 }
 
 ValueHandle mlir::edsc::op::operator/(ValueHandle lhs, ValueHandle rhs) {
-  return createBinaryHandle<DivISOp, DivFOp>(
+  return createBinaryHandle<SignedDivIOp, DivFOp>(
       lhs, rhs, [](AffineExpr d0, AffineExpr d1) -> AffineExpr {
         llvm_unreachable("only exprs of non-index type support operator/");
       });
 }
 
 ValueHandle mlir::edsc::op::operator%(ValueHandle lhs, ValueHandle rhs) {
-  return createBinaryHandle<RemISOp, RemFOp>(
+  return createBinaryHandle<SignedRemIOp, RemFOp>(
       lhs, rhs, [](AffineExpr d0, AffineExpr d1) { return d0 % d1; });
 }
 
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 419df8d2705..3691aee4870 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -850,7 +850,7 @@ static Value *ceilDivPositive(OpBuilder &builder, Location loc, Value *dividend,
   Value *divisorMinusOneCst = builder.create<ConstantIndexOp>(loc, divisor - 1);
   Value *divisorCst = builder.create<ConstantIndexOp>(loc, divisor);
   Value *sum = builder.create<AddIOp>(loc, dividend, divisorMinusOneCst);
-  return builder.create<DivISOp>(loc, sum, divisorCst);
+  return builder.create<SignedDivIOp>(loc, sum, divisorCst);
 }
 
 // Build the IR that performs ceil division of a positive value by another
@@ -864,7 +864,7 @@ static Value *ceilDivPositive(OpBuilder &builder, Location loc, Value *dividend,
   Value *cstOne = builder.create<ConstantIndexOp>(loc, 1);
   Value *divisorMinusOne = builder.create<SubIOp>(loc, divisor, cstOne);
   Value *sum = builder.create<AddIOp>(loc, dividend, divisorMinusOne);
-  return builder.create<DivISOp>(loc, sum, divisor);
+  return builder.create<SignedDivIOp>(loc, sum, divisor);
 }
 
 // Hoist the ops within `outer` that appear before `inner`.
@@ -1084,12 +1084,12 @@ void mlir::coalesceLoops(MutableArrayRef<loop::ForOp> loops) {
   for (unsigned i = 0, e = loops.size(); i < e; ++i) {
     unsigned idx = loops.size() - i - 1;
     if (i != 0)
-      previous =
-          builder.create<DivISOp>(loc, previous, loops[idx + 1].upperBound());
+      previous = builder.create<SignedDivIOp>(loc, previous,
+                                              loops[idx + 1].upperBound());
 
     Value *iv = (i == e - 1) ? previous
-                             : builder.create<RemISOp>(loc, previous,
-                                                       loops[idx].upperBound());
+                             : builder.create<SignedRemIOp>(
+                                   loc, previous, loops[idx].upperBound());
     replaceAllUsesInRegionWith(loops[idx].getInductionVar(), iv,
                                loops.back().region());
   }