| from enum import Enum |
| import functools, sys, ctypes, os, errno |
| import numpy as np |
| from functools import partialmethod |
| from mlir import ir |
| from mlir.dialects import arith, func, gpu, memref, nvgpu, scf, nvvm |
| from mlir.extras import types as T |
| from mlir import runtime as rt |
| from tools import nvgpucompiler |
| |
| MLIR_DYNAMIC = -9223372036854775808 |
| |
| |
| def const(value: int, ty=None): |
| ty = T.index() if ty is None else ty |
| if isinstance(value, ir.Value) and ( |
| value.type.isinstance(value.type) or T.bool().isinstance(value.type) |
| ): |
| return value |
| return arith.constant(ty, value) |
| |
| |
| def get_type_size(ty): |
| if ir.MemRefType.isinstance(ty): |
| size = get_type_size(ty.element_type) |
| for sz in ty.shape: |
| size *= sz |
| return size |
| if ir.FloatType.isinstance(ty): |
| return ir.FloatType(ty).width // 8 |
| if ir.IntegerType.isinstance(ty): |
| return ir.IntegerType(ty).width // 8 |
| raise NotImplementedError(ty) |
| |
| |
| def get_mlir_func_obj_ty(inputArgs): |
| args = [] |
| c_int_p = ctypes.c_int * 1 |
| c_float_p = ctypes.c_float * 1 |
| c_bool_p = ctypes.c_bool * 1 |
| for arg in inputArgs: |
| if isinstance(arg, bool): |
| args.append(c_bool_p(arg)) |
| elif isinstance(arg, int): |
| args.append(c_int_p(arg)) |
| elif isinstance(arg, float): |
| args.append(c_float_p(arg)) |
| elif isinstance(arg, np.ndarray): |
| args.append( |
| ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(arg))) |
| ) |
| else: |
| raise NotImplementedError(arg) |
| return args |
| |
| |
| class Mbarriers: |
| def __init__(self, number_of_barriers=1): |
| self.mbar_ty = ir.Type.parse( |
| "!nvgpu.mbarrier.group<memorySpace=#gpu.address_space<workgroup>, num_barriers = " |
| + str(number_of_barriers) |
| + ">" |
| ) |
| self.mbar_group_op = nvgpu.mbarrier_create(self.mbar_ty) |
| self.number_of_barriers = number_of_barriers |
| |
| def __getitem__(self, key): |
| self.id_op = const(key) |
| return self |
| |
| def init(self, count: int, predicate=None): |
| count_op = const(count) |
| if predicate is None: |
| nvgpu.mbarrier_init(self.mbar_group_op, count_op, self.id_op) |
| else: |
| nvgpu.mbarrier_init( |
| self.mbar_group_op, count_op, self.id_op, predicate=predicate |
| ) |
| |
| def arrive(self, txcount: int = 0, predicate=None): |
| if txcount != 0: |
| txcount_op = const(txcount) |
| nvgpu.mbarrier_arrive_expect_tx( |
| self.mbar_group_op, txcount_op, self.id_op, predicate=predicate |
| ) |
| else: |
| nvgpu.mbarrier_arrive( |
| ir.Type.parse("!nvgpu.mbarrier.token"), self.mbar_group_op, self.id_op |
| ) |
| |
| def try_wait(self, phase: bool = False, ticks: int = 10000000): |
| ticks_op = const(ticks) |
| phase_op = const(phase, T.bool()) |
| nvgpu.MBarrierTryWaitParityOp( |
| self.mbar_group_op, |
| phase_op, |
| ticks_op, |
| mbarId=self.id_op, |
| ) |
| |
| |
| class TMA: |
| """A class that builds a TMA descriptor.""" |
| |
| def __init__( |
| self, |
| tma_box_shape, |
| memref_ty, |
| swizzle=nvgpu.TensorMapSwizzleKind.SWIZZLE_NONE, |
| l2promo=nvgpu.TensorMapL2PromoKind.L2PROMO_NONE, |
| oob=nvgpu.TensorMapOOBKind.OOB_ZERO, |
| interleave=nvgpu.TensorMapInterleaveKind.INTERLEAVE_NONE, |
| ): |
| self.swizzle = swizzle # mlir.nvgpu.TensorMapSwizzleKind |
| self.l2promo = l2promo # mlir.nvgpu.TensorMapL2PromoKind |
| self.oob = oob # mlir.nvgpu.TensorMapOOBKind |
| self.interleave = interleave # mlir.nvgpu.TensorMapInterleaveKind |
| self.tma_box_shape = tma_box_shape |
| self.memref_ty = memref_ty # MemRefType |
| self.tma_memref = ir.MemRefType.get(tma_box_shape, memref_ty.element_type) |
| |
| @property |
| def tensormap_descriptor_ty(self): |
| """Returns a tensormap descriptor type.""" |
| tensorMemrefType = ir.MemRefType.get( |
| self.tma_box_shape, |
| self.memref_ty.element_type, |
| memory_space=ir.Attribute.parse("3"), |
| ) |
| return nvgpu.TensorMapDescriptorType.get( |
| tensorMemrefType, |
| self.swizzle, |
| self.l2promo, |
| self.oob, |
| self.interleave, |
| ) |
| |
| def create_descriptor(self, device_ptr): |
| tma_descriptor_ty = self.tensormap_descriptor_ty |
| device_unranked_memref = memref.CastOp( |
| ir.UnrankedMemRefType.get( |
| self.memref_ty.element_type, self.memref_ty.memory_space |
| ), |
| device_ptr, |
| ) |
| self.tma_descriptor = nvgpu.TmaCreateDescriptorOp( |
| tma_descriptor_ty, device_unranked_memref, map(const, self.tma_box_shape) |
| ) |
| return self.tma_descriptor.result |
| |
| def prefetch(self, predicate=None): |
| nvgpu.tma_prefetch_descriptor(self.tma_descriptor, predicate=predicate) |
| |
| def load(self, dest, mbarrier: Mbarriers, coords=[0], predicate=None): |
| nvgpu.TmaAsyncLoadOp( |
| dest, |
| mbarrier.mbar_group_op, |
| self.tma_descriptor, |
| coordinates=map(const, coords), |
| mbarId=mbarrier.id_op, |
| predicate=predicate, |
| ) |
| |
| |
| WARP_GROUP_SIZE = 128 # Number of threads in a warpgroup |
| |
| |
| class Warpgroup: |
| def __init__(self, primary_thread, register_size): |
| assert (primary_thread % WARP_GROUP_SIZE) == 0 |
| tidx = gpu.thread_id(gpu.Dimension.x) |
| self.primary_thread = primary_thread |
| self.register_size = register_size |
| self.is_wg_primary = (tidx % WARP_GROUP_SIZE) == 0 |
| self.wg_id = tidx / WARP_GROUP_SIZE |
| self.is_me = self.wg_id == (primary_thread // WARP_GROUP_SIZE) |
| |
| def __enter__(self): |
| if_op = scf.IfOp(self.is_me) |
| self.ipoint_op = ir.InsertionPoint(if_op.then_block) |
| self.ipoint_op.__enter__() |
| if self.register_size < 64: |
| nvvm.setmaxregister(self.register_size, nvvm.SetMaxRegisterAction.decrease) |
| else: |
| nvvm.setmaxregister(self.register_size, nvvm.SetMaxRegisterAction.increase) |
| |
| def __exit__(self, exc_type, exc_value, traceback): |
| scf.yield_([]) |
| self.ipoint_op.__exit__(exc_type, exc_value, traceback) |
| return True |
| |
| |
| class WGMMAType(Enum): |
| Accumulator = 1 |
| Descriptor = 2 |
| |
| |
| class WGMMAMatrix: |
| def __init__( |
| self, |
| matrix_type: WGMMAType, |
| shape: list = None, |
| desc: TMA = None, |
| smem=None, |
| ty=None, |
| acc_op=None, |
| ): |
| if acc_op is None: |
| self.M = shape[0] |
| self.N = shape[1] |
| self.ty = ty |
| self.matrix_type = matrix_type |
| self.desc = desc |
| self.smem = smem |
| if matrix_type is WGMMAType.Accumulator: |
| self.acc_op = nvgpu.warpgroup_mma_init_accumulator(self.acc_ty) |
| elif acc_op: |
| self.acc_op = acc_op |
| self.matrix_type = WGMMAType.Accumulator |
| |
| @property |
| def acc_ty(self): |
| parse_str = f"!nvgpu.warpgroup.accumulator<fragmented=vector<{self.M}x{self.N}x{self.ty}>>" |
| return ir.Type.parse(parse_str) |
| |
| @property |
| def wgmma_ty(self): |
| parse_str = f"!nvgpu.warpgroup.descriptor<tensor=memref<{self.M}x{self.N}x{self.desc.memref_ty.element_type}, #gpu.address_space<workgroup>>>" |
| return ir.Type.parse(parse_str) |
| |
| def store_accumulator(self, dest): |
| assert self.matrix_type == WGMMAType.Accumulator |
| nvgpu.warpgroup_mma_store(self.acc_op, dest) |
| |
| def update_smem(self, smem): |
| self.smem = smem |
| |
| def update_accumulator(self, acc_op): |
| self.acc_op = acc_op |
| |
| def __matmul__(self, rhs): |
| lhs = nvgpu.warpgroup_generate_descriptor( |
| self.wgmma_ty, self.smem, self.desc.tma_descriptor |
| ) |
| rhs = nvgpu.warpgroup_generate_descriptor( |
| rhs.wgmma_ty, rhs.smem, rhs.desc.tma_descriptor |
| ) |
| return [lhs, rhs] |
| |
| def __iadd__(self, matmulResult): |
| lhs = matmulResult[0] |
| rhs = matmulResult[1] |
| acc_op = nvgpu.WarpgroupMmaOp( |
| self.acc_op.type, lhs, rhs, self.acc_op, transposeB=True |
| ) |
| return WGMMAMatrix(WGMMAType.Accumulator, acc_op=acc_op) |
| |
| |
| def get_dynamic_shared_memory(shape=None, ty=None, offset: int = 0): |
| smem_space_str = "#gpu.address_space<workgroup>" |
| smem_space = ir.Attribute.parse(smem_space_str) |
| dynamic_smem = gpu.dynamic_shared_memory( |
| ir.MemRefType.get((MLIR_DYNAMIC,), T.i8(), memory_space=smem_space) |
| ) |
| if shape is None: |
| return dynamic_smem |
| memref_ty = ir.MemRefType.get(shape, ty, memory_space=smem_space) |
| return memref.view( |
| ir.MemRefType.get( |
| memref_ty.shape, memref_ty.element_type, memory_space=smem_space |
| ), |
| dynamic_smem, |
| const(offset), |
| [], |
| ) |
| |
| |
| def get_mlir_ty(arg): |
| def get_mlir_ty_from_np(dtype): |
| if dtype == np.float16: |
| return T.f16() |
| if dtype == np.float32: |
| return T.f32() |
| if dtype == np.float64: |
| return T.f64() |
| if dtype == np.int32: |
| return T.i32() |
| if dtype == np.int64: |
| return T.i64() |
| raise NotImplementedError(dtype) |
| |
| if isinstance(arg, bool): |
| return T.bool() |
| elif isinstance(arg, int): |
| return T.index() |
| elif isinstance(arg, float): |
| return T.f32() |
| elif isinstance(arg, np.ndarray): |
| descriptor = rt.get_ranked_memref_descriptor(arg) |
| dtype = get_mlir_ty_from_np(arg.dtype) |
| shape = descriptor.shape |
| return memref.MemRefType.get(shape, dtype) |
| raise NotImplementedError(arg) |
| |
| |
| class NVDSL: |
| @staticmethod |
| def mlir_gpu_launch(grid=(1, 1, 1), block=(1, 1, 1), smem=0): |
| def decorator(func): |
| @functools.wraps(func) |
| def wrapper(*args, **kwargs): |
| launch_op = gpu.LaunchOp( |
| None, |
| [], |
| *map(const, grid), |
| *map(const, block), |
| dynamicSharedMemorySize=arith.constant(T.i32(), smem), |
| ) |
| launch_op.body.blocks.append(*([T.index()] * 12)) |
| with ir.InsertionPoint(launch_op.body.blocks[0]): |
| result = func(*args, **kwargs) |
| gpu.terminator() |
| return result |
| |
| return wrapper |
| |
| return decorator |
| |
| @staticmethod |
| def mlir_func(funcBody): |
| @functools.wraps(funcBody) |
| def wrapper(*args, **kwargs): |
| function_name = funcBody.__name__ |
| |
| def saveIR(module): |
| """Save generated IR""" |
| if True: # self.saveIR: |
| # print(mlir_nvgpu_module) |
| original_stdout = sys.stdout |
| with open("nvdsl.mlir", "w") as f: |
| sys.stdout = f |
| print(module) |
| sys.stdout = original_stdout |
| |
| def _binary_op(lhs, rhs, op: str, predAtt="") -> "ArithValue": |
| """Generate MLIR's Arith dialects binary operations.""" |
| rhs = const(rhs) |
| if arith._is_float_type(lhs.type) and arith._is_float_type(rhs.type): |
| op += "F" |
| if op.startswith("Cmp"): |
| predicateAttr = getattr(arith, f"CmpFPredicate").__dict__[ |
| predAtt |
| ] |
| elif arith._is_integer_like_type( |
| lhs.type |
| ) and arith._is_integer_like_type(lhs.type): |
| if op == "Div" or op == "Rem": |
| op += "U" |
| op += "I" |
| if op.startswith("Cmp"): |
| predicateAttr = getattr(arith, f"CmpIPredicate").__dict__[ |
| predAtt |
| ] |
| else: |
| raise NotImplementedError( |
| f"Unsupported '{op}' operands: {lhs}, {rhs}" |
| ) |
| |
| if op.startswith("Cmp"): |
| op = getattr(arith, f"{op}Op") |
| |
| return op(predicateAttr, lhs, rhs).result |
| else: |
| op = getattr(arith, f"{op}Op") |
| return op(lhs, rhs).result |
| |
| @ir.register_value_caster(ir.IndexType.static_typeid) |
| @ir.register_value_caster(ir.F32Type.static_typeid) |
| @ir.register_value_caster(ir.F16Type.static_typeid) |
| @ir.register_value_caster(ir.F64Type.static_typeid) |
| @ir.register_value_caster(ir.IntegerType.static_typeid) |
| class ArithValue(ir.Value): |
| """Overloads operators for MLIR's Arith dialects binary operations.""" |
| |
| def __init__(self, v): |
| super().__init__(v) |
| |
| __add__ = partialmethod(_binary_op, op="Add") |
| __sub__ = partialmethod(_binary_op, op="Sub") |
| __mul__ = partialmethod(_binary_op, op="Mul") |
| __truediv__ = partialmethod(_binary_op, op="Div") |
| __mod__ = partialmethod(_binary_op, op="Rem") |
| __xor__ = partialmethod(_binary_op, op="XOr") |
| __lt__ = partialmethod(_binary_op, op="Cmp", predAtt="ult") |
| __le__ = partialmethod(_binary_op, op="Cmp", predAtt="ule") |
| __eq__ = partialmethod(_binary_op, op="Cmp", predAtt="eq") |
| __ne__ = partialmethod(_binary_op, op="Cmp", predAtt="ne") |
| __gt__ = partialmethod(_binary_op, op="Cmp", predAtt="ugt") |
| __ge__ = partialmethod(_binary_op, op="Cmp", predAtt="uge") |
| __and__ = partialmethod(_binary_op, op="And") |
| __or__ = partialmethod(_binary_op, op="Or") |
| |
| def __str__(self): |
| return ( |
| super() |
| .__str__() |
| .replace(ir.Value.__name__, ArithValue.__name__) |
| ) |
| |
| # Generate MLIR Context and start generating IR |
| with ir.Context(), ir.Location.unknown(): |
| types = [] |
| for arg in args: |
| types.append(get_mlir_ty(arg)) |
| |
| # Build IR |
| module = ir.Module.create() |
| with ir.InsertionPoint(module.body): |
| fop = func.FuncOp(function_name, (types, [])) |
| fop.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get() |
| with ir.InsertionPoint(fop.add_entry_block()): |
| fargs = [] |
| for i, a in enumerate(types): |
| fargs.append(fop.arguments[i]) |
| |
| # Call user function body |
| result = funcBody(*fargs, **kwargs) |
| func.ReturnOp([]) |
| |
| # Save IR in a file |
| # saveIR(module) |
| |
| # Verify the module |
| # module.operation.verify() |
| |
| # Compile and JIT MLIR module |
| options = f"cubin-chip=sm_90a cubin-features=+ptx80 opt-level=3" |
| support_lib = os.getenv("SUPPORT_LIB") |
| if not os.path.exists(support_lib): |
| raise FileNotFoundError( |
| errno.ENOENT, os.strerror(errno.ENOENT), support_lib |
| ) |
| compiler = nvgpucompiler.NvgpuCompiler( |
| options, opt_level=3, shared_libs=[support_lib] |
| ) |
| engine = compiler.compile_and_jit(module) |
| |
| # Convert input arguments to MLIR arguments |
| newArgs = get_mlir_func_obj_ty(args) |
| |
| # Run the compiled program |
| engine.invoke(function_name, *newArgs) |
| |
| return result |
| |
| return wrapper |
