1// RUN: mlir-opt %s -split-input-file --sparse-reinterpret-map | FileCheck %s 2 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 blocks 14 map = (i, j) -> 15 ( i floordiv 2 : dense 16 , j floordiv 4 : compressed 17 , i mod 2 : dense 18 , j mod 4 : dense 19 ) 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_mul 37 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 : f32 41 %2 = arith.mulf %1, %z : f32 42 linalg.yield %2 : f32 43 } -> 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 : dense 55 ) 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 : dense 93 ) 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), %val 108 : (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: return 120// 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), index 130 return %rd, %rp, %ri : tensor<?xf64>, tensor<?xindex>, tensor<?xindex> 131} 132