Add verifiers and tests for all GraphBLAS ops (#65)

Author: paul-tqh-nguyen

mlir_graphblas/src/include/GraphBLAS/GraphBLASDialect.td

@@ -22,11 +22,4 @@ def GraphBLAS_Dialect : Dialect {
     let cppNamespace = "::mlir::graphblas";
 }
 
-//===--------------------------------------------------------------------===//
-// Base graphblas operation definition.
-//===--------------------------------------------------------------------===//
-
-class GraphBLAS_Op<string mnemonic, list<OpTrait> traits = []> :
-        Op<GraphBLAS_Dialect, mnemonic, traits>;
-
 #endif // GRAPHBLAS_DIALECT

mlir_graphblas/src/include/GraphBLAS/GraphBLASOps.td

@@ -10,6 +10,10 @@
 include "GraphBLASDialect.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
+class GraphBLAS_Op<string mnemonic, list<OpTrait> traits = []> : Op<GraphBLAS_Dialect, mnemonic, traits> {
+    let verifier = [{ return ::verify(*this); }];
+}
+
 def GraphBLAS_TransposeOp : GraphBLAS_Op<"transpose", [NoSideEffect]> {
     let summary = "transpose operation";
     let description = [{
@@ -31,9 +35,6 @@ def GraphBLAS_TransposeOp : GraphBLAS_Op<"transpose", [NoSideEffect]> {
     let assemblyFormat = [{
            $input attr-dict `:` type($input) `to` type($output)
     }];
-
-    // TODO add custom verifier sanity checking the input and output types are sane
-    // let verifier = 
 }
 
 def GraphBLAS_MatrixSelectOp : GraphBLAS_Op<"matrix_select", [NoSideEffect, SameOperandsAndResultType]> {
@@ -55,9 +56,6 @@ def GraphBLAS_MatrixSelectOp : GraphBLAS_Op<"matrix_select", [NoSideEffect, Same
     let assemblyFormat = [{
            $input attr-dict `:` type($input)
     }];
-
-    // TODO add custom verifier sanity checking the selector attribute is sane
-    // let verifier = 
 }
 
 def GraphBLAS_MatrixReduceToScalarOp : GraphBLAS_Op<"matrix_reduce_to_scalar", [NoSideEffect]> {
@@ -78,9 +76,6 @@ def GraphBLAS_MatrixReduceToScalarOp : GraphBLAS_Op<"matrix_reduce_to_scalar", [
     let assemblyFormat = [{
            $input attr-dict `:` type($input) `to` type($output)
     }];
-
-    // TODO add custom verifier sanity checking the type of $output is sane ; sanity check the aggregator
-    // let verifier = 
 }
 
 def GraphBLAS_MatrixApplyOp : GraphBLAS_Op<"matrix_apply", [NoSideEffect]> {
@@ -105,9 +100,6 @@ def GraphBLAS_MatrixApplyOp : GraphBLAS_Op<"matrix_apply", [NoSideEffect]> {
     let assemblyFormat = [{
            $input `,` $thunk attr-dict `:` `(` type($input) `,` type($thunk) `)` `to` type($output)
     }];
-
-    // TODO add custom verifier sanity checking the types of $output and $thunk are sane ; sanity check the apply operator
-    // let verifier = 
 }
 
 def GraphBLAS_MatrixMultiplyOp : GraphBLAS_Op<"matrix_multiply", [NoSideEffect]> {
@@ -131,9 +123,6 @@ def GraphBLAS_MatrixMultiplyOp : GraphBLAS_Op<"matrix_multiply", [NoSideEffect]>
     let assemblyFormat = [{
            $a `,` $b (`,` $mask^)? attr-dict `:` `(` type($a) `,` type($b)  (`,` type($mask)^)? `)` `to` type($output)
     }];
-
-    // TODO add custom verifier sanity checking the types of the inputs and outputs and mask are sane ; sanity check the semiring
-    // let verifier = 
 }
 
 #endif // GRAPHBLAS_OPS

mlir_graphblas/src/include/GraphBLAS/GraphBLASUtils.h

@@ -5,9 +5,10 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "llvm/ADT/APInt.h"
 
 
-mlir::RankedTensorType getCSRTensorType(mlir::MLIRContext *context, mlir::Type valueType);
+mlir::RankedTensorType getCSRTensorType(mlir::MLIRContext *context, llvm::ArrayRef<int64_t> shape, mlir::Type valueType);
 mlir::CallOp callEmptyLike(mlir::OpBuilder &builder, mlir::ModuleOp &mod, mlir::Location loc, mlir::Value tensor);
 mlir::CallOp callDupTensor(mlir::OpBuilder &builder, mlir::ModuleOp &mod, mlir::Location loc, mlir::Value tensor);
 
@@ -20,4 +21,4 @@ mlir::CallOp callResizeIndex(mlir::OpBuilder &builder, mlir::ModuleOp &mod, mlir
 mlir::CallOp callResizeValues(mlir::OpBuilder &builder, mlir::ModuleOp &mod, mlir::Location loc,
                               mlir::Value tensor, mlir::Value size);
 
-#endif // GRAPHBLAS_GRAPHBLASUTILS_H
+#endif // GRAPHBLAS_GRAPHBLASUTILS_H

mlir_graphblas/src/lib/GraphBLAS/GraphBLASOps.cpp

@@ -4,9 +4,281 @@
 //
 //===--------------------------------------------------------------------===//
 
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "GraphBLAS/GraphBLASOps.h"
 #include "GraphBLAS/GraphBLASDialect.h"
 #include "mlir/IR/OpImplementation.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/None.h"
+
+using namespace mlir;
+using namespace mlir::graphblas;
+
+//===--------------------------------------------------------------------===//
+// Helpers
+//===--------------------------------------------------------------------===//
+
+enum CompressionType { CSR, CSC, EITHER };
+
+static llvm::Optional<std::string> checkCompressedSparseTensor(
+        Type inputType,
+        int inputIndex,
+	CompressionType compressionType
+    ) {
+  /*
+     Negative values for inputIndex indicate that the input type is the return type.
+     Otherwise, inputIndex indicates which arg inputType corresponds to.
+
+     Returns llvm::None if the given tensor is valid.
+     Returns a string explaining the problem otherwise.
+   */
+
+  std::string inputName = inputIndex < 0 ? "Return value" : "Operand #"+std::to_string(inputIndex);
+
+  mlir::sparse_tensor::SparseTensorEncodingAttr sparseEncoding =
+    mlir::sparse_tensor::getSparseTensorEncoding(inputType);
+  if (!sparseEncoding)
+    return inputName+" must be a sparse tensor.";
+
+  RankedTensorType inputTensorType = inputType.dyn_cast<RankedTensorType>();
+  if (inputTensorType.getRank() != 2)
+    return inputName+" must have rank 2.";
+
+  ArrayRef<mlir::sparse_tensor::SparseTensorEncodingAttr::DimLevelType> compression =
+    sparseEncoding.getDimLevelType();
+  if (compression[0] != mlir::sparse_tensor::SparseTensorEncodingAttr::DimLevelType::Dense ||
+      compression[1] != mlir::sparse_tensor::SparseTensorEncodingAttr::DimLevelType::Compressed)
+    return inputName+" must have CSR or CSC compression, i.e. must have "
+      "dimLevelType = [ \"dense\", \"compressed\" ] in the sparse encoding.";
+
+  if (compressionType != EITHER) {
+    
+    AffineMap dimOrdering = sparseEncoding.getDimOrdering();
+    unsigned dimOrdering0 = dimOrdering.getDimPosition(0); 
+    unsigned dimOrdering1 = dimOrdering.getDimPosition(1);
+
+    assert(compressionType == CSR || compressionType == CSC);
+    
+    if (compressionType == CSR) {
+      if (dimOrdering0 != 0 || dimOrdering1 != 1)
+	return inputName+" must have CSR compression.";
+    } else if (compressionType == CSC) {
+      if (dimOrdering0 != 1 || dimOrdering1 != 0)
+	return inputName+" must have CSC compression.";
+    }
+  }
+
+  return llvm::None;
+}
+
+//===--------------------------------------------------------------------===//
+// GraphBLAS Ops Methods
+//===--------------------------------------------------------------------===//
+
+static LogicalResult verify(MatrixApplyOp op) {
+  Type inputType = op.input().getType();
+  Type thunkType = op.thunk().getType();
+  Type resultType = op.getResult().getType();
+
+  llvm::Optional<std::string> inputCompressionErrorMessage = checkCompressedSparseTensor(inputType, 0, EITHER);
+  if (inputCompressionErrorMessage)
+    return op.emitError(inputCompressionErrorMessage.getValue());
+
+  llvm::Optional<std::string> resultCompressionErrorMessage = checkCompressedSparseTensor(resultType, -1, EITHER);
+  if (resultCompressionErrorMessage)
+    return op.emitError(resultCompressionErrorMessage.getValue());
+
+  RankedTensorType inputTensorType = inputType.dyn_cast<RankedTensorType>();
+  RankedTensorType resultTensorType = resultType.dyn_cast<RankedTensorType>();
+
+  if (inputTensorType.getElementType() != thunkType)
+    return op.emitError("Element type of input tensor does not match type of thunk.");
+
+  if (resultTensorType.getElementType() != thunkType)
+    // TODO this is not always correct, e.g. matrix_apply_less_than(tensor<f64>, 2.3) -> tensor<i1>.
+    return op.emitError("Element type of result tensor does not match type of thunk.");
+
+  ArrayRef<int64_t> inputShape = inputTensorType.getShape();
+  ArrayRef<int64_t> resultShape = resultTensorType.getShape();
+
+  // TODO intelligently handle arbitrarily shaped tensors, i.e. tensors with shapes using "?"
+  if (inputShape[0] != resultShape[0] || inputShape[1] != resultShape[1])
+    return op.emitError("Input shape does not match output shape.");
+
+  static const std::vector<std::string> supportedOperators{"min"};
+  std::string applyOperator = op.apply_operator().str();
+  bool operatorSupported = std::find(supportedOperators.begin(), supportedOperators.end(), applyOperator)
+    != supportedOperators.end();
+  if (!operatorSupported)
+    return op.emitError("\""+applyOperator+"\" is not a supported operator.");
+
+  return success();
+}
+
+static LogicalResult verify(MatrixMultiplyOp op) {
+  Type aType = op.a().getType();
+  Type bType = op.b().getType();
+  Type resultType = op.getResult().getType();
+
+  llvm::Optional<std::string> aCompressionErrorMessage = checkCompressedSparseTensor(aType, 0, CSR);
+  if (aCompressionErrorMessage)
+    return op.emitError(aCompressionErrorMessage.getValue());
+
+  llvm::Optional<std::string> bCompressionErrorMessage = checkCompressedSparseTensor(bType, 1, CSC);
+  if (bCompressionErrorMessage)
+    return op.emitError(bCompressionErrorMessage.getValue());
+
+  llvm::Optional<std::string> resultCompressionErrorMessage = checkCompressedSparseTensor(resultType, -1, CSR);
+  if (resultCompressionErrorMessage)
+    return op.emitError(resultCompressionErrorMessage.getValue());
+
+  static const std::vector<std::string> supportedSemirings{"plus_times", "plus_pair", "plus_plus"};
+  std::string semiring = op.semiring().str();
+  bool semiringSupported = std::find(supportedSemirings.begin(), supportedSemirings.end(), semiring)
+    != supportedSemirings.end();
+  if (!semiringSupported)
+    return op.emitError("\""+semiring+"\" is not a supported semiring.");
+
+  RankedTensorType aTensorType = aType.dyn_cast<RankedTensorType>();
+  RankedTensorType bTensorType = bType.dyn_cast<RankedTensorType>();
+  RankedTensorType resultTensorType = resultType.dyn_cast<RankedTensorType>();
+
+  ArrayRef<int64_t> aShape = aTensorType.getShape();
+  ArrayRef<int64_t> bShape = bTensorType.getShape();
+  ArrayRef<int64_t> resultShape = resultTensorType.getShape();
+  // TODO intelligently handle arbitrarily shaped tensors, i.e. tensors with shapes using "?"
+  if (aShape[1] != bShape[0])
+    return op.emitError("Operand shapes are incompatible.");
+  if (resultShape[0] != aShape[0] || resultShape[1] != bShape[1])
+    return op.emitError("Operand shapes incompatible with output shape.");
+
+  if (aTensorType.getElementType() != bTensorType.getElementType())
+    return op.emitError("Operand element types must be identical.");
+  if (aTensorType.getElementType() != resultTensorType.getElementType())
+    return op.emitError("Result element type differs from the input element types.");
+
+  Value mask = op.mask();
+  if (mask) {
+    Type maskType = mask.getType();
+    llvm::Optional<std::string> maskCompressionErrorMessage = checkCompressedSparseTensor(maskType, 2, CSR);
+    if (maskCompressionErrorMessage)
+      return op.emitError(maskCompressionErrorMessage.getValue());
+
+    RankedTensorType maskTensorType = maskType.dyn_cast<RankedTensorType>();
+    ArrayRef<int64_t> maskShape = maskTensorType.getShape();
+    if (resultShape[0] != maskShape[0] || resultShape[1] != maskShape[1])
+      return op.emitError("Mask shape must match output shape.");
+  }
+
+  return success();
+}
+
+static LogicalResult verify(MatrixReduceToScalarOp op) {
+  Type operandType = op.input().getType();
+
+  llvm::Optional<std::string> compressionErrorMessage = checkCompressedSparseTensor(operandType, 0, EITHER);
+  if (compressionErrorMessage)
+    return op.emitError(compressionErrorMessage.getValue());
+
+  static const std::vector<std::string> supportedAggregators{"sum"};
+  std::string aggregator = op.aggregator().str();
+  bool aggregatorSupported = std::find(supportedAggregators.begin(), supportedAggregators.end(), aggregator)
+    != supportedAggregators.end();
+  if (!aggregatorSupported)
+    return op.emitError("\""+aggregator+"\" is not a supported aggregator.");
+
+  Type resultType = op.getResult().getType();
+  RankedTensorType operandTensorType = operandType.dyn_cast<RankedTensorType>();
+  if (resultType != operandTensorType.getElementType())
+    return op.emitError("Operand and output types are incompatible.");
+
+  return success();
+}
+
+static LogicalResult verify(MatrixSelectOp op) {
+  // input and result types are already guaranteed to be the same
+  Type resultType = op.getResult().getType();
+
+  llvm::Optional<std::string> resultCompressionErrorMessage = checkCompressedSparseTensor(resultType, -1, EITHER);
+  if (resultCompressionErrorMessage)
+    return op.emitError(resultCompressionErrorMessage.getValue());
+
+  static const std::vector<std::string> supportedSelectors{"triu", "tril", "gt0"};
+  std::string selector = op.selector().str();
+  bool selectorSupported = std::find(supportedSelectors.begin(), supportedSelectors.end(), selector)
+    != supportedSelectors.end();
+  if (!selectorSupported)
+    return op.emitError("\""+selector+"\" is not a supported selector.");
+
+  return success();
+}
+
+static LogicalResult verify(TransposeOp op) {
+  Type inputType = op.input().getType();
+  Type resultType = op.getResult().getType();
+
+  llvm::Optional<std::string> inputCompressionErrorMessage = checkCompressedSparseTensor(inputType, 0, EITHER);
+  if (inputCompressionErrorMessage)
+    return op.emitError(inputCompressionErrorMessage.getValue());
+
+  llvm::Optional<std::string> resultCompressionErrorMessage = checkCompressedSparseTensor(resultType, -1, EITHER);
+  if (resultCompressionErrorMessage)
+    return op.emitError(resultCompressionErrorMessage.getValue());
+
+  // TODO intelligently handle arbitrarily shaped tensors, i.e. tensors with shapes using "?"
+
+  RankedTensorType inputTensorType = inputType.dyn_cast<RankedTensorType>();
+  RankedTensorType resultTensorType = resultType.dyn_cast<RankedTensorType>();
+
+  if (inputTensorType.getElementType() != resultTensorType.getElementType())
+    return op.emitError("Input and output tensors have different element types.");
+
+  ArrayRef<int64_t> inputShape = inputTensorType.getShape();
+  ArrayRef<int64_t> resultShape = resultTensorType.getShape();
+
+  mlir::sparse_tensor::SparseTensorEncodingAttr inputSparseEncoding =
+    mlir::sparse_tensor::getSparseTensorEncoding(inputType);
+
+  mlir::sparse_tensor::SparseTensorEncodingAttr resultSparseEncoding =
+    mlir::sparse_tensor::getSparseTensorEncoding(resultType);
+
+  bool swapSizes = op.swap_sizes();
+  if (swapSizes) {
+    if (inputShape[0] != resultShape[1] || inputShape[1] != resultShape[0])
+      return op.emitError("Input and output shapes are expected to be swapped.");
+    if (inputSparseEncoding != resultSparseEncoding)
+      return op.emitError("Input and output tensors are expected to have identical sparse encodings.");
+  } else {
+    if (inputShape[0] != resultShape[0] || inputShape[1] != resultShape[1])
+      return op.emitError("Input and output shapes are expected to be the same.");
+
+    AffineMap inputDimOrdering = inputSparseEncoding.getDimOrdering();
+    AffineMap resultDimOrdering = resultSparseEncoding.getDimOrdering();
+    unsigned inputDimOrdering0 = inputDimOrdering.getDimPosition(0);
+    unsigned inputDimOrdering1 = inputDimOrdering.getDimPosition(1);
+    unsigned resultDimOrdering0 = resultDimOrdering.getDimPosition(0);
+    unsigned resultDimOrdering1 = resultDimOrdering.getDimPosition(1);
+    if (inputDimOrdering0 != resultDimOrdering1 || inputDimOrdering1 != resultDimOrdering0)
+      return op.emitError("Sparse encoding dimension orderings of input and result tensors "
+			  "expected to be swapped.");
+
+    // TODO should we be more lenient like the sparse tensor dialect is via isMatchingWidth?
+    // see llvm-project/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+    unsigned inputPointerBitWidth = inputSparseEncoding.getPointerBitWidth();
+    unsigned resultPointerBitWidth = resultSparseEncoding.getPointerBitWidth();
+    if (inputPointerBitWidth != resultPointerBitWidth)
+      return op.emitError("Sparse encoding pointer bit widths of input and result tensors do not match.");
+
+    unsigned inputIndexBitWidth = inputSparseEncoding.getIndexBitWidth();
+    unsigned resultIndexBitWidth = resultSparseEncoding.getIndexBitWidth();
+    if (inputIndexBitWidth != resultIndexBitWidth)
+      return op.emitError("Sparse encoding index bit widths of input and result tensors do not match.");
+
+    // dimLevelType values guaranteed to be the same since we already checked earlier
+  }
+
+  return success();
+}
 
 #define GET_OP_CLASSES
 #include "GraphBLAS/GraphBLASOps.cpp.inc"