d4088e7d5f
Updates: 1. Infer lvlToDim from dimToLvl 2. Add more tests for block sparsity 3. Finish TODOs related to lvlToDim, including adding lvlToDim to python binding Verification of lvlToDim that user provides will be implemented in the next PR.
77 lines
2.7 KiB
C
77 lines
2.7 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/RegisterEverything.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<{ "
|
|
"map = [s0](d0, d1) -> (s0 : dense, d0 : compressed, d1 : compressed), "
|
|
"posWidth = 32, crdWidth = 64 }>";
|
|
// clang-format on
|
|
MlirAttribute originalAttr =
|
|
mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(originalAsm));
|
|
// CHECK: isa: 1
|
|
fprintf(stderr, "isa: %d\n",
|
|
mlirAttributeIsASparseTensorEncodingAttr(originalAttr));
|
|
MlirAffineMap dimToLvl =
|
|
mlirSparseTensorEncodingAttrGetDimToLvl(originalAttr);
|
|
// CHECK: (d0, d1)[s0] -> (s0, d0, d1)
|
|
mlirAffineMapDump(dimToLvl);
|
|
// CHECK: level_type: 4
|
|
// CHECK: level_type: 8
|
|
// CHECK: level_type: 8
|
|
MlirAffineMap lvlToDim =
|
|
mlirSparseTensorEncodingAttrGetLvlToDim(originalAttr);
|
|
int lvlRank = mlirSparseTensorEncodingGetLvlRank(originalAttr);
|
|
enum MlirSparseTensorDimLevelType *lvlTypes =
|
|
malloc(sizeof(enum MlirSparseTensorDimLevelType) * lvlRank);
|
|
for (int l = 0; l < lvlRank; ++l) {
|
|
lvlTypes[l] = mlirSparseTensorEncodingAttrGetLvlType(originalAttr, l);
|
|
fprintf(stderr, "level_type: %d\n", lvlTypes[l]);
|
|
}
|
|
// CHECK: posWidth: 32
|
|
int posWidth = mlirSparseTensorEncodingAttrGetPosWidth(originalAttr);
|
|
fprintf(stderr, "posWidth: %d\n", posWidth);
|
|
// CHECK: crdWidth: 64
|
|
int crdWidth = mlirSparseTensorEncodingAttrGetCrdWidth(originalAttr);
|
|
fprintf(stderr, "crdWidth: %d\n", crdWidth);
|
|
MlirAttribute newAttr = mlirSparseTensorEncodingAttrGet(
|
|
ctx, lvlRank, lvlTypes, dimToLvl, lvlToDim, posWidth, crdWidth);
|
|
mlirAttributeDump(newAttr); // For debugging filecheck output.
|
|
// CHECK: equal: 1
|
|
fprintf(stderr, "equal: %d\n", mlirAttributeEqual(originalAttr, newAttr));
|
|
free(lvlTypes);
|
|
return 0;
|
|
}
|
|
|
|
int main(void) {
|
|
MlirContext ctx = mlirContextCreate();
|
|
mlirDialectHandleRegisterDialect(mlirGetDialectHandle__sparse_tensor__(),
|
|
ctx);
|
|
if (testRoundtripEncoding(ctx))
|
|
return 1;
|
|
|
|
mlirContextDestroy(ctx);
|
|
return 0;
|
|
}
|