[mlir][sparse] add shift ops support

[aartbik]

Summary

Arbitrary shifts have some complications, but shift by invariants
(viz. tensor index exp only at left hand side) can be easily
handled with the conjunctive rule.

[gussmith23]

This looks good -- I have some minor comments, do with them what you will!

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

211–214

Are c and inv in these comments referencing the same thing? An invariant? Mostly just wondering if we should use the same symbol for both.

238

Is this tested anywhere? I can't help but think the new way of storing the op symbols may lead to a bug someday (if you forget to add a symbol for a new op, for example).

[aartbik]

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

211–214

For divs, it is really a nonzero const, and for shifts, any invariant. Hence, this difference. I will add some asserts in a follow up CL to make sure we keep that property.

238

Agreed. I will do a follow up CL that adds a unit test + some asserts on this.

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

enum Kind {
   kDivU, // unsigned
   kAddF,
   kAddI,
   kSubF,
   kSubI,
   kAndI,
   kOrI,
   kXorI,
+  kShrS, // signed
+  kShrU, // unsigned
+  kShlI,
 };
 
 /// Children subexpressions of tensor operations.
 struct Children {
   unsigned e0;
   unsigned e1;
 };
 
#endif
   /// Returns index of the root expression on success.
   Optional<unsigned> buildTensorExpFromLinalg(linalg::GenericOp op);
 
   /// Rebuilds SSA format from a tensor expression.
   Value buildExp(PatternRewriter &rewriter, Location loc, unsigned e, Value v0,
                  Value v1);
 
 private:
-  bool maybeZero(unsigned e);
+  bool maybeZero(unsigned e) const;
+  bool isInvariant(unsigned e) const;
 
   /// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
   Optional<unsigned> buildTensorExp(linalg::GenericOp op, Value v);
 
   const unsigned outTensor;
   const unsigned syntheticTensor;
   const unsigned numTensors;
   const unsigned numLoops;

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

bool Merger::isConjunction(unsigned t, unsigned e) const {
   switch (tensorExps[e].kind) {
   case Kind::kTensor:
     return tensorExps[e].tensor == t;
   case kAbsF:
   case kCeilF:
   case kFloorF:
   case kNegF:
   case kNegI:
+  case Kind::kDivF: // note: x / c only
+  case Kind::kDivS:
+  case Kind::kDivU:
+  case Kind::kShrS: // note: x >> inv only
+  case Kind::kShrU:
+  case Kind::kShlI:
     return isConjunction(t, tensorExps[e].children.e0);
   case Kind::kMulF:
   case Kind::kMulI:
   case Kind::kAndI:
-  case Kind::kDivF: // note: x / c only
-  case Kind::kDivS:
-  case Kind::kDivU:
     return isConjunction(t, tensorExps[e].children.e0) ||
            isConjunction(t, tensorExps[e].children.e1);
   default:
     return false;
   }
 }
 
 #ifndef NDEBUG
 
 //
 // Print methods (for debugging).
 //
 
-static const char *kOpSymbols[] = {"",  "",  "abs", "ceil", "floor", "-",
-                                   "-", "*", "*",   "/",    "/",     "+",
-                                   "+", "-", "-",   "&",    "|",     "^"};
+static const char *kOpSymbols[] = {
+    "",  "",  "abs", "ceil", "floor", "-", "-", "*",   "*",  "/", "/",
+    "+", "+", "-",   "-",    "&",     "|", "^", "a>>", ">>", "<<"};
 
 void Merger::dumpExp(unsigned e) const {
   switch (tensorExps[e].kind) {
   case Kind::kTensor:
     if (tensorExps[e].tensor == syntheticTensor)
       llvm::dbgs() << "synthetic_";
     else if (tensorExps[e].tensor == outTensor)
       llvm::dbgs() << "output_";
     llvm::dbgs() << "tensor_" << tensorExps[e].tensor;
     break;
case Kind::kXorI:
     //
     //  x+y|!y | y |    x-y|!y | y |
     //  ---+---+---+    ---+---+---+
     //  !x | 0 | y |    !x | 0 |-y |
     //   x | x |x+y|     x | x |x-y|
     return takeDisj(kind, // take binary disjunction
                     buildLattices(tensorExps[e].children.e0, i),
                     buildLattices(tensorExps[e].children.e1, i));
+  case Kind::kShrS:
+  case Kind::kShrU:
+  case Kind::kShlI:
+    // A shift operation by an invariant amount (viz. tensor expressions
+    // can only occur at the left-hand-side of the operator) can be handled
+    // with the conjuction rule.
+    return takeConj(kind, // take binary conjunction
+                    buildLattices(tensorExps[e].children.e0, i),
+                    buildLattices(tensorExps[e].children.e1, i));
   }
   llvm_unreachable("unexpected expression kind");
 }
 
 Optional<unsigned> Merger::buildTensorExpFromLinalg(linalg::GenericOp op) {
   Operation *yield = op.region().front().getTerminator();
   return buildTensorExp(op, yield->getOperand(0));
 }
 
