41 lines · plain
1// RUN: mlir-opt %s -sparse-storage-specifier-to-llvm --cse --canonicalize | FileCheck %s2 3#CSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}>4 5// CHECK-LABEL: func.func @sparse_metadata_init() -> !llvm.struct<(array<2 x i64>, array<3 x i64>)> {6// CHECK-DAG: %[[STRUCT:.*]] = llvm.mlir.poison : !llvm.struct<(array<2 x i64>, array<3 x i64>)>7// CHECK-DAG: %[[CST0:.*]] = arith.constant 0 : i648// CHECK: %[[VAL_1:.*]] = llvm.insertvalue %[[CST0]], %[[STRUCT]][1, 0] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>9// CHECK: %[[VAL_2:.*]] = llvm.insertvalue %[[CST0]], %[[VAL_1]][1, 1] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>10// CHECK: %[[VAL_3:.*]] = llvm.insertvalue %[[CST0]], %[[VAL_2]][1, 2] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>11// CHECK: return %[[VAL_3]] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>12// CHECK: }13func.func @sparse_metadata_init() -> !sparse_tensor.storage_specifier<#CSR> {14 %0 = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#CSR>15 return %0 : !sparse_tensor.storage_specifier<#CSR>16}17 18// CHECK-LABEL: func.func @sparse_get_md(19// CHECK-SAME: %[[VAL_0:.*]]: !llvm.struct<(array<2 x i64>, array<3 x i64>)>) -> index {20// CHECK: %[[VAL_1:.*]] = llvm.extractvalue %[[VAL_0]][0, 0] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>21// CHECK: %[[CAST:.*]] = arith.index_cast %[[VAL_1]] : i64 to index22// CHECK: return %[[CAST]] : index23func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#CSR>) -> index {24 %0 = sparse_tensor.storage_specifier.get %arg0 lvl_sz at 025 : !sparse_tensor.storage_specifier<#CSR>26 return %0 : index27}28 29// CHECK-LABEL: func.func @sparse_set_md(30// CHECK-SAME: %[[VAL_0:.*]]: !llvm.struct<(array<2 x i64>, array<3 x i64>)>,31// CHECK-SAME: %[[VAL_1:.*]]: index) -> !llvm.struct<(array<2 x i64>, array<3 x i64>)> {32// CHECK: %[[CAST:.*]] = arith.index_cast %[[VAL_1]] : index to i6433// CHECK: %[[VAL_2:.*]] = llvm.insertvalue %[[CAST]], %[[VAL_0]][0, 0] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>34// CHECK: return %[[VAL_2]] : !llvm.struct<(array<2 x i64>, array<3 x i64>)>35func.func @sparse_set_md(%arg0: !sparse_tensor.storage_specifier<#CSR>, %arg1: index)36 -> !sparse_tensor.storage_specifier<#CSR> {37 %0 = sparse_tensor.storage_specifier.set %arg0 lvl_sz at 0 with %arg138 : !sparse_tensor.storage_specifier<#CSR>39 return %0 : !sparse_tensor.storage_specifier<#CSR>40}41