[mlir][sparse] Moved a conditional from the RT library to the generated MLIR.

[wrengr]

Summary

When generating code to add an element to SparseTensorCOO (e.g., when doing dense=>sparse conversion), we used to check for nonzero values on the runtime side, whereas now we generate MLIR code to do that check.

[wrengr]

A bunch of integration tests are currently failing. I'll look into those and post an update tomorrow (or later this week).

[aartbik]

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

182

this seems the right decision right, so no TODO required?

188

note that there is a rewriter.getZeroAttr(vtp) utility that works for both FP and INT (doing the same under water). Perhaps you can simply start with

Value zero = rewriter.getZeroAttr(vtp);
if (fp)
if(int)
etc. to make the CMP builders fit on one line for readability?

189

Note that LLVM style (and thus clang-tidy) prefers

if (x)

return...

if (y)

return ...

<anything else>

instead of the else-s

193

so no else, but just unreachable

[aartbik]

As for the tests, I have several CHECK tests that verify the IR is as we expect for the conversion, e.g.

func @sparse_convert_1d(%arg0: tensor<?xi32>) -> tensor<?xi32, #SparseVector> {

%0 = sparse_tensor.convert %arg0 : tensor<?xi32> to tensor<?xi32, #SparseVector>
return %0 : tensor<?xi32, #SparseVector>

}

In this case, you will need to add the "if val != 0" to the pattern in the CHECKs.

[wrengr]

simplified getting zeroAttr(). And resolving clang-tidy.

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

182

I was just wanting confirmation, so sgtm.

188

Nice! I thought something like getZeroAttr ought to exist, but I just couldn't seem to find it.

[wrengr]

Marking the newly generated comparison ops as legal.

[aartbik]

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

1 ↗

(On Diff #374085)

Ah, good catch. Can you also fix the spares -> sparse while you are on this line ;-)

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

554–555

please rebase to main. you will most likely get conflicts on this file (due to the ABI fixes that just went in), but it will be extremely easy to resolve since you just need to find the two lines where you deleted this return ;-)

[wrengr]

typo fix

[wrengr]

rebase

[aartbik]

In D110121#3016086, @wrengr wrote:

rebase

Is this rebased with a "pulled" mainline? I don't seen the memref ABI fixes in this revision yet?
No big deal, but better to rebase now then have failed merge due to conflicts later ;-)

[wrengr]

Pretty sure I pulled before rebasing, but I guess it didn't take for some reason.

[wrengr]

re-rebased

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

static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
   params.push_back(ptr);
   // Generate the call to create new tensor.
   StringRef name = "newSparseTensor";
   auto call =
       rewriter.create<CallOp>(loc, pTp, getFunc(op, name, pTp, params), params);
   return call.getResult(0);
 }
 
+/// Generates the comparison `v != 0` where `v` is of numeric type `t`.
+/// For floating types, we use the "unordered" comparator (i.e., returns
+/// true if `v` is NaN).
+// TODO: should we use the ordered comparator instead?
+static Value genIsNonzero(ConversionPatternRewriter &rewriter, Location loc,
+                          Type t, Value v) {
+  if (t.isa<FloatType>())
+    return rewriter.create<CmpFOp>(
+        loc, CmpFPredicate::UNE, v,
+        rewriter.create<ConstantOp>(loc, rewriter.getFloatAttr(t, 0)));
+  else if (t.isIntOrIndex())
+    return rewriter.create<CmpIOp>(
+        loc, CmpIPredicate::ne, v,
+        rewriter.create<ConstantOp>(loc, rewriter.getIntegerAttr(t, 0)));
+  else
+    llvm_unreachable("Unknown element type");
+}
+
 /// Generates a call that adds one element to a coordinate scheme.
+/// In particular, this generates code like the following:
+///   val = a[i1,..,ik];
+///   if val != 0
+///     t->add(val, [i1,..,ik], [p1,..,pk]);
 static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
                           Value ptr, Value tensor, Value ind, Value perm,
                           ValueRange ivs) {
-  Location loc = op->getLoc();
   StringRef name;
   Type eltType = tensor.getType().cast<ShapedType>().getElementType();
   if (eltType.isF64())
     name = "addEltF64";
   else if (eltType.isF32())
     name = "addEltF32";
   else if (eltType.isInteger(64))
     name = "addEltI64";
   else if (eltType.isInteger(32))
     name = "addEltI32";
   else if (eltType.isInteger(16))
     name = "addEltI16";
   else if (eltType.isInteger(8))
     name = "addEltI8";
   else
     llvm_unreachable("Unknown element type");
+  Location loc = op->getLoc();
   Value val = rewriter.create<tensor::ExtractOp>(loc, tensor, ivs);
-  // TODO: add if here?
+  Value cond = genIsNonzero(rewriter, loc, eltType, val);
+  scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, cond, /*else*/ false);
+  rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
   unsigned i = 0;
   for (auto iv : ivs) {
     Value idx = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(i++));
     rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
   }
   SmallVector<Value, 8> params;
   params.push_back(ptr);
   params.push_back(val);
matchAndRewrite(ConvertOp op, ArrayRef<Value> operands,
     //     for ik in dimk
     //       val = a[i1,..,ik]
     //       if val != 0
     //         t->add(val, [i1,..,ik], [p1,..,pk])
     //   s = newSparseTensor(t)
     // Note that the dense tensor traversal code is actually implemented
     // using MLIR IR to avoid having to expose too much low-level
     // memref traversal details to the runtime support library.
+    // Also note that the code below only generates the "new" ops and
+    // the loop-nest per se; whereas the entire body of the innermost
+    // loop is generated by genAddElt().
     Location loc = op->getLoc();
     ShapedType shape = resType.cast<ShapedType>();
     auto memTp =
         MemRefType::get({ShapedType::kDynamicSize}, rewriter.getIndexType());
     Value perm;
     Value ptr = genNewCall(rewriter, op, encDst, 2, perm);
     Value tensor = operands[0];
     Value arg = rewriter.create<ConstantOp>(

mlir/lib/ExecutionEngine/SparseUtils.cpp

#define IMPL2(RET, NAME, TYPE, LIB)                                            \
   }
 
 #define IMPL3(NAME, TYPE)                                                      \
   void *NAME(void *tensor, TYPE value, uint64_t *ibase, uint64_t *idata,       \
              uint64_t ioff, uint64_t isize, uint64_t istride, uint64_t *pbase, \
              uint64_t *pdata, uint64_t poff, uint64_t psize,                   \
              uint64_t pstride) {                                               \
     assert(istride == 1 && pstride == 1 && isize == psize);                    \
     uint64_t *indx = idata + ioff;                                             \
-    if (!value)                                                                \
-      return tensor;                                                           \
     uint64_t *perm = pdata + poff;                                             \
     std::vector<uint64_t> indices(isize);                                      \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indices[perm[r]] = indx[r];                                              \
     static_cast<SparseTensorCOO<TYPE> *>(tensor)->add(indices, value);         \
     return tensor;                                                             \
   }
 
 TEMPLATE(MemRef1DU64, uint64_t);