xref: /llvm-project/mlir/test/Dialect/SparseTensor/specifier_to_llvm.mlir (revision 4bce270157f9a81bd7e38dc589a2970a445d1e96)
1083ddffeSPeiming Liu// RUN: mlir-opt %s -sparse-storage-specifier-to-llvm --cse --canonicalize | FileCheck %s
2083ddffeSPeiming Liu
3e2e429d9SYinying Li#CSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}>
4083ddffeSPeiming Liu
5083ddffeSPeiming Liu// CHECK-LABEL:   func.func @sparse_metadata_init() -> !llvm.struct<(array<2 x i64>, array<3 x i64>)> {
6*4bce2701SGuy David// CHECK-DAG:       %[[STRUCT:.*]] = llvm.mlir.undef : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
7*4bce2701SGuy David// CHECK-DAG:       %[[CST0:.*]] = arith.constant 0 : i64
8*4bce2701SGuy David// CHECK:           %[[VAL_1:.*]] = llvm.insertvalue %[[CST0]], %[[STRUCT]][1, 0] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
9*4bce2701SGuy David// CHECK:           %[[VAL_2:.*]] = llvm.insertvalue %[[CST0]], %[[VAL_1]][1, 1] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
10*4bce2701SGuy David// CHECK:           %[[VAL_3:.*]] = llvm.insertvalue %[[CST0]], %[[VAL_2]][1, 2] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
11*4bce2701SGuy David// CHECK:           return %[[VAL_3]] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
12083ddffeSPeiming Liu// CHECK:         }
13083ddffeSPeiming Liufunc.func @sparse_metadata_init() -> !sparse_tensor.storage_specifier<#CSR> {
14083ddffeSPeiming Liu  %0 = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#CSR>
15083ddffeSPeiming Liu  return %0 : !sparse_tensor.storage_specifier<#CSR>
16083ddffeSPeiming Liu}
17083ddffeSPeiming Liu
18083ddffeSPeiming Liu// CHECK-LABEL:   func.func @sparse_get_md(
1944ff23d5SPeiming Liu// CHECK-SAME:      %[[VAL_0:.*]]: !llvm.struct<(array<2 x i64>, array<3 x i64>)>) -> index {
20083ddffeSPeiming Liu// CHECK:           %[[VAL_1:.*]] = llvm.extractvalue %[[VAL_0]][0, 0] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
2144ff23d5SPeiming Liu// CHECK:           %[[CAST:.*]] = arith.index_cast %[[VAL_1]] : i64 to index
2244ff23d5SPeiming Liu// CHECK:           return %[[CAST]] : index
2344ff23d5SPeiming Liufunc.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#CSR>) -> index {
2484cd51bbSwren romano  %0 = sparse_tensor.storage_specifier.get %arg0 lvl_sz at 0
2544ff23d5SPeiming Liu       : !sparse_tensor.storage_specifier<#CSR>
2644ff23d5SPeiming Liu  return %0 : index
27083ddffeSPeiming Liu}
28083ddffeSPeiming Liu
29083ddffeSPeiming Liu// CHECK-LABEL:   func.func @sparse_set_md(
30083ddffeSPeiming Liu// CHECK-SAME:      %[[VAL_0:.*]]: !llvm.struct<(array<2 x i64>, array<3 x i64>)>,
3144ff23d5SPeiming Liu// CHECK-SAME:      %[[VAL_1:.*]]: index) -> !llvm.struct<(array<2 x i64>, array<3 x i64>)> {
3244ff23d5SPeiming Liu// CHECK:           %[[CAST:.*]] = arith.index_cast %[[VAL_1]] : index to i64
3344ff23d5SPeiming Liu// CHECK:           %[[VAL_2:.*]] = llvm.insertvalue %[[CAST]], %[[VAL_0]][0, 0] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
34083ddffeSPeiming Liu// CHECK:           return %[[VAL_2]] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>
3544ff23d5SPeiming Liufunc.func @sparse_set_md(%arg0: !sparse_tensor.storage_specifier<#CSR>, %arg1: index)
36083ddffeSPeiming Liu          -> !sparse_tensor.storage_specifier<#CSR> {
3784cd51bbSwren romano  %0 = sparse_tensor.storage_specifier.set %arg0 lvl_sz at 0 with %arg1
3844ff23d5SPeiming Liu       : !sparse_tensor.storage_specifier<#CSR>
39083ddffeSPeiming Liu  return %0 : !sparse_tensor.storage_specifier<#CSR>
40083ddffeSPeiming Liu}
41