[mlir][sparse] misc code cleanup

[wrengr]

Summary

• Flattening/simplifying some nested conditionals
• const-ifying some local variables

Depends On D143800

[wrengr]

Using the getRankedTensorType method in lieu of static_cast

[wrengr]

git-clang-format

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

for (const auto &it : llvm::zip(op.getInputs(), adaptor.getInputs())) {
             });
       }
       // Accumulate offset.
       // TODO: avoid calling sparseDimSize multiple times by caching the result!
       Value curDim =
           createOrFoldDimCall(rewriter, loc, srcTp, adaptedOp, concatDim);
       offset = rewriter.create<arith::AddIOp>(loc, offset, curDim);
     }
-    if (dstTp.hasEncoding()) {
-      if (!allDense) {
-        // In sparse output case, the destination holds the COO.
-        Value coo = dst;
-        dst = params.genNewCall(Action::kFromCOO, coo);
-        // Release resources.
-        genDelCOOCall(rewriter, loc, elemTp, coo);
-      } else {
-        dst = dstTensor;
-      }
+    if (!dstTp.hasEncoding()) {
+      rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(
+          op, dstTp.getRankedTensorType(), dst);
+    } else if (allDense) {
+      rewriter.replaceOp(op, dstTensor);
+    } else {
       rewriter.replaceOp(op, dst);
-    } else {
-      rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(
-          op, dstTp.getRankedTensorType(), dst);
     }
     return success();
   }
 };
 
 /// Sparse conversion rule for the output operator.
 class SparseTensorOutConverter : public OpConversionPattern<OutOp> {
 public:

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

if (dstTp.hasEncoding()) {
           const auto stt = getSparseTensorType(i);
           allOrdered = stt.isAllOrdered() && stt.isIdentity();
           if (!allOrdered)
             break;
         }
       }
 
       needTmpCOO = !allDense && !allOrdered;
+      const RankedTensorType tp = needTmpCOO ? getUnorderedCOOFromType(dstTp)
+                                             : dstTp.getRankedTensorType();
+      encDst = needTmpCOO ? getSparseTensorEncoding(tp) : encDst;
       SmallVector<Value> dynSizes;
       getDynamicSizes(dstTp, sizes, dynSizes);
-      RankedTensorType tp = dstTp;
-      if (needTmpCOO) {
-        tp = getUnorderedCOOFromType(dstTp);
-        encDst = getSparseTensorEncoding(tp);
-      }
       dst = rewriter.create<AllocTensorOp>(loc, tp, dynSizes).getResult();
       if (allDense) {
         // Create a view of the values buffer to match the unannotated dense
         // tensor.
         Value valuesBuffer = genToValues(rewriter, loc, dst);
         Value idxBuffer =
             genAlloca(rewriter, loc, dimRank, rewriter.getIndexType(),
                       /*staticShape=*/true);
for (Value input : op.getInputs()) {
         dst = foreachOp.getResult(0);
         initArgs[0] = dst;
       }
     }
 
     // Temp variable to avoid needing to call `getRankedTensorType`
     // in the three use-sites below.
     const RankedTensorType dstRTT = dstTp;
-    if (encDst) {
-      if (!allDense) {
-        dst = rewriter.create<LoadOp>(loc, dst, true);
-        if (needTmpCOO) {
-          Value tmpCoo = dst;
-          dst = rewriter.create<ConvertOp>(loc, dstRTT, tmpCoo).getResult();
-          rewriter.create<DeallocTensorOp>(loc, tmpCoo);
-        }
-      } else {
-        dst = rewriter.create<ConvertOp>(loc, dstRTT, annotatedDenseDst)
-                  .getResult();
-      }
+    if (!encDst) {
+      rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstRTT, dst);
+    } else if (allDense) {
+      rewriter.replaceOp(
+          op, rewriter.create<ConvertOp>(loc, dstRTT, annotatedDenseDst)
+                  .getResult());
+    } else {
       rewriter.replaceOp(op, dst);
-    } else {
-      rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstRTT, dst);
     }
     return success();
   }
 };
 
 /// Sparse rewriting rule for the convert operator.
 struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
   using OpRewritePattern::OpRewritePattern;