| // RUN: mlir-opt %s --linalg-generalize-named-ops \ |
| // RUN: --pre-sparsification-rewrite \ |
| // RUN: --sparse-reinterpret-map \ |
| // RUN: --sparsification="parallelization-strategy=dense-outer-loop" \ |
| // RUN: --sparse-gpu-codegen | FileCheck %s |
| |
| #CSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }> |
| |
| // |
| // CHECK-LABEL: gpu.module @sparse_kernels |
| // CHECK: gpu.func @kernel1 |
| // CHECK: gpu.func @kernel0 |
| // |
| // CHECK-LABEL: func.func @matmuls |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: %[[T1:.*]] = gpu.launch_func async @sparse_kernels::@kernel1 blocks |
| // CHECK: gpu.memcpy async [%[[T1]]] |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.wait |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: gpu.alloc async |
| // CHECK: gpu.memcpy async |
| // CHECK: %[[T0:.*]] = gpu.launch_func async @sparse_kernels::@kernel0 blocks |
| // CHECK: gpu.memcpy async [%[[T0]]] |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.dealloc async |
| // CHECK: gpu.wait |
| // |
| func.func @matmuls(%A: tensor<1024x8xf64>, |
| %B: tensor<8x1024xf64, #CSR>, |
| %C: tensor<1024x1024xf64, #CSR>) -> tensor<1024x1024xf64> { |
| %Z = arith.constant dense<0.0> : tensor<1024x1024xf64> |
| %T = linalg.matmul |
| ins(%A, %B: tensor<1024x8xf64>, tensor<8x1024xf64, #CSR>) |
| outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64> |
| %D = linalg.matmul |
| ins(%T, %C: tensor<1024x1024xf64>, tensor<1024x1024xf64, #CSR>) |
| outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64> |
| return %D : tensor<1024x1024xf64> |
| } |
