170 lines · plain
1// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification="parallelization-strategy=none" | \2// RUN: FileCheck %s --check-prefix=CHECK-PAR03// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification="parallelization-strategy=dense-outer-loop" | \4// RUN: FileCheck %s --check-prefix=CHECK-PAR15// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification="parallelization-strategy=any-storage-outer-loop" | \6// RUN: FileCheck %s --check-prefix=CHECK-PAR27// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification="parallelization-strategy=dense-any-loop" | \8// RUN: FileCheck %s --check-prefix=CHECK-PAR39// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification="parallelization-strategy=any-storage-any-loop" | \10// RUN: FileCheck %s --check-prefix=CHECK-PAR411 12#DenseMatrix = #sparse_tensor.encoding<{13 map = (d0, d1) -> (d0 : dense, d1 : dense)14}>15 16#SparseMatrix = #sparse_tensor.encoding<{17 map = (d0, d1) -> (d0 : compressed, d1 : compressed)18}>19 20#CSR = #sparse_tensor.encoding<{21 map = (d0, d1) -> (d0 : dense, d1 : compressed)22}>23 24#trait_dd = {25 indexing_maps = [26 affine_map<(i,j) -> (i,j)>, // A27 affine_map<(i,j) -> (i,j)> // X (out)28 ],29 iterator_types = ["parallel", "parallel"],30 doc = "X(i,j) = A(i,j) * SCALE"31}32 33//34// CHECK-PAR0-LABEL: func @scale_dd35// CHECK-PAR0: scf.for36// CHECK-PAR0: scf.for37// CHECK-PAR0: return38//39// CHECK-PAR1-LABEL: func @scale_dd40// CHECK-PAR1: scf.parallel41// CHECK-PAR1: scf.for42// CHECK-PAR1: return43//44// CHECK-PAR2-LABEL: func @scale_dd45// CHECK-PAR2: scf.parallel46// CHECK-PAR2: scf.for47// CHECK-PAR2: return48//49// CHECK-PAR3-LABEL: func @scale_dd50// CHECK-PAR3: scf.parallel51// CHECK-PAR3: scf.parallel52// CHECK-PAR3: return53//54// CHECK-PAR4-LABEL: func @scale_dd55// CHECK-PAR4: scf.parallel56// CHECK-PAR4: scf.parallel57// CHECK-PAR4: return58//59func.func @scale_dd(%scale: f32,60 %arga: tensor<?x?xf32, #DenseMatrix>,61 %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {62 %0 = linalg.generic #trait_dd63 ins(%arga: tensor<?x?xf32, #DenseMatrix>)64 outs(%argx: tensor<?x?xf32>) {65 ^bb(%a: f32, %x: f32):66 %0 = arith.mulf %a, %scale : f3267 linalg.yield %0 : f3268 } -> tensor<?x?xf32>69 return %0 : tensor<?x?xf32>70}71 72#trait_ss = {73 indexing_maps = [74 affine_map<(i,j) -> (i,j)>, // A75 affine_map<(i,j) -> (i,j)> // X (out)76 ],77 iterator_types = ["parallel", "parallel"],78 doc = "X(i,j) = A(i,j) * SCALE"79}80 81//82// CHECK-PAR0-LABEL: func @scale_ss83// CHECK-PAR0: scf.for84// CHECK-PAR0: scf.for85// CHECK-PAR0: return86//87// CHECK-PAR1-LABEL: func @scale_ss88// CHECK-PAR1: scf.for89// CHECK-PAR1: scf.for90// CHECK-PAR1: return91//92// CHECK-PAR2-LABEL: func @scale_ss93// CHECK-PAR2: scf.parallel94// CHECK-PAR2: scf.for95// CHECK-PAR2: return96//97// CHECK-PAR3-LABEL: func @scale_ss98// CHECK-PAR3: scf.for99// CHECK-PAR3: scf.for100// CHECK-PAR3: return101//102// CHECK-PAR4-LABEL: func @scale_ss103// CHECK-PAR4: scf.parallel104// CHECK-PAR4: scf.parallel105// CHECK-PAR4: return106//107func.func @scale_ss(%scale: f32,108 %arga: tensor<?x?xf32, #SparseMatrix>,109 %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {110 %0 = linalg.generic #trait_ss111 ins(%arga: tensor<?x?xf32, #SparseMatrix>)112 outs(%argx: tensor<?x?xf32>) {113 ^bb(%a: f32, %x: f32):114 %0 = arith.mulf %a, %scale : f32115 linalg.yield %0 : f32116 } -> tensor<?x?xf32>117 return %0 : tensor<?x?xf32>118}119 120#trait_matvec = {121 indexing_maps = [122 affine_map<(i,j) -> (i,j)>, // A123 affine_map<(i,j) -> (j)>, // b124 affine_map<(i,j) -> (i)> // x (out)125 ],126 iterator_types = ["parallel", "reduction"],127 doc = "x(i) += A(i,j) * b(j)"128}129 130//131// CHECK-PAR0-LABEL: func @matvec132// CHECK-PAR0: scf.for133// CHECK-PAR0: scf.for134// CHECK-PAR0: return135//136// CHECK-PAR1-LABEL: func @matvec137// CHECK-PAR1: scf.parallel138// CHECK-PAR1: scf.for139// CHECK-PAR1: return140//141// CHECK-PAR2-LABEL: func @matvec142// CHECK-PAR2: scf.parallel143// CHECK-PAR2: scf.for144// CHECK-PAR2: return145//146// CHECK-PAR3-LABEL: func @matvec147// CHECK-PAR3: scf.parallel148// CHECK-PAR3: scf.for149// CHECK-PAR3: return150//151// CHECK-PAR4-LABEL: func @matvec152// CHECK-PAR4: scf.parallel153// CHECK-PAR4: scf.parallel154// CHECK-PAR4: scf.reduce155// CHECK-PAR4: return156//157func.func @matvec(%arga: tensor<16x32xf32, #CSR>,158 %argb: tensor<32xf32>,159 %argx: tensor<16xf32>) -> tensor<16xf32> {160 %0 = linalg.generic #trait_matvec161 ins(%arga, %argb : tensor<16x32xf32, #CSR>, tensor<32xf32>)162 outs(%argx: tensor<16xf32>) {163 ^bb(%A: f32, %b: f32, %x: f32):164 %0 = arith.mulf %A, %b : f32165 %1 = arith.addf %0, %x : f32166 linalg.yield %1 : f32167 } -> tensor<16xf32>168 return %0 : tensor<16xf32>169}170