127 lines · plain
1// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification | FileCheck %s2 3#DenseMatrix = #sparse_tensor.encoding<{4 map = (d0, d1) -> (d0 : dense, d1 : dense)5}>6 7#SparseMatrix = #sparse_tensor.encoding<{8 map = (d0, d1) -> (d0 : compressed, d1 : compressed)9}>10 11#trait = {12 indexing_maps = [13 affine_map<(i,j) -> (i,j)>, // A14 affine_map<(i,j) -> (i,j)> // X (out)15 ],16 iterator_types = ["parallel", "parallel"],17 doc = "X(i,j) = A(i,j) * i * j"18}19 20// CHECK-LABEL: func.func @dense_index(21// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xi64, #sparse{{[0-9]*}}>22// CHECK-DAG: %[[VAL_1:.*]] = arith.constant 0 : index23// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index24// CHECK-DAG: %[[VAL_3:.*]] = sparse_tensor.lvl %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse{{[0-9]*}}>25// CHECK-DAG: %[[VAL_4:.*]] = sparse_tensor.lvl %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse{{[0-9]*}}>26// CHECK-DAG: %[[VAL_5:.*]] = tensor.empty(%[[VAL_3]], %[[VAL_4]]) : tensor<?x?xi64, #sparse{{[0-9]*}}>27// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xi64, #sparse{{[0-9]*}}>28// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.lvl %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse{{[0-9]*}}>29// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.lvl %[[VAL_0]], %[[VAL_2]] : tensor<?x?xi64, #sparse{{[0-9]*}}>30// CHECK-DAG: %[[VAL_24:.*]] = sparse_tensor.lvl %[[VAL_5]], %[[VAL_2]] : tensor<?x?xi64, #sparse{{[0-9]*}}>31// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_5]] : tensor<?x?xi64, #sparse{{[0-9]*}}>32// CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_1]] to %[[VAL_7]] step %[[VAL_2]] {33// CHECK: %[[VAL_12:.*]] = arith.muli %[[VAL_10]], %[[VAL_8]] : index34// CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_10]], %[[VAL_24]] : index35// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_1]] to %[[VAL_8]] step %[[VAL_2]] {36// CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_11]], %[[VAL_12]] : index37// CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_11]], %[[VAL_14]] : index38// CHECK: %[[VAL_16:.*]] = arith.index_cast %[[VAL_11]] : index to i6439// CHECK: %[[VAL_17:.*]] = arith.index_cast %[[VAL_10]] : index to i6440// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xi64>41// CHECK: %[[VAL_19:.*]] = arith.muli %[[VAL_17]], %[[VAL_18]] : i6442// CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_16]], %[[VAL_19]] : i6443// CHECK: memref.store %[[VAL_20]], %[[VAL_9]]{{\[}}%[[VAL_15]]] : memref<?xi64>44// CHECK: }45// CHECK: }46// CHECK: %[[VAL_21:.*]] = sparse_tensor.load %[[VAL_5]] : tensor<?x?xi64, #sparse{{[0-9]*}}>47// CHECK: return %[[VAL_21]] : tensor<?x?xi64, #sparse{{[0-9]*}}>48// CHECK: }49func.func @dense_index(%arga: tensor<?x?xi64, #DenseMatrix>)50 -> tensor<?x?xi64, #DenseMatrix> {51 %c0 = arith.constant 0 : index52 %c1 = arith.constant 0 : index53 %0 = sparse_tensor.lvl %arga, %c0 : tensor<?x?xi64, #DenseMatrix>54 %1 = sparse_tensor.lvl %arga, %c1 : tensor<?x?xi64, #DenseMatrix>55 %init = tensor.empty(%0, %1) : tensor<?x?xi64, #DenseMatrix>56 %r = linalg.generic #trait57 ins(%arga: tensor<?x?xi64, #DenseMatrix>)58 outs(%init: tensor<?x?xi64, #DenseMatrix>) {59 ^bb(%a: i64, %x: i64):60 %i = linalg.index 0 : index61 %j = linalg.index 1 : index62 %ii = arith.index_cast %i : index to i6463 %jj = arith.index_cast %j : index to i6464 %m1 = arith.muli %ii, %a : i6465 %m2 = arith.muli %jj, %m1 : i6466 linalg.yield %m2 : i6467 } -> tensor<?x?xi64, #DenseMatrix>68 return %r : tensor<?x?xi64, #DenseMatrix>69}70 71 72// CHECK-LABEL: func.func @sparse_index(73// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xi64, #sparse{{[0-9]*}}>74// CHECK-DAG: %[[VAL_1:.*]] = arith.constant 0 : index75// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index76// CHECK-DAG: %[[VAL_3:.*]] = sparse_tensor.lvl %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse{{[0-9]*}}>77// CHECK-DAG: %[[VAL_4:.*]] = sparse_tensor.lvl %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse{{[0-9]*}}>78// CHECK-DAG: %[[VAL_5:.*]] = tensor.empty(%[[VAL_3]], %[[VAL_4]]) : tensor<?x?xi64, #sparse{{[0-9]*}}>79// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<?x?xi64, #sparse{{[0-9]*}}>80// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<?x?xi64, #sparse{{[0-9]*}}>81// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<?x?xi64, #sparse{{[0-9]*}}>82// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 1 : index} : tensor<?x?xi64, #sparse{{[0-9]*}}>83// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xi64, #sparse{{[0-9]*}}>84// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_1]]] : memref<?xindex>85// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_2]]] : memref<?xindex>86// CHECK: %[[T:.*]] = scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_2]] {{.*}} {87// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>88// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<?xindex>89// CHECK: %[[VAL_16:.*]] = arith.addi %[[VAL_13]], %[[VAL_2]] : index90// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xindex>91// CHECK: %[[L:.*]] = scf.for %[[VAL_18:.*]] = %[[VAL_15]] to %[[VAL_17]] step %[[VAL_2]] {{.*}} {92// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xindex>93// CHECK: %[[VAL_20:.*]] = arith.index_cast %[[VAL_19]] : index to i6494// CHECK: %[[VAL_21:.*]] = arith.index_cast %[[VAL_14]] : index to i6495// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_18]]] : memref<?xi64>96// CHECK: %[[VAL_23:.*]] = arith.muli %[[VAL_21]], %[[VAL_22]] : i6497// CHECK: %[[VAL_24:.*]] = arith.muli %[[VAL_20]], %[[VAL_23]] : i6498// CHECK: %[[Y:.*]] = tensor.insert %[[VAL_24]] into %{{.*}}[%[[VAL_14]], %[[VAL_19]]] : tensor<?x?xi64, #sparse{{[0-9]*}}>99// CHECK: scf.yield %[[Y]]100// CHECK: }101// CHECK: scf.yield %[[L]]102// CHECK: }103// CHECK: %[[VAL_25:.*]] = sparse_tensor.load %[[T]] hasInserts : tensor<?x?xi64, #sparse{{[0-9]*}}>104// CHECK: return %[[VAL_25]] : tensor<?x?xi64, #sparse{{[0-9]*}}>105// CHECK: }106func.func @sparse_index(%arga: tensor<?x?xi64, #SparseMatrix>)107 -> tensor<?x?xi64, #SparseMatrix> {108 %c0 = arith.constant 0 : index109 %c1 = arith.constant 0 : index110 %0 = sparse_tensor.lvl %arga, %c0 : tensor<?x?xi64, #SparseMatrix>111 %1 = sparse_tensor.lvl %arga, %c1 : tensor<?x?xi64, #SparseMatrix>112 %init = tensor.empty(%0, %1) : tensor<?x?xi64, #SparseMatrix>113 %r = linalg.generic #trait114 ins(%arga: tensor<?x?xi64, #SparseMatrix>)115 outs(%init: tensor<?x?xi64, #SparseMatrix>) {116 ^bb(%a: i64, %x: i64):117 %i = linalg.index 0 : index118 %j = linalg.index 1 : index119 %ii = arith.index_cast %i : index to i64120 %jj = arith.index_cast %j : index to i64121 %m1 = arith.muli %ii, %a : i64122 %m2 = arith.muli %jj, %m1 : i64123 linalg.yield %m2 : i64124 } -> tensor<?x?xi64, #SparseMatrix>125 return %r : tensor<?x?xi64, #SparseMatrix>126}127