[mlir][CAPI] Add CAPI bindings for the sparse_tensor dialect.

[stellaraccident]

Summary

• Adds dialect registration, hand coded 'encoding' attribute and test.
• An MLIR CAPI tablegen backend for attributes does not exist, and this is a relatively complicated case. I opted to hand code it in a canonical way for now, which will provide a reasonable blueprint for building out the tablegen version in the future.
• Also added a (local) CMake function for declaring new CAPI tests, since it was getting repetitive/buggy.

[stellaraccident]

Also adding Alex for a look, since I think this is the first such case (dialect attribute) we have in tree.

[ftynse]

mlir/test/CAPI/sparse_tensor.c

26–30

Nit: can this use a raw literal or be wrapper in // clang-format off / // clang-format on to avoid awkward string reformatting?

[stellaraccident]

mlir/test/CAPI/sparse_tensor.c

26–30

It's C, so opted for the clang-format option and reflowed.

[stellaraccident]

This broke the windows build bot, and I missed it. Rolling forward fix in: https://reviews.llvm.org/D102189

mlir/include/mlir-c/Dialect/SparseTensor.h

+//===-- mlir-c/Dialect/SparseTensor.h - C API for SparseTensor ----*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_SPARSE_TENSOR_H
+#define MLIR_C_DIALECT_SPARSE_TENSOR_H
+
+#include "mlir-c/AffineMap.h"
+#include "mlir-c/Registration.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor);
+
+/// Dimension level types that define sparse tensors:
+///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE      - dimension is dense, every
+///   entry is stored
+///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED - dimension is sparse,
+///   only nonzeros are stored.
+///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON  - dimension contains single
+///   coordinate, no siblings.
+///
+/// These correspond to SparseTensorEncodingAttr::DimLevelType in the C++ API.
+/// If updating, keep them in sync and update the static_assert in the impl
+/// file.
+enum MlirSparseTensorDimLevelType {
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON,
+};
+
+//===----------------------------------------------------------------------===//
+// SparseTensorEncodingAttr
+//===----------------------------------------------------------------------===//
+
+/// Checks whether the given attribute is a sparse_tensor.encoding attribute.
+MLIR_CAPI_EXPORTED bool
+mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr);
+
+/// Creates a sparse_tensor.encoding attribute with the given parameters.
+MLIR_CAPI_EXPORTED MlirAttribute mlirSparseTensorEncodingAttrGet(
+    MlirContext ctx, intptr_t numDimLevelTypes,
+    enum MlirSparseTensorDimLevelType const *dimLevelTypes,
+    MlirAffineMap dimOrdering, int pointerBitWidth, int indexBitWidth);
+
+/// Returns the number of dim level types in a sparse_tensor.encoding attribute.
+MLIR_CAPI_EXPORTED intptr_t
+mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr);
+
+/// Returns a specified dim level type in a sparse_tensor.encoding attribute.
+MLIR_CAPI_EXPORTED enum MlirSparseTensorDimLevelType
+mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t pos);
+
+/// Returns the dimension ordering in a sparse_tensor.encoding attribute.
+MLIR_CAPI_EXPORTED MlirAffineMap
+mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr);
+
+/// Returns the pointer bit width in a sparse_tensor.encoding attribute.
+MLIR_CAPI_EXPORTED int
+mlirSparseTensorEncodingAttrGetPointerBitWidth(MlirAttribute attr);
+
+/// Returns the index bit width in a sparse_tensor.encoding attribute.
+MLIR_CAPI_EXPORTED int
+mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_SPARSE_TENSOR_H

mlir/lib/CAPI/Dialect/CMakeLists.txt

-# TODO: Make the check source feature optional as an argument on *_add_library.
-set(LLVM_OPTIONAL_SOURCES
-  Async.cpp
-  AsyncPasses.cpp
-  GPU.cpp
-  GPUPasses.cpp
-  Linalg.cpp
-  LinalgPasses.cpp
-  SCF.cpp
-  Shape.cpp
-  Standard.cpp
-  Tensor.cpp
-)
-
 add_mlir_public_c_api_library(MLIRCAPIAsync
   Async.cpp
   AsyncPasses.cpp
 
+  PARTIAL_SOURCES_INTENDED
   DEPENDS
   MLIRAsyncPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRAsync
   MLIRAsyncTransforms
   MLIRPass
 )
 
 add_mlir_public_c_api_library(MLIRCAPIGPU
   GPU.cpp
   GPUPasses.cpp
 
+  PARTIAL_SOURCES_INTENDED
   DEPENDS
   MLIRGPUPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRGPU
   MLIRPass
 )
 
 add_mlir_public_c_api_library(MLIRCAPILinalg
   Linalg.cpp
   LinalgPasses.cpp
 
+  PARTIAL_SOURCES_INTENDED
   DEPENDS
   MLIRLinalgPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRLinalg
   MLIRPass
   MLIRLinalgTransforms
 )
 
 add_mlir_public_c_api_library(MLIRCAPISCF
   SCF.cpp
 
+  PARTIAL_SOURCES_INTENDED
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRSCF
 )
 
 add_mlir_public_c_api_library(MLIRCAPIShape
   Shape.cpp
 
+  PARTIAL_SOURCES_INTENDED
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRShape
 )
 
