[mlir][sparse] support constant affine expression on dense dimension

[Peiming]

update test cases.

[Peiming]

update comments

[Peiming]

small fixes.

[Peiming]

small fix

[Peiming]

clang-format

[Peiming]

rebase.

[Peiming]

rebase.

[aartbik]

This comment was removed by aartbik.

[Peiming]

update comments.

[aartbik]

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

1177

what does it mean if we reach here?

1301

Until now

1302

addressed -> addresses

dependend -> dependent

[Peiming]

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

1177

Good question, I should assert(false) here, it should have been rejected... But that code is in the upcoming patch (see D138171)

// TODO: support compound affine expression on sparse output.
if (getNumCompoundAffineOnSparseDims(op.getMatchingIndexingMap(lhs),
                                     lhs->get()) != 0)
  return false;

[Peiming]

address comments.

[aartbik]

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

1177

then it is ok, since that will go in shortly too ;-)

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

 
 /// Generates a for-loop on a single index.
 static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, bool isOuter, bool isInner,
                          unsigned idx, size_t tid, size_t dim,
                          ArrayRef<size_t> extraTids,
                          ArrayRef<size_t> extraDims) {
   Location loc = op.getLoc();
-  auto iteratorTypes = op.getIteratorTypesArray();
   bool isSparse = isCompressedDLT(merger.getDimLevelType(tid, idx)) ||
                   isSingletonDLT(merger.getDimLevelType(tid, idx));
   bool isParallel = isParallelFor(codegen, isOuter, isSparse);
 
   Operation *loop =
       genLoopBoundary(codegen, merger, [&](MutableArrayRef<Value> reduc) {
         return codegen.loopEmitter.enterLoopOverTensorAtDim(
             builder, loc, tid, dim, reduc, isParallel, extraTids, extraDims);
static bool startLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   return false;
 }
 
 static void translateBitsToTidDimPairs(
     Merger &merger, CodeGen &codegen, linalg::GenericOp op, unsigned li,
     unsigned idx, SmallVectorImpl<size_t> &condTids,
     SmallVectorImpl<size_t> &condDims, SmallVectorImpl<size_t> &extraTids,
     SmallVectorImpl<size_t> &extraDims, SmallVectorImpl<size_t> &affineTids,
     SmallVectorImpl<size_t> &affineDims, SmallVectorImpl<AffineExpr> &exps) {
 
   const BitVector &all = merger.lat(li).bits;
   const BitVector &simple = merger.lat(li).simple;
 
   // Converts bits to array + dim pair
   merger.foreachTidDimPairInBits(all, [&, idx](unsigned b, unsigned tid,
                                                Optional<unsigned> dim,
                                                DimLevelType dlt) {
if (simple.test(b)) {
         // when
         // 1. The previous level is an (at-level) invariant compound dense
         // affine (with no corresponding loop idx); or
         // 2. The previous level is being generated right now.
         if (exp.isa<AffineConstantExpr>()) {
           // TODO:  Should we come up with a more adhersive way to handle
           // constant expression? We now requires two (somehow ad-hoc) code for
           // it.
-          assert(false && "do not support constant affine");
+          if (i != 0 && // i == 0 cases are handled in genConstantDenseAddress
+              ((!affineTids.empty() && affineTids.back() == tid &&
+                affineDims.back() == i - 1) ||          // Condition 1
+               merger.getLoopIdx(tid, i - 1) == idx)) { // Condition 2
+            affineTids.push_back(tid);
+            affineDims.push_back(i);
+            exps.push_back(exp);
           }
-
+        } else {
           bool atLevel = false;
           if (isInvariantAffine(codegen, exp, idx, atLevel) && atLevel) {
             // If the compound affine is invariant and we are right at the
             // level. We need to generate the address according to the affine
             // expression. This is also the best place we can do it to avoid
             // putting it inside inner loops.
             // NOTE: It assumes that the levels of the input tensor are
             // initialized in order, another more admissible approach might be
             // accepting out-of-order access between consecutive dense levels.
             affineTids.push_back(tid);
             affineDims.push_back(i);
             exps.push_back(exp);
           }
         }
       }
+    }
   });
 
   if (isDenseDLT(merger.getDimLevelType(merger.getOutTensorID(), idx))) {
     // Note that we generate dense indices of the output tensor
     // unconditionally, since they may not appear in the lattice, but may be
     // needed for linearized codegen.
     // Only dense dimensions should be optimized from conditions.
     auto dim = merger.getDimNum(merger.getOutTensorID(), idx).value();
static Operation *startLoop(Merger &merger, CodeGen &codegen,
   Operation *loop = genLoop(merger, codegen, builder, op, at, needsUniv,
                             condTids, condDims, extraTids, extraDims);
 
   for (auto [tid, dim, exp] : llvm::zip(affineTids, affineDims, affines)) {
     codegen.loopEmitter.genDenseAffineAddressAtCurLevel(builder, op.getLoc(),
                                                         tid, dim, exp);
   }
   return loop;
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.
 static bool endLoop(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                     linalg::GenericOp op, Operation *loop, unsigned idx,
                     unsigned li, bool needsUniv) {
   // End a while-loop.
   if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
     finalizeWhileOp(merger, codegen, rewriter, op, idx, needsUniv,
                     merger.lat(li).bits, whileOp);
static void genResult(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
   } else {
     // To rematerialize an non-annotated tensor, simply load it
     // from the bufferized value.
     Value val = codegen.loopEmitter.getValBuffer().back(); // value array
     rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, val);
   }
 }
 
+static void genConstantDenseAddress(CodeGen &codegen, RewriterBase &rewriter,
+                                    linalg::GenericOp op) {
+  // We can generates address for constant affine expression before any loops
+  // starting from the first level as they do not depend on any thing.
+  // E.g., [Dense, Dense, Sparse] -> (1, 2, d0), the addresses for the first two
+  // levels can be determined before loops.
+  for (OpOperand *input : op.getDpsInputOperands()) {
+    ArrayRef<AffineExpr> affines =
+        op.getMatchingIndexingMap(input).getResults();
+    auto enc = getSparseTensorEncoding(input->get().getType());
+    if (enc) {
+      for (unsigned i = 0, e = affines.size(); i < e; i++) {
+        AffineExpr affine = affines[toOrigDim(enc, i)];
+        if (isDenseDLT(getDimLevelType(enc, i)) &&
+            affine.isa<AffineConstantExpr>()) {
+          codegen.loopEmitter.genDenseAffineAddressAtCurLevel(
+              rewriter, op.getLoc(), input->getOperandNumber(), i, affine);
+        } else {
+          // Breaks on first non-dense non-constant level.
+          break;
+        }
+      }
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Sparse compiler rewriting methods.
 //===----------------------------------------------------------------------===//
 
 namespace {
-
 /// Sparse rewriting rule for generic Lingalg operation.
 struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
 public:
   GenericOpSparsifier(MLIRContext *context, SparsificationOptions o)
       : OpRewritePattern<linalg::GenericOp>(context), options(o) {}
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
LogicalResult matchAndRewrite(linalg::GenericOp op,
     SmallVector<Value> tensors;
     for (OpOperand &t : op->getOpOperands())
       tensors.push_back(t.get());
 
     // Recursively generates code if admissible.
     CodeGen codegen(options, tensors, numTensors, numLoops, sparseOut,
                     outerParNest, topSort);
     genBuffers(merger, codegen, rewriter, op);
+    genConstantDenseAddress(codegen, rewriter, op);
     genStmt(merger, codegen, rewriter, op, exp, 0);
     genResult(merger, codegen, rewriter, op);
     return success();
   }
 
 private:
   // Last resort cycle resolution.
   LogicalResult resolveCycle(Merger &merger, PatternRewriter &rewriter,