incorporate code review feedback

Author: seibert

mlir_graphblas/src/include/GraphBLAS/GraphBLASUtils.h

@@ -14,6 +14,9 @@ bool typeIsCSC(mlir::Type inputType);
 mlir::RankedTensorType getCSRTensorType(mlir::MLIRContext *context, llvm::ArrayRef<int64_t> shape, mlir::Type valueType);
 mlir::RankedTensorType getCSCTensorType(mlir::MLIRContext *context, llvm::ArrayRef<int64_t> shape, mlir::Type valueType);
 
+int64_t getRank(mlir::Type inputType);
+int64_t getRank(mlir::Value inputValue);
+
 mlir::Value convertToExternalCSR(mlir::OpBuilder &builder, mlir::ModuleOp &mod, mlir::Location loc, mlir::Value input);
 mlir::Value convertToExternalCSC(mlir::OpBuilder &builder, mlir::ModuleOp &mod, mlir::Location loc, mlir::Value input);
 mlir::Value callEmptyLike(mlir::OpBuilder &builder, mlir::ModuleOp &mod, mlir::Location loc, mlir::Value tensor);

mlir_graphblas/src/lib/GraphBLAS/GraphBLASOps.cpp

@@ -13,6 +13,7 @@
 #include "llvm/ADT/None.h"
 
 #include "GraphBLAS/GraphBLASOpsEnums.cpp.inc"
+#include "GraphBLAS/GraphBLASUtils.h"
 
 using namespace mlir;
 using namespace mlir::graphblas;
@@ -99,17 +100,6 @@ static llvm::Optional<std::string> checkCompressedVector(
   return llvm::None;
 }
 
-static int64_t getRank(Type inputType)
-{
-  mlir::sparse_tensor::SparseTensorEncodingAttr sparseEncoding =
-    mlir::sparse_tensor::getSparseTensorEncoding(inputType);
-  if (!sparseEncoding)
-    return -1;
-
-  RankedTensorType inputTensorType = inputType.dyn_cast<RankedTensorType>();
-  return inputTensorType.getRank();
-}
-
 //===--------------------------------------------------------------------===//
 // GraphBLAS Ops Methods
 //===--------------------------------------------------------------------===//
@@ -310,7 +300,6 @@ static LogicalResult verifyMatrixMultiplyArgs(T op, bool checkResultTensorType)
       resultTensorType = resultType.dyn_cast<RankedTensorType>();
       resultShape = resultTensorType.getShape();
       resultRank = getRank(resultType);
-      RankedTensorType resultTensorType = resultType.dyn_cast<RankedTensorType>();
       resultElementType = resultTensorType.getElementType();
     }
   }

mlir_graphblas/src/lib/GraphBLAS/GraphBLASOptimizePass.cpp

@@ -97,6 +97,9 @@ class FuseMatrixMultiplyReduceRewrite : public OpRewritePattern<graphblas::Matri
     if (predecessor != nullptr && predecessor->hasOneUse()) {
       Location loc = op->getLoc();
 
+      if (getRank(predecessor.a()) < 2 || getRank(predecessor.b()) < 2)
+        return failure();
+
       // Build new MatrixMultiplyReduceToScalarGeneric op with the operands and regions of the multiply,
       // then add in the aggregator from the reduce
       ValueRange operands = predecessor.getOperands();
@@ -190,7 +193,6 @@ class FuseMatrixMultiplyApplyRewrite : public OpRewritePattern<graphblas::Matrix
       RegionRange applyExtensions = op.extensions();
 
       RankedTensorType tensorType = predecessor.a().getType().dyn_cast<RankedTensorType>();
-      Type valueType = tensorType.getElementType();
 
       graphblas::MatrixMultiplyGenericOp newMultOp = rewriter.create<graphblas::MatrixMultiplyGenericOp>(loc,
                                 op->getResultTypes(), operands, attributes.getAttrs(),

mlir_graphblas/src/lib/GraphBLAS/GraphBLASUtils.cpp

@@ -92,6 +92,23 @@ bool typeIsCSC(Type inputType) {
   return true;
 }
 
+int64_t getRank(Type inputType)
+{
+  mlir::sparse_tensor::SparseTensorEncodingAttr sparseEncoding =
+      mlir::sparse_tensor::getSparseTensorEncoding(inputType);
+  if (!sparseEncoding)
+    return -1;
+
+  RankedTensorType inputTensorType = inputType.dyn_cast<RankedTensorType>();
+  return inputTensorType.getRank();
+}
+
+int64_t getRank(Value inputValue)
+{
+  Type inputType = inputValue.getType();
+  return getRank(inputType);
+}
+
 // make Compressed Vector type
 RankedTensorType getCompressedVectorType(MLIRContext *context, ArrayRef<int64_t> shape, Type valueType)
 {