+add_mlir_public_c_api_library(MLIRCAPISparseTensor
+  SparseTensor.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRSparseTensor
+)
+
 add_mlir_public_c_api_library(MLIRCAPIStandard
   Standard.cpp
 
+  PARTIAL_SOURCES_INTENDED
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRStandard
 )
 
 add_mlir_public_c_api_library(MLIRCAPITensor
   Tensor.cpp
 
+  PARTIAL_SOURCES_INTENDED
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRTensor
 )

mlir/lib/CAPI/Dialect/SparseTensor.cpp

+//===- Tensor.cpp - C API for SparseTensor dialect ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/SparseTensor.h"
+#include "mlir-c/IR.h"
+#include "mlir/CAPI/AffineMap.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Support/LLVM.h"
+
+using namespace llvm;
+using namespace mlir::sparse_tensor;
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor,
+                                      mlir::sparse_tensor::SparseTensorDialect)
+
+// Ensure the C-API enums are int-castable to C++ equivalents.
+static_assert(
+    static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) ==
+            static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Dense) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::Compressed) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) ==
+            static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Singleton),
+    "MlirSparseTensorDimLevelType (C-API) and DimLevelType (C++) mismatch");
+
+bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) {
+  return unwrap(attr).isa<SparseTensorEncodingAttr>();
+}
+
+MlirAttribute mlirSparseTensorEncodingAttrGet(
+    MlirContext ctx, intptr_t numDimLevelTypes,
+    MlirSparseTensorDimLevelType const *dimLevelTypes,
+    MlirAffineMap dimOrdering, int pointerBitWidth, int indexBitWidth) {
+  SmallVector<SparseTensorEncodingAttr::DimLevelType> cppDimLevelTypes;
+  cppDimLevelTypes.resize(numDimLevelTypes);
+  for (intptr_t i = 0; i < numDimLevelTypes; ++i)
+    cppDimLevelTypes[i] =
+        static_cast<SparseTensorEncodingAttr::DimLevelType>(dimLevelTypes[i]);
+  return wrap(SparseTensorEncodingAttr::get(unwrap(ctx), cppDimLevelTypes,
+                                            unwrap(dimOrdering),
+                                            pointerBitWidth, indexBitWidth));
+}
+
+MlirAffineMap mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr) {
+  return wrap(unwrap(attr).cast<SparseTensorEncodingAttr>().getDimOrdering());
+}
+
+intptr_t mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr) {
+  return unwrap(attr).cast<SparseTensorEncodingAttr>().getDimLevelType().size();
+}
+
+MlirSparseTensorDimLevelType
+mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t pos) {
+  return static_cast<MlirSparseTensorDimLevelType>(
+      unwrap(attr).cast<SparseTensorEncodingAttr>().getDimLevelType()[pos]);
+}
+
+int mlirSparseTensorEncodingAttrGetPointerBitWidth(MlirAttribute attr) {
+  return unwrap(attr).cast<SparseTensorEncodingAttr>().getPointerBitWidth();
+}
+
+int mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr) {
+  return unwrap(attr).cast<SparseTensorEncodingAttr>().getIndexBitWidth();
+}

mlir/test/CAPI/CMakeLists.txt

-set(LLVM_OPTIONAL_SOURCES
-  execution_engine.c
-  ir.c
-  pass.c
-)
-
+function(_add_capi_test_executable name)
+  cmake_parse_arguments(ARG
+    ""
+    ""
+    ""
+    ${ARGN})
   set(LLVM_LINK_COMPONENTS
     )
-
-add_llvm_executable(mlir-capi-ir-test
-  ir.c
-  )
-llvm_update_compile_flags(mlir-capi-ir-test)
-
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-target_link_libraries(mlir-capi-ir-test
+  add_llvm_executable(${name}
+    PARTIAL_SOURCES_INTENDED
+    ${ARG_UNPARSED_ARGUMENTS})
+  llvm_update_compile_flags(${name})
+  target_link_libraries(${name}
     PRIVATE
-  MLIRPublicAPI
-  )
-
+    MLIRPublicAPI)
+endfunction(_add_capi_test_executable)
 
