[Mlir-commits] [mlir] [mlir][sparse] test for linalg tensor semantics (PR #70254)

[Aart Bik]

https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/70254

This test used to be here, but somehow got lost while linalg rewrote their interfaces. It is essential to test this on entry of sparsification, however, since all subsequent analysis simply assumes tensor types.

Fixes:
https://github.com/llvm/llvm-project/issues/64325

>From d3229adac5d3d9dbb0c22a4469a7cbd8079d0a01 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Wed, 25 Oct 2023 13:59:06 -0700
Subject: [PATCH] [mlir][sparse] test for linalg tensor semantics

This test used to be here, but somehow got lost while
linalg rewrote their interfaces. It is essential to
test this on entry of sparsification, however, since
all subsequent analysis simply assumes tensor types.

Fixes:
https://github.com/llvm/llvm-project/issues/64325
---
 .../SparseTensor/Transforms/Sparsification.cpp   | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index e47131bb78318ab..ede4d38bb281277 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -1939,9 +1939,12 @@ struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
-    // Only accept single output operations without affine index on sparse
-    // output.
-    if (op.getNumDpsInits() != 1 || hasNonTrivialAffineOnSparseOut(op))
+    // Only accept single output operations with pure tensor semantics.
+    if (op.getNumDpsInits() != 1 || !op.hasTensorSemantics())
+      return failure();
+
+    // Only accept trivial affine indices.
+    if (hasNonTrivialAffineOnSparseOut(op))
       return failure();
 
     // Sets up a code generation environment.
@@ -1951,7 +1954,6 @@ struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
     // TODO: we should probably always use slice-based codegen whenever
     // possible, we can even intermix slice-based and filter-loop based codegen.
     bool idxReducBased = options.enableIndexReduction && numFilterLoops != 0;
-
     // If we have indexing map like (d0) -> (0, d0), there might be more
     // levels then loops because of the constant index, that means we can not
     // use numLoops as the upper bound for ranks of all tensors.
@@ -1964,9 +1966,7 @@ struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
         maxLvlRank = std::max(maxLvlRank, SparseTensorType(rtp).getLvlRank());
       }
     }
-
-    // If we uses slice based algorithm for affine index, we do not need filter
-    // loop.
+    // A slice based algorithm for affine indices does not need filter loops.
     CodegenEnv env(op, options, numTensors, numLoops,
                    /*numFilterLoops=*/idxReducBased ? 0 : numFilterLoops,
                    maxLvlRank);
@@ -2006,7 +2006,6 @@ struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
     // to resolve cycles by inserting a conversion.
     bool isAdmissible = false;
     bool hasCycle = true;
-
     // A const list of all masks that we used for iteration graph
     // computation. Must be ordered from more strict to less strict.
     // Ideally (though might not be guaranteed), the earlier a constraint mask
@@ -2030,7 +2029,6 @@ struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
                  ? failure() // TODO: should cycle be resolved differently?
                  : resolveCycle(env, rewriter); // one last shot
     }
-
     if (!isAdmissible)
       return failure(); // inadmissible expression, reject