[mlir][sparse][gpu] guard matvec COO AoS

[aartbik]

rebased with main

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

return aTp.isDenseLvl(0) && aTp.isCompressedLvl(1) && aTp.isOrderedLvl(1) &&
          (aTp.getElementType().isF64() || aTp.getElementType().isF32()) &&
          (aTp.getCrdWidth() == 0 || aTp.getCrdWidth() == 32 ||
           aTp.getCrdWidth() == 64);
 }
 
 /// Test for admissible types on operands (with output parameter `isCOO`).
 static bool areAdmissibleTypes(SparseTensorType aTp, SparseTensorType bTp,
                                SparseTensorType cTp, bool enableRT,
-                               bool &isCOO) {
+                               bool isMatVec, bool &isCOO) {
   if (bTp.hasEncoding() || cTp.hasEncoding())
     return false;
   if (isAdmissibleCOO(aTp)) {
     isCOO = true;
 #ifdef CUSPARSE_COO_AOS
-    return true;
+    return isMatVec;
 #else
     return enableRT;
 #endif
   }
   return isAdmissibleCSR(aTp);
 }
 
 /// Generates the first positions/coordinates of a sparse matrix.
static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
   Value y = op.getOperand(2); // we have y = Ax
   SmallVector<Value> tokens;
 
   // Only admissible sparse matrix format and dense vectors.
   bool isCOO = false;
   SparseTensorType aTp = getSparseTensorType(a);
   SparseTensorType xTp = getSparseTensorType(x);
   SparseTensorType yTp = getSparseTensorType(y);
-  if (!areAdmissibleTypes(aTp, xTp, yTp, enableRT, isCOO))
+  if (!areAdmissibleTypes(aTp, xTp, yTp, enableRT, /*isMatVec=*/true, isCOO))
     return failure();
 
   // Start sparse kernel and copy data from host to device.
   //   a : memR/memC/memV -> rowA,colA,valA
   //   x : memX           -> vecX
   //   y : memY           -> vecY
   Value nseA = rewriter.create<NumberOfEntriesOp>(loc, a);
   Value szY = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
   Value c = op.getOperand(2); // we have C = AB
   SmallVector<Value> tokens;
 
   // Only admissible sparse matrix format and dense matrices.
   bool isCOO = false;
   SparseTensorType aTp = getSparseTensorType(a);
   SparseTensorType bTp = getSparseTensorType(b);
   SparseTensorType cTp = getSparseTensorType(c);
-  if (!areAdmissibleTypes(aTp, bTp, cTp, enableRT, isCOO))
+  if (!areAdmissibleTypes(aTp, bTp, cTp, enableRT, /*isMatVec=*/false, isCOO))
     return failure();
 
   // Start sparse kernel and copy data from host to device.
   //   a : memR/memC/memV -> rowA,colA,valA
   //   b : bufB           -> matA
   //   c : bufC           -> matC
   Value nseA = rewriter.create<NumberOfEntriesOp>(loc, a);
   Value szm = linalg::createOrFoldDimOp(rewriter, loc, a, 0);