-add_llvm_executable(mlir-capi-pass-test
-  pass.c
-  )
-llvm_update_compile_flags(mlir-capi-pass-test)
+_add_capi_test_executable(mlir-capi-execution-engine-test
+  execution_engine.c
+DEPENDS
+  MLIRConversionPassIncGen
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-target_link_libraries(mlir-capi-pass-test
-  PRIVATE
-  MLIRPublicAPI
-  )
+_add_capi_test_executable(mlir-capi-ir-test
+  ir.c
 
-add_llvm_executable(mlir-capi-execution-engine-test
-  execution_engine.c
-DEPENDS
-  MLIRConversionPassIncGen
-  )
-llvm_update_compile_flags(mlir-capi-execution-engine-test)
+_add_capi_test_executable(mlir-capi-pass-test
+  pass.c
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-target_link_libraries(mlir-capi-execution-engine-test
-  PRIVATE
-  MLIRPublicAPI
-  )
+_add_capi_test_executable(mlir-capi-sparse-tensor-test
+  sparse_tensor.c

mlir/test/CAPI/sparse_tensor.c

+//===- sparse_tensor.c - Test of sparse_tensor APIs -----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: mlir-capi-sparse-tensor-test 2>&1 | FileCheck %s
+
+#include "mlir-c/Dialect/SparseTensor.h"
+#include "mlir-c/IR.h"
+#include "mlir-c/Registration.h"
+
+#include <assert.h>
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+// CHECK-LABEL: testRoundtripEncoding()
+static int testRoundtripEncoding(MlirContext ctx) {
+  fprintf(stderr, "testRoundtripEncoding()\n");
+  // clang-format off
+  const char *originalAsm =
+    "#sparse_tensor.encoding<{ "
+    "dimLevelType = [ \"dense\", \"compressed\", \"singleton\"], "
+    "dimOrdering = affine_map<(d0, d1, d2) -> (d0, d1, d2)>, "
+    "pointerBitWidth = 32, indexBitWidth = 64 }>";
+  // clang-format on
+  MlirAttribute originalAttr =
+      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(originalAsm));
+  // CHECK: isa: 1
+  fprintf(stderr, "isa: %d\n",
+          mlirAttributeIsASparseTensorEncodingAttr(originalAttr));
+  MlirAffineMap dimOrdering =
+      mlirSparseTensorEncodingAttrGetDimOrdering(originalAttr);
+  // CHECK: (d0, d1, d2) -> (d0, d1, d2)
+  mlirAffineMapDump(dimOrdering);
+  // CHECK: level_type: 0
+  // CHECK: level_type: 1
+  // CHECK: level_type: 2
+  int numLevelTypes = mlirSparseTensorEncodingGetNumDimLevelTypes(originalAttr);
+  enum MlirSparseTensorDimLevelType *levelTypes =
+      alloca(sizeof(enum MlirSparseTensorDimLevelType) * numLevelTypes);
+  for (int i = 0; i < numLevelTypes; ++i) {
+    levelTypes[i] =
+        mlirSparseTensorEncodingAttrGetDimLevelType(originalAttr, i);
+    fprintf(stderr, "level_type: %d\n", levelTypes[i]);
+  }
+  // CHECK: pointer: 32
+  int pointerBitWidth =
+      mlirSparseTensorEncodingAttrGetPointerBitWidth(originalAttr);
+  fprintf(stderr, "pointer: %d\n", pointerBitWidth);
+  // CHECK: index: 64
+  int indexBitWidth =
+      mlirSparseTensorEncodingAttrGetIndexBitWidth(originalAttr);
+  fprintf(stderr, "index: %d\n", indexBitWidth);
+
+  MlirAttribute newAttr = mlirSparseTensorEncodingAttrGet(
+      ctx, numLevelTypes, levelTypes, dimOrdering, pointerBitWidth,
+      indexBitWidth);
+  mlirAttributeDump(newAttr); // For debugging filecheck output.
+  // CHECK: equal: 1
+  fprintf(stderr, "equal: %d\n", mlirAttributeEqual(originalAttr, newAttr));
+  return 0;
+}
+
+int main() {
+  MlirContext ctx = mlirContextCreate();
+  mlirDialectHandleRegisterDialect(mlirGetDialectHandle__sparse_tensor__(),
+                                   ctx);
+  if (testRoundtripEncoding(ctx))
+    return 1;
+
+  mlirContextDestroy(ctx);
+  return 0;
+}

mlir/test/CMakeLists.txt

configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/Unit/lit.cfg.py
   )
 
 set(MLIR_TEST_DEPENDS
   FileCheck count not
   mlir-capi-execution-engine-test
   mlir-capi-ir-test
   mlir-capi-pass-test
+  mlir-capi-sparse-tensor-test
   mlir-cpu-runner
   mlir-edsc-builder-api-test
   mlir-linalg-ods-gen
   mlir-lsp-server
   mlir-opt
   mlir-reduce
   mlir-sdbm-api-test
   mlir-tblgen