-bool Merger::maybeZero(unsigned e) {
+bool Merger::maybeZero(unsigned e) const {
   if (tensorExps[e].kind == Kind::kInvariant) {
     if (auto c = tensorExps[e].val.getDefiningOp<ConstantIntOp>())
       return c.getValue() == 0;
     if (auto c = tensorExps[e].val.getDefiningOp<ConstantFloatOp>())
       return c.getValue().isZero();
   }
   return true;
 }
 
+bool Merger::isInvariant(unsigned e) const {
+  return tensorExps[e].kind == Kind::kInvariant;
+}
+
 Optional<unsigned> Merger::buildTensorExp(linalg::GenericOp op, Value v) {
   if (auto arg = v.dyn_cast<BlockArgument>()) {
     unsigned argN = arg.getArgNumber();
     // Any argument of the generic op that is not marked as a scalar
     // argument is considered a tensor, indexed by the implicit loop
     // bounds. This includes rank-0 tensor arguments.
     if (arg.getOwner()->getParentOp() == op) {
       OpOperand *t = op.getInputAndOutputOperands()[argN];
if (x.hasValue() && y.hasValue()) {
       if (isa<SubIOp>(def))
         return addExp(Kind::kSubI, e0, e1);
       if (isa<AndOp>(def))
         return addExp(Kind::kAndI, e0, e1);
       if (isa<OrOp>(def))
         return addExp(Kind::kOrI, e0, e1);
       if (isa<XOrOp>(def))
         return addExp(Kind::kXorI, e0, e1);
+      if (isa<SignedShiftRightOp>(def) && isInvariant(e1))
+        return addExp(Kind::kShrS, e0, e1);
+      if (isa<UnsignedShiftRightOp>(def) && isInvariant(e1))
+        return addExp(Kind::kShrU, e0, e1);
+      if (isa<ShiftLeftOp>(def) && isInvariant(e1))
+        return addExp(Kind::kShlI, e0, e1);
     }
   }
   // Cannot build.
   return None;
 }
 
 Value Merger::buildExp(PatternRewriter &rewriter, Location loc, unsigned e,
                        Value v0, Value v1) {
Value Merger::buildExp(PatternRewriter &rewriter, Location loc, unsigned e,
   case Kind::kSubI:
     return rewriter.create<SubIOp>(loc, v0, v1);
   case Kind::kAndI:
     return rewriter.create<AndOp>(loc, v0, v1);
   case Kind::kOrI:
     return rewriter.create<OrOp>(loc, v0, v1);
   case Kind::kXorI:
     return rewriter.create<XOrOp>(loc, v0, v1);
+  case Kind::kShrS:
+    return rewriter.create<SignedShiftRightOp>(loc, v0, v1);
+  case Kind::kShrU:
+    return rewriter.create<UnsignedShiftRightOp>(loc, v0, v1);
+  case Kind::kShlI:
+    return rewriter.create<ShiftLeftOp>(loc, v0, v1);
   }
   llvm_unreachable("unexpected expression kind in build");
 }
 
 } // namespace sparse_tensor
 } // namespace mlir

mlir/test/Dialect/SparseTensor/sparse_int_ops.mlir

%0 = linalg.generic #trait2
      ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
     outs(%argx: tensor<32xi64>) {
       ^bb(%a: i64, %b: i64, %x: i64):
         %0 = xor %a, %b : i64
         linalg.yield %0 : i64
   } -> tensor<32xi64>
   return %0 : tensor<32xi64>
 }
