51 lines · python
1# RUN: env SUPPORT_LIB=%mlir_cuda_runtime \2# RUN: %PYTHON %s | FileCheck %s3 4# ===----------------------------------------------------------------------===//5# Chapter 0 : Hello World6# ===----------------------------------------------------------------------===//7#8# This program demonstrates Hello World:9# 1. Build MLIR function with arguments10# 2. Build MLIR GPU kernel11# 3. Print from a GPU thread12# 4. Pass arguments, JIT compile and run the MLIR function13#14# ===----------------------------------------------------------------------===//15 16 17from mlir.dialects import gpu18from tools.nvdsl import *19 20 21# 1. The decorator generates a MLIR func.func.22# Everything inside the Python function becomes the body of the func.23# The decorator also translates `alpha` to an `index` type.24@NVDSL.mlir_func25def main(alpha):26 # 2. The decorator generates a MLIR gpu.launch.27 # Everything inside the Python function becomes the body of the gpu.launch.28 # This allows for late outlining of the GPU kernel, enabling optimizations29 # like constant folding from host to device.30 @NVDSL.mlir_gpu_launch(grid=(1, 1, 1), block=(4, 1, 1))31 def kernel():32 tidx = gpu.thread_id(gpu.Dimension.x)33 # + operator generates arith.addi34 myValue = alpha + tidx35 # Print from a GPU thread36 gpu.printf("GPU thread %llu has %llu\n", [tidx, myValue])37 38 # 3. Call the GPU kernel39 kernel()40 41 42alpha = 10043# 4. The `mlir_func` decorator JIT compiles the IR and executes the MLIR function.44main(alpha)45 46 47# CHECK: GPU thread 0 has 10048# CHECK: GPU thread 1 has 10149# CHECK: GPU thread 2 has 10250# CHECK: GPU thread 3 has 10351