[mlir][sparse] mask reduction update

[aartbik]

Summary

Reduction updates should be masked, just like the load and stores.
Note that alternatively, we could use the fact that masked values are
zero of += updates and mask invariants to get this working but that
would not work for *= updates. Masking the update itself is cleanest.
This change also replaces the constant mask with a broadcast of "true"
since this constant folds much better for various folding patterns.

[aartbik]

better type test reduction end

mlir/lib/Dialect/Linalg/Transforms/Sparsification.cpp

static Value genVectorMask(CodeGen &codegen, PatternRewriter &rewriter,
   // example, "for i = 0, 128, 16"). A constant all-true mask is generated
   // so that all subsequent masked memory operations are immediately folded
   // into unconditional memory operations.
   IntegerAttr loInt, hiInt, stepInt;
   if (matchPattern(lo, m_Constant(&loInt)) &&
       matchPattern(hi, m_Constant(&hiInt)) &&
       matchPattern(step, m_Constant(&stepInt))) {
     if (((hiInt.getInt() - loInt.getInt()) % stepInt.getInt()) == 0)
-      return rewriter.create<vector::ConstantMaskOp>(
-          loc, mtp, rewriter.getI64ArrayAttr(codegen.curVecLength));
+      return rewriter.create<vector::BroadcastOp>(
+          loc, mtp, rewriter.create<ConstantIntOp>(loc, 1, 1));
   }
   // Otherwise, generate a vector mask that avoids overrunning the upperbound
   // during vector execution. Here we rely on subsequent loop optimizations to
   // avoid executing the mask in all iterations, for example, by splitting the
   // loop into an unconditional vector loop and a scalar cleanup loop.
   Value end = rewriter.create<SubIOp>(loc, hi, iv);
   return rewriter.create<vector::CreateMaskOp>(loc, mtp, end);
 }
if (codegen.curVecLength > 1)
     return genVectorLoad(codegen, rewriter, ptr, args);
   return rewriter.create<LoadOp>(loc, ptr, args);
 }
 
 /// Generates a store on a dense tensor.
 static void genTensorStore(Merger &merger, CodeGen &codegen,
                            PatternRewriter &rewriter, linalg::GenericOp op,
                            unsigned tensor, Value rhs) {
+  Location loc = op.getLoc();
   // Test if this is a scalarized reduction.
   unsigned lhs = op.getNumShapedOperands() - 1;
   if (lhs == tensor && codegen.redVal) {
+    if (codegen.curVecLength > 1)
+      rhs = rewriter.create<SelectOp>(loc, codegen.curVecMask, rhs,
+                                      codegen.redVal);
     codegen.redVal = rhs;
     return;
   }
   // Actual store.
   SmallVector<Value, 4> args;
   auto map = op.getIndexingMap(tensor);
   for (unsigned i = 0, m = map.getNumResults(); i < m; ++i) {
     unsigned idx = map.getDimPosition(i);
     args.push_back(codegen.loops[idx]); // universal dense index
   }
-  Location loc = op.getLoc();
   Value ptr = codegen.buffers[tensor];
   if (codegen.curVecLength > 1)
     genVectorStore(codegen, rewriter, rhs, ptr, args);
   else
     rewriter.create<StoreOp>(loc, rhs, ptr, args);
 }
 
 /// Generates a pointer/index load from the sparse storage scheme.

 /// Generates end of a reduction.
 static void genReductionEnd(Merger &merger, CodeGen &codegen,
                             PatternRewriter &rewriter, linalg::GenericOp op) {
   Value red = codegen.redVal;
   if (!red)
     return;
   codegen.redVal = merger.exp(codegen.redExp).val = Value(); // end chain
   unsigned lhs = op.getNumShapedOperands() - 1;
-  if (codegen.curVecLength > 1) {
+  if (red.getType().isa<VectorType>()) {
     // TODO: assumes + reductions for now
     codegen.curVecLength = 1;
     Value ld = genTensorLoad(merger, codegen, rewriter, op, codegen.redExp);
     red = rewriter.create<vector::ReductionOp>(
         op.getLoc(), ld.getType(), rewriter.getStringAttr("add"), red, ld);
   }
   genTensorStore(merger, codegen, rewriter, op, lhs, red);
 }

mlir/test/Dialect/Linalg/sparse_vector.mlir

outs(%argx: tensor<f32>) {
       ^bb(%a: f32, %b: f32, %x: f32):
         %0 = mulf %a, %b : f32
         %1 = addf %x, %0 : f32
         linalg.yield %1 : f32
   } -> tensor<f32>
   return %0 : tensor<f32>
 }
 
+//
+// CHECK-VEC1-LABEL: func @reduction_17
+// CHECK-VEC1-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-VEC1-DAG:   %[[c16:.*]] = constant 16 : index
+// CHECK-VEC1-DAG:   %[[c17:.*]] = constant 17 : index
+// CHECK-VEC1-DAG:   %[[v0:.*]] = constant dense<0.000000e+00> : vector<16xf32>
+// CHECK-VEC1:       %[[red:.*]] = scf.for %[[i:.*]] = %[[c0]] to %[[c17]] step %[[c16]] iter_args(%[[red_in:.*]] = %[[v0]]) -> (vector<16xf32>) {
+// CHECK-VEC1:         %[[sub:.*]] = subi %[[c17]], %[[i]] : index
+// CHECK-VEC1:         %[[mask:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
+// CHECK-VEC1:         %[[la:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<17xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+// CHECK-VEC1:         %[[lb:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<17xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+// CHECK-VEC1:         %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32>
+// CHECK-VEC1:         %[[a:.*]] = addf %[[red_in]], %[[m]] : vector<16xf32>
+// CHECK-VEC1:         %[[s:.*]] = select %[[mask]], %[[a]], %[[red_in]] : vector<16xi1>, vector<16xf32>
+// CHECK-VEC1:         scf.yield %[[s]] : vector<16xf32>
+// CHECK-VEC1:       }
+// CHECK-VEC1:       %{{.*}} = vector.reduction "add", %[[red]], %{{.*}} : vector<16xf32> into f32
+// CHECK-VEC1:       return
+//
+func @reduction_17(%arga: tensor<17xf32>, %argb: tensor<17xf32>, %argx: tensor<f32>) -> tensor<f32> {
+  %0 = linalg.generic #trait_reduction_d
+    ins(%arga, %argb: tensor<17xf32>, tensor<17xf32>)
+    outs(%argx: tensor<f32>) {
+      ^bb(%a: f32, %b: f32, %x: f32):
+        %0 = mulf %a, %b : f32
+        %1 = addf %x, %0 : f32
+        linalg.yield %1 : f32
+  } -> tensor<f32>
+  return %0 : tensor<f32>
+}
+
 #trait_mul_ds = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // a
     affine_map<(i,j) -> (i,j)>,  // b
     affine_map<(i,j) -> (i,j)>   // x (out)
   ],
   sparse = [
     [ "D", "S" ],  // a