| // RUN: mlir-opt %s -split-input-file --sparse-reinterpret-map | FileCheck %s |
| |
| #trait_mul = { |
| indexing_maps = [ |
| affine_map<(i,j) -> (i,j)>, // A (in) |
| affine_map<(i,j) -> (j,i)>, // B (in, transposed) |
| affine_map<(i,j) -> (i,j)> // X (out) |
| ], |
| iterator_types = ["parallel", "parallel"], |
| doc = "X(i,j) *= A(i,j) * B(j,i)" |
| } |
| |
| #BSR = #sparse_tensor.encoding<{ // 2x4 blocks |
| map = (i, j) -> |
| ( i floordiv 2 : dense |
| , j floordiv 4 : compressed |
| , i mod 2 : dense |
| , j mod 4 : dense |
| ) |
| }> |
| |
| // CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0 * 2 + d2, d1 * 4 + d3)> |
| // CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2, d3) -> (d1 * 4 + d3, d0 * 2 + d2)> |
| // CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> |
| // CHECK-LABEL: func @mul( |
| // CHECK-SAME: %[[A0:.*0]]: tensor<32x32xf32>, |
| // CHECK-SAME: %[[A1:.*1]]: tensor<32x32xf32>, |
| // CHECK-SAME: %[[A2:.*2]]: tensor<32x32xf32, #sparse{{[0-9]*}}>) |
| // CHECK: %[[T0:.*]] = sparse_tensor.reinterpret_map %[[A2]] |
| // CHECK: %[[T1:.*]] = linalg.generic {doc = {{.*}} indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} |
| // CHECK: %[[T2:.*]] = sparse_tensor.reinterpret_map %[[T1]] |
| // CHECK: return %[[T2]] : tensor<32x32xf32, #sparse{{[0-9]*}}> |
| func.func @mul(%arg0: tensor<32x32xf32>, |
| %arg1: tensor<32x32xf32>, |
| %arg2: tensor<32x32xf32, #BSR>) -> tensor<32x32xf32, #BSR> { |
| %0 = linalg.generic #trait_mul |
| ins(%arg0, %arg1: tensor<32x32xf32>, tensor<32x32xf32>) |
| outs(%arg2: tensor<32x32xf32, #BSR>) { |
| ^bb(%x: f32, %y : f32, %z : f32): |
| %1 = arith.mulf %x, %y : f32 |
| %2 = arith.mulf %1, %z : f32 |
| linalg.yield %2 : f32 |
| } -> tensor<32x32xf32, #BSR> |
| return %0 : tensor<32x32xf32, #BSR> |
| } |
| |
| // ----- |
| |
| #BSR = #sparse_tensor.encoding<{ |
| map = ( i, j ) -> |
| ( i floordiv 2 : dense, |
| j floordiv 2 : compressed, |
| i mod 2 : dense, |
| j mod 2 : dense |
| ) |
| }> |
| |
| // 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) }> |
| // CHECK-DAG: #[[$demap:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2, d3) -> (d0 : dense, d1 : compressed, d2 : dense, d3 : dense) }> |
| // CHECK-LABEL: func.func @sparse_foreach_reinterpret_map( |
| // CHECK-SAME: %[[VAL_0:.*]]: tensor<2x4xf64, #[[$remap]]> |
| // CHECK: %[[VAL_1:.*]] = bufferization.alloc_tensor() : tensor<1x2x2x2xf64, #[[$demap]]> |
| // CHECK: %[[VAL_2:.*]] = sparse_tensor.reinterpret_map %[[VAL_0]] : tensor<2x4xf64, #[[$remap]]> to tensor<1x2x2x2xf64, #[[$demap]]> |
| // CHECK: %[[VAL_4:.*]] = sparse_tensor.foreach in %[[VAL_2]] init(%[[VAL_1]]) |
| // CHECK: ^bb0(%[[VAL_5:.*]]: index, %[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index, %[[VAL_8:.*]]: index, %[[VAL_9:.*]]: f64, %[[VAL_10:.*]]: tensor<1x2x2x2xf64, #[[$demap]]> |
| // CHECK: %[[VAL_11:.*]] = tensor.insert %[[VAL_9]] into %[[VAL_10]]{{\[}}%[[VAL_5]], %[[VAL_6]], %[[VAL_7]], %[[VAL_8]]] : tensor<1x2x2x2xf64, #[[$demap]]> |
| // CHECK: sparse_tensor.yield %[[VAL_11]] : tensor<1x2x2x2xf64, #sparse{{[0-9]*}}> |
| // CHECK: } |
| // CHECK: %[[VAL_12:.*]] = sparse_tensor.reinterpret_map %[[VAL_4]] : tensor<1x2x2x2xf64, #[[$demap]]> to tensor<2x4xf64, #[[$remap]]> |
| // CHECK: %[[VAL_13:.*]] = sparse_tensor.load %[[VAL_12]] hasInserts : tensor<2x4xf64, #[[$remap]]> |
| // CHECK: return %[[VAL_13]] : tensor<2x4xf64, #sparse{{[0-9]*}}> |
| // CHECK: } |
| func.func @sparse_foreach_reinterpret_map(%6 : tensor<2x4xf64, #BSR>) -> tensor<2x4xf64, #BSR> { |
| %7 = bufferization.alloc_tensor() : tensor<2x4xf64, #BSR> |
| %8 = sparse_tensor.foreach in %6 init(%7) : tensor<2x4xf64, #BSR>, tensor<2x4xf64, #BSR> -> tensor<2x4xf64, #BSR> do { |
| ^bb0(%arg0: index, %arg1: index, %arg2: f64, %arg3: tensor<2x4xf64, #BSR>): |
| %inserted = tensor.insert %arg2 into %arg3[%arg0, %arg1] : tensor<2x4xf64, #BSR> |
| sparse_tensor.yield %inserted : tensor<2x4xf64, #BSR> |
| } |
| %9 = sparse_tensor.load %8 hasInserts : tensor<2x4xf64, #BSR> |
| return %9 : tensor<2x4xf64, #BSR> |
| } |
| |
| |
| // ----- |
| |
| #BSR = #sparse_tensor.encoding<{ |
| map = ( i, j ) -> |
| ( i floordiv 2 : dense, |
| j floordiv 2 : compressed, |
| i mod 2 : dense, |
| j mod 2 : dense |
| ) |
| }> |
| // 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) }> |
| // CHECK-DAG: #[[$demap:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2, d3) -> (d0 : dense, d1 : compressed, d2 : dense, d3 : dense) }> |
| |
| // CHECK-LABEL: func.func @sparse_assemble_reinterpret_map( |
| // CHECK-SAME: %[[VAL_0:.*]]: tensor<?xf64>, |
| // CHECK-SAME: %[[VAL_1:.*]]: tensor<?xindex>, |
| // CHECK-SAME: %[[VAL_2:.*]]: tensor<?xindex>) -> tensor<2x4xf64, #[[$remap]]> { |
| // CHECK: %[[VAL_3:.*]] = sparse_tensor.assemble {{.*}} to tensor<1x2x2x2xf64, #[[$demap]]> |
| // CHECK: %[[VAL_4:.*]] = sparse_tensor.reinterpret_map %[[VAL_3]] : tensor<1x2x2x2xf64, #[[$demap]]> to tensor<2x4xf64, #[[$remap]]> |
| // CHECK: return %[[VAL_4]] : tensor<2x4xf64, #[[$remap]]> |
| // CHECK: } |
| func.func @sparse_assemble_reinterpret_map(%val : tensor<?xf64>, %pos:tensor<?xindex>, %crd:tensor<?xindex>) -> tensor<2x4xf64, #BSR> { |
| %0 = sparse_tensor.assemble (%pos, %crd), %val |
| : (tensor<?xindex>, tensor<?xindex>), tensor<?xf64> to tensor<2x4xf64, #BSR> |
| return %0 : tensor<2x4xf64, #BSR> |
| } |
| |
| // CHECK-LABEL: func.func @sparse_disassemble_reinterpret_map( |
| // CHECK-SAME: %[[VAL_0:.*]]: tensor<2x4xf64, #[[$remap]]>, |
| // CHECK-SAME: %[[VAL_1:.*]]: tensor<?xf64>, |
| // CHECK-SAME: %[[VAL_2:.*]]: tensor<?xindex>, |
| // CHECK-SAME: %[[VAL_3:.*]]: tensor<?xindex>) -> (tensor<?xf64>, tensor<?xindex>, tensor<?xindex>) { |
| // CHECK: %[[VAL_4:.*]] = sparse_tensor.reinterpret_map %[[VAL_0]] : tensor<2x4xf64, #[[$remap]]> to tensor<1x2x2x2xf64, #[[$demap]]> |
| // CHECK: %{{.*}} = sparse_tensor.disassemble %[[VAL_4]] : tensor<1x2x2x2xf64, #[[$demap]]> |
| // CHECK: return |
| // CHECK: } |
| func.func @sparse_disassemble_reinterpret_map(%sp : tensor<2x4xf64, #BSR>, |
| %od : tensor<?xf64>, |
| %op : tensor<?xindex>, |
| %oi : tensor<?xindex>) |
| -> (tensor<?xf64>, tensor<?xindex>, tensor<?xindex>) { |
| %rp, %ri, %rd, %dl, %pl, %il = sparse_tensor.disassemble %sp : tensor<2x4xf64, #BSR> |
| out_lvls(%op, %oi : tensor<?xindex>, tensor<?xindex>) |
| out_vals(%od : tensor<?xf64>) |
| -> (tensor<?xindex>, tensor<?xindex>), tensor<?xf64>, (index, index), index |
| return %rd, %rp, %ri : tensor<?xf64>, tensor<?xindex>, tensor<?xindex> |
| } |