+
+// CHECK-LABEL:   func @ashrbyc(
+// CHECK-SAME:                  %[[VAL_0:.*]]: tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>>,
+// CHECK-SAME:                  %[[VAL_1:.*]]: tensor<32xi64> {linalg.inplaceable = true}) -> tensor<32xi64> {
+// CHECK:           %[[VAL_2:.*]] = constant 2 : i64
+// CHECK:           %[[VAL_3:.*]] = constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = constant 1 : index
+// CHECK:           %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64>
+// CHECK:           %[[VAL_8:.*]] = memref.buffer_cast %[[VAL_1]] : memref<32xi64>
+// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
+// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
+// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
+// CHECK:             %[[VAL_14:.*]] = shift_right_signed %[[VAL_13]], %[[VAL_2]] : i64
+// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
+// CHECK:           }
+// CHECK:           %[[VAL_15:.*]] = memref.tensor_load %[[VAL_8]] : memref<32xi64>
+// CHECK:           return %[[VAL_15]] : tensor<32xi64>
+// CHECK:         }
+func @ashrbyc(%arga: tensor<32xi64, #SV>,
+              %argx: tensor<32xi64> {linalg.inplaceable = true}) -> tensor<32xi64> {
+  %c = constant 2 : i64
+  %0 = linalg.generic #traitc
+     ins(%arga: tensor<32xi64, #SV>)
+    outs(%argx: tensor<32xi64>) {
+      ^bb(%a: i64, %x: i64):
+        %0 = shift_right_signed %a, %c : i64
+        linalg.yield %0 : i64
+  } -> tensor<32xi64>
+  return %0 : tensor<32xi64>
+}
+
+// CHECK-LABEL:   func @lsrbyc(
+// CHECK-SAME:                 %[[VAL_0:.*]]: tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>>,
+// CHECK-SAME:                 %[[VAL_1:.*]]: tensor<32xi64> {linalg.inplaceable = true}) -> tensor<32xi64> {
+// CHECK:           %[[VAL_2:.*]] = constant 2 : i64
+// CHECK:           %[[VAL_3:.*]] = constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = constant 1 : index
+// CHECK:           %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64>
+// CHECK:           %[[VAL_8:.*]] = memref.buffer_cast %[[VAL_1]] : memref<32xi64>
+// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
+// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
+// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
+// CHECK:             %[[VAL_14:.*]] = shift_right_unsigned %[[VAL_13]], %[[VAL_2]] : i64
+// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
+// CHECK:           }
+// CHECK:           %[[VAL_15:.*]] = memref.tensor_load %[[VAL_8]] : memref<32xi64>
+// CHECK:           return %[[VAL_15]] : tensor<32xi64>
+// CHECK:         }
+func @lsrbyc(%arga: tensor<32xi64, #SV>,
+             %argx: tensor<32xi64> {linalg.inplaceable = true}) -> tensor<32xi64> {
+  %c = constant 2 : i64
+  %0 = linalg.generic #traitc
+     ins(%arga: tensor<32xi64, #SV>)
+    outs(%argx: tensor<32xi64>) {
+      ^bb(%a: i64, %x: i64):
+        %0 = shift_right_unsigned %a, %c : i64
+        linalg.yield %0 : i64
+  } -> tensor<32xi64>
+  return %0 : tensor<32xi64>
+}
+
+// CHECK-LABEL:   func @lslbyc(
+// CHECK-SAME:                 %[[VAL_0:.*]]: tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>>,
+// CHECK-SAME:                 %[[VAL_1:.*]]: tensor<32xi64> {linalg.inplaceable = true}) -> tensor<32xi64> {
+// CHECK:           %[[VAL_2:.*]] = constant 2 : i64
+// CHECK:           %[[VAL_3:.*]] = constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = constant 1 : index
+// CHECK:           %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64>
+// CHECK:           %[[VAL_8:.*]] = memref.buffer_cast %[[VAL_1]] : memref<32xi64>
+// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
+// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
+// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
+// CHECK:             %[[VAL_14:.*]] = shift_left %[[VAL_13]], %[[VAL_2]] : i64
+// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
+// CHECK:           }
+// CHECK:           %[[VAL_15:.*]] = memref.tensor_load %[[VAL_8]] : memref<32xi64>
+// CHECK:           return %[[VAL_15]] : tensor<32xi64>
+// CHECK:         }
+func @lslbyc(%arga: tensor<32xi64, #SV>,
+             %argx: tensor<32xi64> {linalg.inplaceable = true}) -> tensor<32xi64> {
+  %c = constant 2 : i64
+  %0 = linalg.generic #traitc
+     ins(%arga: tensor<32xi64, #SV>)
+    outs(%argx: tensor<32xi64>) {
+      ^bb(%a: i64, %x: i64):
+        %0 = shift_left %a, %c : i64
+        linalg.yield %0 : i64
+  } -> tensor<32xi64>
+  return %0 : tensor<32xi64>
+}
+