64 lines · plain
1// RUN: mlir-opt %s --linalg-generalize-named-ops \2// RUN: --pre-sparsification-rewrite \3// RUN: --sparse-reinterpret-map \4// RUN: --sparsification="parallelization-strategy=dense-outer-loop" \5// RUN: --sparse-gpu-codegen | FileCheck %s6 7#CSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }>8 9//10// CHECK-LABEL: gpu.module @sparse_kernels11// CHECK: gpu.func @kernel112// CHECK: gpu.func @kernel013//14// CHECK-LABEL: func.func @matmuls15// CHECK: gpu.alloc async16// CHECK: gpu.memcpy async17// CHECK: gpu.alloc async18// CHECK: gpu.memcpy async19// CHECK: gpu.alloc async20// CHECK: gpu.memcpy async21// CHECK: gpu.alloc async22// CHECK: gpu.memcpy async23// CHECK: gpu.alloc async24// CHECK: gpu.memcpy async25// CHECK: %[[T1:.*]] = gpu.launch_func async @sparse_kernels::@kernel1 blocks26// CHECK: gpu.memcpy async [%[[T1]]]27// CHECK: gpu.dealloc async28// CHECK: gpu.dealloc async29// CHECK: gpu.dealloc async30// CHECK: gpu.dealloc async31// CHECK: gpu.dealloc async32// CHECK: gpu.wait33// CHECK: gpu.alloc async34// CHECK: gpu.memcpy async35// CHECK: gpu.alloc async36// CHECK: gpu.memcpy async37// CHECK: gpu.alloc async38// CHECK: gpu.memcpy async39// CHECK: gpu.alloc async40// CHECK: gpu.memcpy async41// CHECK: gpu.alloc async42// CHECK: gpu.memcpy async43// CHECK: %[[T0:.*]] = gpu.launch_func async @sparse_kernels::@kernel0 blocks44// CHECK: gpu.memcpy async [%[[T0]]]45// CHECK: gpu.dealloc async46// CHECK: gpu.dealloc async47// CHECK: gpu.dealloc async48// CHECK: gpu.dealloc async49// CHECK: gpu.dealloc async50// CHECK: gpu.wait51//52func.func @matmuls(%A: tensor<1024x8xf64>,53 %B: tensor<8x1024xf64, #CSR>,54 %C: tensor<1024x1024xf64, #CSR>) -> tensor<1024x1024xf64> {55 %Z = arith.constant dense<0.0> : tensor<1024x1024xf64>56 %T = linalg.matmul57 ins(%A, %B: tensor<1024x8xf64>, tensor<8x1024xf64, #CSR>)58 outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64>59 %D = linalg.matmul60 ins(%T, %C: tensor<1024x1024xf64>, tensor<1024x1024xf64, #CSR>)61 outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64>62 return %D : tensor<1024x1024xf64>63}64