132 lines · plain
1// RUN: mlir-opt %s -split-input-file --sparse-reinterpret-map | FileCheck %s2 3#trait_mul = {4 indexing_maps = [5 affine_map<(i,j) -> (i,j)>, // A (in)6 affine_map<(i,j) -> (j,i)>, // B (in, transposed)7 affine_map<(i,j) -> (i,j)> // X (out)8 ],9 iterator_types = ["parallel", "parallel"],10 doc = "X(i,j) *= A(i,j) * B(j,i)"11}12 13#BSR = #sparse_tensor.encoding<{ // 2x4 blocks14 map = (i, j) ->15 ( i floordiv 2 : dense16 , j floordiv 4 : compressed17 , i mod 2 : dense18 , j mod 4 : dense19 )20}>21 22// CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0 * 2 + d2, d1 * 4 + d3)>23// CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2, d3) -> (d1 * 4 + d3, d0 * 2 + d2)>24// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>25// CHECK-LABEL: func @mul(26// CHECK-SAME: %[[A0:.*0]]: tensor<32x32xf32>,27// CHECK-SAME: %[[A1:.*1]]: tensor<32x32xf32>,28// CHECK-SAME: %[[A2:.*2]]: tensor<32x32xf32, #sparse{{[0-9]*}}>)29// CHECK: %[[T0:.*]] = sparse_tensor.reinterpret_map %[[A2]]30// CHECK: %[[T1:.*]] = linalg.generic {doc = {{.*}} indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}31// CHECK: %[[T2:.*]] = sparse_tensor.reinterpret_map %[[T1]]32// CHECK: return %[[T2]] : tensor<32x32xf32, #sparse{{[0-9]*}}>33func.func @mul(%arg0: tensor<32x32xf32>,34 %arg1: tensor<32x32xf32>,35 %arg2: tensor<32x32xf32, #BSR>) -> tensor<32x32xf32, #BSR> {36 %0 = linalg.generic #trait_mul37 ins(%arg0, %arg1: tensor<32x32xf32>, tensor<32x32xf32>)38 outs(%arg2: tensor<32x32xf32, #BSR>) {39 ^bb(%x: f32, %y : f32, %z : f32):40 %1 = arith.mulf %x, %y : f3241 %2 = arith.mulf %1, %z : f3242 linalg.yield %2 : f3243 } -> tensor<32x32xf32, #BSR>44 return %0 : tensor<32x32xf32, #BSR>45}46 47// -----48 49#BSR = #sparse_tensor.encoding<{50 map = ( i, j ) ->51 ( i floordiv 2 : dense,52 j floordiv 2 : compressed,53 i mod 2 : dense,54 j mod 2 : dense55 )56}>57 58// CHECK-DAG: #[[$remap:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 2 : compressed, d0 mod 2 : dense, d1 mod 2 : dense) }>59// CHECK-DAG: #[[$demap:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2, d3) -> (d0 : dense, d1 : compressed, d2 : dense, d3 : dense) }>60// CHECK-LABEL: func.func @sparse_foreach_reinterpret_map(61// CHECK-SAME: %[[VAL_0:.*]]: tensor<2x4xf64, #[[$remap]]>62// CHECK: %[[VAL_1:.*]] = bufferization.alloc_tensor() : tensor<1x2x2x2xf64, #[[$demap]]>63// CHECK: %[[VAL_2:.*]] = sparse_tensor.reinterpret_map %[[VAL_0]] : tensor<2x4xf64, #[[$remap]]> to tensor<1x2x2x2xf64, #[[$demap]]>64// CHECK: %[[VAL_4:.*]] = sparse_tensor.foreach in %[[VAL_2]] init(%[[VAL_1]])65// CHECK: ^bb0(%[[VAL_5:.*]]: index, %[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index, %[[VAL_8:.*]]: index, %[[VAL_9:.*]]: f64, %[[VAL_10:.*]]: tensor<1x2x2x2xf64, #[[$demap]]>66// CHECK: %[[VAL_11:.*]] = tensor.insert %[[VAL_9]] into %[[VAL_10]]{{\[}}%[[VAL_5]], %[[VAL_6]], %[[VAL_7]], %[[VAL_8]]] : tensor<1x2x2x2xf64, #[[$demap]]>67// CHECK: sparse_tensor.yield %[[VAL_11]] : tensor<1x2x2x2xf64, #sparse{{[0-9]*}}>68// CHECK: }69// CHECK: %[[VAL_12:.*]] = sparse_tensor.reinterpret_map %[[VAL_4]] : tensor<1x2x2x2xf64, #[[$demap]]> to tensor<2x4xf64, #[[$remap]]>70// CHECK: %[[VAL_13:.*]] = sparse_tensor.load %[[VAL_12]] hasInserts : tensor<2x4xf64, #[[$remap]]>71// CHECK: return %[[VAL_13]] : tensor<2x4xf64, #sparse{{[0-9]*}}>72// CHECK: }73func.func @sparse_foreach_reinterpret_map(%6 : tensor<2x4xf64, #BSR>) -> tensor<2x4xf64, #BSR> {74 %7 = bufferization.alloc_tensor() : tensor<2x4xf64, #BSR>75 %8 = sparse_tensor.foreach in %6 init(%7) : tensor<2x4xf64, #BSR>, tensor<2x4xf64, #BSR> -> tensor<2x4xf64, #BSR> do {76 ^bb0(%arg0: index, %arg1: index, %arg2: f64, %arg3: tensor<2x4xf64, #BSR>):77 %inserted = tensor.insert %arg2 into %arg3[%arg0, %arg1] : tensor<2x4xf64, #BSR>78 sparse_tensor.yield %inserted : tensor<2x4xf64, #BSR>79 }80 %9 = sparse_tensor.load %8 hasInserts : tensor<2x4xf64, #BSR>81 return %9 : tensor<2x4xf64, #BSR>82}83 84 85// -----86 87#BSR = #sparse_tensor.encoding<{88 map = ( i, j ) ->89 ( i floordiv 2 : dense,90 j floordiv 2 : compressed,91 i mod 2 : dense,92 j mod 2 : dense93 )94}>95// CHECK-DAG: #[[$remap:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 2 : compressed, d0 mod 2 : dense, d1 mod 2 : dense) }>96// CHECK-DAG: #[[$demap:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2, d3) -> (d0 : dense, d1 : compressed, d2 : dense, d3 : dense) }>97 98// CHECK-LABEL: func.func @sparse_assemble_reinterpret_map(99// CHECK-SAME: %[[VAL_0:.*]]: tensor<?xf64>,100// CHECK-SAME: %[[VAL_1:.*]]: tensor<?xindex>,101// CHECK-SAME: %[[VAL_2:.*]]: tensor<?xindex>) -> tensor<2x4xf64, #[[$remap]]> {102// CHECK: %[[VAL_3:.*]] = sparse_tensor.assemble {{.*}} to tensor<1x2x2x2xf64, #[[$demap]]>103// CHECK: %[[VAL_4:.*]] = sparse_tensor.reinterpret_map %[[VAL_3]] : tensor<1x2x2x2xf64, #[[$demap]]> to tensor<2x4xf64, #[[$remap]]>104// CHECK: return %[[VAL_4]] : tensor<2x4xf64, #[[$remap]]>105// CHECK: }106func.func @sparse_assemble_reinterpret_map(%val : tensor<?xf64>, %pos:tensor<?xindex>, %crd:tensor<?xindex>) -> tensor<2x4xf64, #BSR> {107 %0 = sparse_tensor.assemble (%pos, %crd), %val108 : (tensor<?xindex>, tensor<?xindex>), tensor<?xf64> to tensor<2x4xf64, #BSR>109 return %0 : tensor<2x4xf64, #BSR>110}111 112// CHECK-LABEL: func.func @sparse_disassemble_reinterpret_map(113// CHECK-SAME: %[[VAL_0:.*]]: tensor<2x4xf64, #[[$remap]]>,114// CHECK-SAME: %[[VAL_1:.*]]: tensor<?xf64>,115// CHECK-SAME: %[[VAL_2:.*]]: tensor<?xindex>,116// CHECK-SAME: %[[VAL_3:.*]]: tensor<?xindex>) -> (tensor<?xf64>, tensor<?xindex>, tensor<?xindex>) {117// CHECK: %[[VAL_4:.*]] = sparse_tensor.reinterpret_map %[[VAL_0]] : tensor<2x4xf64, #[[$remap]]> to tensor<1x2x2x2xf64, #[[$demap]]>118// CHECK: %{{.*}} = sparse_tensor.disassemble %[[VAL_4]] : tensor<1x2x2x2xf64, #[[$demap]]>119// CHECK: return120// CHECK: }121func.func @sparse_disassemble_reinterpret_map(%sp : tensor<2x4xf64, #BSR>,122 %od : tensor<?xf64>,123 %op : tensor<?xindex>,124 %oi : tensor<?xindex>)125 -> (tensor<?xf64>, tensor<?xindex>, tensor<?xindex>) {126 %rp, %ri, %rd, %dl, %pl, %il = sparse_tensor.disassemble %sp : tensor<2x4xf64, #BSR>127 out_lvls(%op, %oi : tensor<?xindex>, tensor<?xindex>)128 out_vals(%od : tensor<?xf64>)129 -> (tensor<?xindex>, tensor<?xindex>), tensor<?xf64>, (index, index), index130 return %rd, %rp, %ri : tensor<?xf64>, tensor<?xindex>, tensor<?xindex>131}132