[sparse][mlir][vectorization] add support for shift-by-invariant

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp

 
 /// Helper to test for given index value.
 static bool isIntValue(Value val, int64_t idx) {
   if (auto ival = getConstantIntValue(val))
     return *ival == idx;
   return false;
 }
 
+/// Helper test for invariant value (defined outside given block).
+static bool isInvariantValue(Value val, Block *block) {
+  return val.getDefiningOp() && val.getDefiningOp()->getBlock() != block;
+}
+
+/// Helper test for invariant argument (defined outside given block).
+static bool isInvariantArg(BlockArgument arg, Block *block) {
+  return arg.getOwner() != block;
+}
+
 /// Constructs vector type for element type.
 static VectorType vectorType(VL vl, Type etp) {
   unsigned numScalableDims = vl.enableVLAVectorization;
   return VectorType::get(vl.vectorLength, etp, numScalableDims);
 }
 
 /// Constructs vector type from pointer.
 static VectorType vectorType(VL vl, Value ptr) {

 /// penalty. Indices that already are 64-bit, in theory, cannot express the
 /// full range since the LLVM backend defines addressing in terms of an
 /// unsigned pointer/signed index pair.
 static bool vectorizeSubscripts(PatternRewriter &rewriter, scf::ForOp forOp,
                                 VL vl, ValueRange subs, bool codegen,
                                 Value vmask, SmallVectorImpl<Value> &idxs) {
   unsigned d = 0;
   unsigned dim = subs.size();
+  Block *block = &forOp.getRegion().front();
   for (auto sub : subs) {
     bool innermost = ++d == dim;
     // Invariant subscripts in outer dimensions simply pass through.
     // Note that we rely on LICM to hoist loads where all subscripts
     // are invariant in the innermost loop.
-    if (sub.getDefiningOp() &&
-        sub.getDefiningOp()->getBlock() != &forOp.getRegion().front()) {
+    // Example:
+    //   a[inv][i] for inv
+    if (isInvariantValue(sub, block)) {
       if (innermost)
         return false;
       if (codegen)
         idxs.push_back(sub);
       continue; // success so far
     }
     // Invariant block arguments (including outer loop indices) in outer
     // dimensions simply pass through. Direct loop indices in the
     // innermost loop simply pass through as well.
-    if (auto barg = sub.dyn_cast<BlockArgument>()) {
-      bool invariant = barg.getOwner() != &forOp.getRegion().front();
-      if (invariant == innermost)
+    // Example:
+    //   a[i][j] for both i and j
+    if (auto arg = sub.dyn_cast<BlockArgument>()) {
+      if (isInvariantArg(arg, block) == innermost)
         return false;
       if (codegen)
         idxs.push_back(sub);
       continue; // success so far
     }
     // Look under the hood of casting.
     auto cast = sub;
     while (true) {
for (auto sub : subs) {
     // and 16-bit values, an 32-bit index width suffices. For 32-bit values,
     // zero extending the elements into 64-bit loses some performance since
     // the 32-bit indexed gather/scatter is more efficient than the 64-bit
     // index variant (if the negative 32-bit index space is unused, the
     // enableSIMDIndex32 flag can preserve this performance). For 64-bit
     // values, there is no good way to state that the indices are unsigned,
     // which creates the potential of incorrect address calculations in the
     // unlikely case we need such extremely large offsets.
+    // Example:
+    //    a[ ind[i] ]
     if (auto load = cast.getDefiningOp<memref::LoadOp>()) {
       if (!innermost)
         return false;
       if (codegen) {
         SmallVector<Value> idxs2(load.getIndices()); // no need to analyze
         Location loc = forOp.getLoc();
         Value vload =
             genVectorLoad(rewriter, loc, vl, load.getMemRef(), idxs2, vmask);
         Type etp = vload.getType().cast<VectorType>().getElementType();
         if (!etp.isa<IndexType>()) {
           if (etp.getIntOrFloatBitWidth() < 32)
             vload = rewriter.create<arith::ExtUIOp>(
                 loc, vectorType(vl, rewriter.getI32Type()), vload);
           else if (etp.getIntOrFloatBitWidth() < 64 && !vl.enableSIMDIndex32)
             vload = rewriter.create<arith::ExtUIOp>(
                 loc, vectorType(vl, rewriter.getI64Type()), vload);
         }
         idxs.push_back(vload);
       }
       continue; // success so far
     }
     // Address calculation 'i = add inv, idx' (after LICM).
+    // Example:
+    //    a[base + i]
     if (auto load = cast.getDefiningOp<arith::AddIOp>()) {
       Value inv = load.getOperand(0);
       Value idx = load.getOperand(1);
-      if (inv.getDefiningOp() &&
-          inv.getDefiningOp()->getBlock() != &forOp.getRegion().front() &&
-          idx.dyn_cast<BlockArgument>()) {
-        if (!innermost)
+      if (isInvariantValue(inv, block)) {
+        if (auto arg = idx.dyn_cast<BlockArgument>()) {
+          if (isInvariantArg(arg, block) || !innermost)
             return false;
           if (codegen)
             idxs.push_back(
                 rewriter.create<arith::AddIOp>(forOp.getLoc(), inv, idx));
           continue; // success so far
         }
       }
+    }
     return false;
   }
   return true;
 }
 
 #define UNAOP(xxx)                                                             \
   if (isa<xxx>(def)) {                                                         \
     if (codegen)                                                               \
if (arg == forOp.getInductionVar()) {
       // invariant value, and rely on later replacing and folding to
       // construct a proper reduction chain for the latter case.
       if (codegen)
         vexp = genVectorInvariantValue(rewriter, vl, exp);
       return true;
   }
   // Something defined outside the loop-body is invariant.
   Operation *def = exp.getDefiningOp();
-  if (def->getBlock() != &forOp.getRegion().front()) {
+  Block *block = &forOp.getRegion().front();
+  if (def->getBlock() != block) {
     if (codegen)
       vexp = genVectorInvariantValue(rewriter, vl, exp);
     return true;
   }
   // Proper load operations. These are either values involved in the
   // actual computation, such as a[i] = b[i] becomes a[lo:hi] = b[lo:hi],
   // or index values inside the computation that are now fetched from
   // the sparse storage index arrays, such as a[i] = i becomes
if (vectorizeExpr(rewriter, forOp, vl, def->getOperand(0), codegen, vmask,
       // TODO: complex?
     }
   } else if (def->getNumOperands() == 2) {
     Value vx, vy;
     if (vectorizeExpr(rewriter, forOp, vl, def->getOperand(0), codegen, vmask,
                       vx) &&
         vectorizeExpr(rewriter, forOp, vl, def->getOperand(1), codegen, vmask,
                       vy)) {
+      // We only accept shift-by-invariant (where the same shift factor applies
+      // to all packed elements). In the vector dialect, this is still
+      // represented with an expanded vector at the right-hand-side, however,
+      // so that we do not have to special case the code generation.
+      if (isa<arith::ShLIOp>(def) || isa<arith::ShRUIOp>(def) ||
+          isa<arith::ShRSIOp>(def)) {
+        Value shiftFactor = def->getOperand(1);
+        if (!isInvariantValue(shiftFactor, block))
+          return false;
+      }
+      // Generate code.
       BINOP(arith::MulFOp)
       BINOP(arith::MulIOp)
       BINOP(arith::DivFOp)
       BINOP(arith::DivSIOp)
       BINOP(arith::DivUIOp)
       BINOP(arith::AddFOp)
       BINOP(arith::AddIOp)
       BINOP(arith::SubFOp)
       BINOP(arith::SubIOp)
       BINOP(arith::AndIOp)
       BINOP(arith::OrIOp)
       BINOP(arith::XOrIOp)
+      BINOP(arith::ShLIOp)
+      BINOP(arith::ShRUIOp)
+      BINOP(arith::ShRSIOp)
       // TODO: complex?
-      // TODO: shift by invariant?
     }
   }
   return false;
 }
 
 #undef UNAOP
 #undef TYPEDUNAOP
 #undef BINOP

mlir/test/Dialect/SparseTensor/sparse_vector_ops.mlir

#trait = {
   iterator_types = ["parallel"],
   doc = "x(i) = a(i) ops b(i)"
 }
 
 // CHECK-LABEL: func.func @vops
 // CHECK-DAG:       %[[C1:.*]] = arith.constant dense<2.000000e+00> : vector<8xf32>
 // CHECK-DAG:       %[[C2:.*]] = arith.constant dense<1.000000e+00> : vector<8xf32>
 // CHECK-DAG:       %[[C3:.*]] = arith.constant dense<255> : vector<8xi64>
+// CHECK-DAG:       %[[C4:.*]] = arith.constant dense<4> : vector<8xi32>
+// CHECK-DAG:       %[[C5:.*]] = arith.constant dense<1> : vector<8xi32>
 // CHECK:           scf.for
 // CHECK:             %[[VAL_14:.*]] = vector.load
 // CHECK:             %[[VAL_15:.*]] = math.absf %[[VAL_14]] : vector<8xf32>
 // CHECK:             %[[VAL_16:.*]] = math.ceil %[[VAL_15]] : vector<8xf32>
 // CHECK:             %[[VAL_17:.*]] = math.floor %[[VAL_16]] : vector<8xf32>
 // CHECK:             %[[VAL_18:.*]] = math.sqrt %[[VAL_17]] : vector<8xf32>
 // CHECK:             %[[VAL_19:.*]] = math.expm1 %[[VAL_18]] : vector<8xf32>
 // CHECK:             %[[VAL_20:.*]] = math.sin %[[VAL_19]] : vector<8xf32>
 // CHECK:             %[[VAL_21:.*]] = math.tanh %[[VAL_20]] : vector<8xf32>
 // CHECK:             %[[VAL_22:.*]] = arith.negf %[[VAL_21]] : vector<8xf32>
 // CHECK:             %[[VAL_23:.*]] = vector.load
 // CHECK:             %[[VAL_24:.*]] = arith.mulf %[[VAL_22]], %[[VAL_23]] : vector<8xf32>
 // CHECK:             %[[VAL_25:.*]] = arith.divf %[[VAL_24]], %[[C1]] : vector<8xf32>
 // CHECK:             %[[VAL_26:.*]] = arith.addf %[[VAL_25]], %[[C1]] : vector<8xf32>
 // CHECK:             %[[VAL_27:.*]] = arith.subf %[[VAL_26]], %[[C2]] : vector<8xf32>
 // CHECK:             %[[VAL_28:.*]] = arith.extf %[[VAL_27]] : vector<8xf32> to vector<8xf64>
 // CHECK:             %[[VAL_29:.*]] = arith.bitcast %[[VAL_28]] : vector<8xf64> to vector<8xi64>
 // CHECK:             %[[VAL_30:.*]] = arith.addi %[[VAL_29]], %[[VAL_29]] : vector<8xi64>
 // CHECK:             %[[VAL_31:.*]] = arith.andi %[[VAL_30]], %[[C3]] : vector<8xi64>
 // CHECK:             %[[VAL_32:.*]] = arith.trunci %[[VAL_31]] : vector<8xi64> to vector<8xi16>
 // CHECK:             %[[VAL_33:.*]] = arith.extsi %[[VAL_32]] : vector<8xi16> to vector<8xi32>
-// CHECK:             %[[VAL_34:.*]] = arith.uitofp %[[VAL_33]] : vector<8xi32> to vector<8xf32>
-// CHECK:             vector.store %[[VAL_34]]
+// CHECK:             %[[VAL_34:.*]] = arith.shrsi %[[VAL_33]], %[[C4]] : vector<8xi32>
+// CHECK:             %[[VAL_35:.*]] = arith.shrui %[[VAL_34]], %[[C4]] : vector<8xi32>
+// CHECK:             %[[VAL_36:.*]] = arith.shli %[[VAL_35]], %[[C5]] : vector<8xi32>
+// CHECK:             %[[VAL_37:.*]] = arith.uitofp %[[VAL_36]] : vector<8xi32> to vector<8xf32>
+// CHECK:             vector.store %[[VAL_37]]
 // CHECK:           }
 func.func @vops(%arga: tensor<1024xf32, #DenseVector>,
                 %argb: tensor<1024xf32, #DenseVector>) -> tensor<1024xf32> {
   %init = bufferization.alloc_tensor() : tensor<1024xf32>
   %o = arith.constant 1.0 : f32
   %c = arith.constant 2.0 : f32
   %i = arith.constant 255 : i64
+  %s = arith.constant 4 : i32
+  %t = arith.constant 1 : i32
   %0 = linalg.generic #trait
     ins(%arga, %argb: tensor<1024xf32, #DenseVector>, tensor<1024xf32, #DenseVector>)
     outs(%init: tensor<1024xf32>) {
       ^bb(%a: f32, %b: f32, %x: f32):
         %0 = math.absf %a : f32
         %1 = math.ceil %0 : f32
         %2 = math.floor %1 : f32
         %3 = math.sqrt %2 : f32
         %4 = math.expm1 %3 : f32
         %5 = math.sin %4 : f32
         %6 = math.tanh %5 : f32
         %7 = arith.negf %6 : f32
         %8 = arith.mulf %7, %b : f32
         %9 = arith.divf %8, %c : f32
         %10 = arith.addf %9, %c : f32
         %11 = arith.subf %10, %o : f32
         %12 = arith.extf %11 : f32 to f64
         %13 = arith.bitcast %12 : f64 to i64
         %14 = arith.addi %13, %13 : i64
         %15 = arith.andi %14, %i : i64
         %16 = arith.trunci %15 : i64 to i16
         %17 = arith.extsi %16 : i16 to i32
-        %18 = arith.uitofp %17 : i32 to f32
-        linalg.yield %18 : f32
+  %18 = arith.shrsi %17, %s : i32
+  %19 = arith.shrui %18, %s : i32
+  %20 = arith.shli %19, %t : i32
+        %21 = arith.uitofp %20 : i32 to f32
+        linalg.yield %21 : f32
   } -> tensor<1024xf32>
   return %0 : tensor<1024xf32>
 }