135 lines · plain
1// NOTE: this test requires gpu-sm80 and cusparselt2//3// DEFINE: %{compile} = mlir-opt %s \4// DEFINE: --sparsifier="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71 gpu-format=%gpu_compilation_format5// DEFINE: %{run} = mlir-runner \6// DEFINE: --shared-libs=%mlir_cuda_runtime \7// DEFINE: --shared-libs=%mlir_c_runner_utils \8// DEFINE: --e main --entry-point-result=void \9// DEFINE: | FileCheck %s10//11// with RT lib:12//13// RUN: %{compile} enable-runtime-library=true" | %{run}14//15// without RT lib:16//17// RUN: %{compile} enable-runtime-library=false" | %{run}18 19#NV_24 = #sparse_tensor.encoding<{20 map = ( i, j ) ->21 ( i : dense,22 j floordiv 4 : dense,23 j mod 4 : structured[2, 4]24 )25}>26 27module {28 29 llvm.func @mgpuCreateSparseLtEnv()30 llvm.func @mgpuDestroySparseLtEnv()31 32 func.func @matmul24(%Ad: tensor<16x16xf16>,33 %B: tensor<16x16xf16>,34 %Cin: tensor<16x16xf16>) -> tensor<16x16xf16> {35 %A = sparse_tensor.convert %Ad : tensor<16x16xf16> to tensor<16x16xf16, #NV_24>36 %C = linalg.matmul37 ins(%A, %B: tensor<16x16xf16, #NV_24>, tensor<16x16xf16>)38 outs(%Cin: tensor<16x16xf16>) -> tensor<16x16xf16>39 return %C : tensor<16x16xf16>40 }41 42 func.func @main() {43 llvm.call @mgpuCreateSparseLtEnv() : () -> ()44 45 %c0 = arith.constant 0 : index46 %c1 = arith.constant 1 : index47 %c16 = arith.constant 16 : index48 49 %f0 = arith.constant 0.0 : f1650 %f1 = arith.constant 1.0 : f1651 %f4 = arith.constant 4.0 : f1652 53 // Initial A, B, C matrices.54 %A = tensor.generate {55 ^bb0(%i: index, %j: index):56 %val = arith.andi %j, %c1 : index57 %cmp = arith.cmpi eq, %val, %c0 : index58 %res = arith.select %cmp, %f4, %f1 : f1659 tensor.yield %res : f1660 } : tensor<16x16xf16>61 %B = tensor.generate {62 ^bb0(%i: index, %j: index):63 %cmp = arith.cmpi eq, %i, %j : index64 %res = arith.select %cmp, %f1, %f0 : f1665 tensor.yield %res : f1666 } : tensor<16x16xf16>67 %C = tensor.generate {68 ^bb0(%i: index, %j: index):69 tensor.yield %f0 : f1670 } : tensor<16x16xf16>71 72 // Call the kernel.73 //74 // By effectively computing D = A B + C with id(B) and zero(C)75 // the resulting matrix returns the pruned A back to the caller.76 //77 %D = call @matmul24(%A, %B, %C): (tensor<16x16xf16>,78 tensor<16x16xf16>,79 tensor<16x16xf16>) -> (tensor<16x16xf16>)80 81 //82 // This was the original matrix.83 //84 // CHECK: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )85 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )86 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )87 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )88 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )89 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )90 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )91 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )92 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )93 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )94 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )95 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )96 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )97 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )98 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )99 // CHECK-NEXT: ( 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 1 )100 //101 scf.for %i = %c0 to %c16 step %c1 {102 %va = vector.transfer_read %A[%i, %c0], %f0 : tensor<16x16xf16>, vector<16xf16>103 vector.print %va : vector<16xf16>104 }105 106 //107 // This is the STRIP-pruned matrix.108 //109 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )110 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )111 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )112 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )113 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )114 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )115 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )116 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )117 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )118 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )119 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )120 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )121 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )122 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )123 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )124 // CHECK-NEXT: ( 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0 )125 //126 scf.for %i = %c0 to %c16 step %c1 {127 %vd = vector.transfer_read %D[%i, %c0], %f0 : tensor<16x16xf16>, vector<16xf16>128 vector.print %vd : vector<16xf16>129 }130 131 llvm.call @mgpuDestroySparseLtEnv() : () -> ()132 return133 }134}135