1 //===- sparse_tensor.c - Test of sparse_tensor APIs -----------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM
4 // Exceptions.
5 // See https://llvm.org/LICENSE.txt for license information.
6 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7 //
8 //===----------------------------------------------------------------------===//
9
10 // RUN: mlir-capi-sparse-tensor-test 2>&1 | FileCheck %s
11
12 #include "mlir-c/Dialect/SparseTensor.h"
13 #include "mlir-c/IR.h"
14 #include "mlir-c/RegisterEverything.h"
15
16 #include <assert.h>
17 #include <inttypes.h>
18 #include <math.h>
19 #include <stdio.h>
20 #include <stdlib.h>
21 #include <string.h>
22
23 // CHECK-LABEL: testRoundtripEncoding()
testRoundtripEncoding(MlirContext ctx)24 static int testRoundtripEncoding(MlirContext ctx) {
25 fprintf(stderr, "testRoundtripEncoding()\n");
26 // clang-format off
27 const char *originalAsm =
28 "#sparse_tensor.encoding<{ "
29 "map = [s0](d0, d1) -> (s0 : dense, d0 : compressed, d1 : compressed), "
30 "posWidth = 32, crdWidth = 64, explicitVal = 1 : i64}>";
31 // clang-format on
32 MlirAttribute originalAttr =
33 mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(originalAsm));
34 // CHECK: isa: 1
35 fprintf(stderr, "isa: %d\n",
36 mlirAttributeIsASparseTensorEncodingAttr(originalAttr));
37 MlirAffineMap dimToLvl =
38 mlirSparseTensorEncodingAttrGetDimToLvl(originalAttr);
39 // CHECK: (d0, d1)[s0] -> (s0, d0, d1)
40 mlirAffineMapDump(dimToLvl);
41 // CHECK: level_type: 65536
42 // CHECK: level_type: 262144
43 // CHECK: level_type: 262144
44 MlirAffineMap lvlToDim =
45 mlirSparseTensorEncodingAttrGetLvlToDim(originalAttr);
46 int lvlRank = mlirSparseTensorEncodingGetLvlRank(originalAttr);
47 MlirSparseTensorLevelType *lvlTypes =
48 malloc(sizeof(MlirSparseTensorLevelType) * lvlRank);
49 for (int l = 0; l < lvlRank; ++l) {
50 lvlTypes[l] = mlirSparseTensorEncodingAttrGetLvlType(originalAttr, l);
51 fprintf(stderr, "level_type: %" PRIu64 "\n", lvlTypes[l]);
52 }
53 // CHECK: posWidth: 32
54 int posWidth = mlirSparseTensorEncodingAttrGetPosWidth(originalAttr);
55 fprintf(stderr, "posWidth: %d\n", posWidth);
56 // CHECK: crdWidth: 64
57 int crdWidth = mlirSparseTensorEncodingAttrGetCrdWidth(originalAttr);
58 fprintf(stderr, "crdWidth: %d\n", crdWidth);
59
60 // CHECK: explicitVal: 1 : i64
61 MlirAttribute explicitVal =
62 mlirSparseTensorEncodingAttrGetExplicitVal(originalAttr);
63 fprintf(stderr, "explicitVal: ");
64 mlirAttributeDump(explicitVal);
65 // CHECK: implicitVal: <<NULL ATTRIBUTE>>
66 MlirAttribute implicitVal =
67 mlirSparseTensorEncodingAttrGetImplicitVal(originalAttr);
68 fprintf(stderr, "implicitVal: ");
69 mlirAttributeDump(implicitVal);
70
71 MlirAttribute newAttr = mlirSparseTensorEncodingAttrGet(
72 ctx, lvlRank, lvlTypes, dimToLvl, lvlToDim, posWidth, crdWidth,
73 explicitVal, implicitVal);
74 mlirAttributeDump(newAttr); // For debugging filecheck output.
75 // CHECK: equal: 1
76 fprintf(stderr, "equal: %d\n", mlirAttributeEqual(originalAttr, newAttr));
77 free(lvlTypes);
78 return 0;
79 }
80
main(void)81 int main(void) {
82 MlirContext ctx = mlirContextCreate();
83 mlirDialectHandleRegisterDialect(mlirGetDialectHandle__sparse_tensor__(),
84 ctx);
85 if (testRoundtripEncoding(ctx))
86 return 1;
87
88 mlirContextDestroy(ctx);
89 return 0;
90 }
91