forked from OSchip/llvm-project
82 lines
2.8 KiB
C
82 lines
2.8 KiB
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 =
|
|
malloc(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));
|
|
|
|
free(levelTypes);
|
|
return 0;
|
|
}
|
|
|
|
int main() {
|
|
MlirContext ctx = mlirContextCreate();
|
|
mlirDialectHandleRegisterDialect(mlirGetDialectHandle__sparse_tensor__(),
|
|
ctx);
|
|
if (testRoundtripEncoding(ctx))
|
|
return 1;
|
|
|
|
mlirContextDestroy(ctx);
|
|
return 0;
|
|
}
|