From b84515703f8c53847eb19f883e6aeecd115f4dc2 Mon Sep 17 00:00:00 2001 From: modularbot <116839051+modularbot@users.noreply.github.com> Date: Wed, 18 Dec 2024 18:34:59 -0600 Subject: [PATCH] [External] [stdlib] Doc improvements and cleanups around low-level `VariadicPack` use (#48996) [External] [stdlib] Doc improvements and cleanups around low-level VariadicPack use #48996 This includes changes from modularml/mojo#3632, minus parts which ended up conflicting with internal changes doing the same refactoring. Context: https://github.com/modularml/mojo/pull/3632#issuecomment-2462756696 Original commit description: [External] [stdlib] Refactor `external_call()` and inline `_LITRefPackHelper` into `VariadicPack` Refactor `external_call()` and inline most of `_LITRefPackHelper` into `VariadicPack`, leave `_LITRefPackHelper` itself intact since I found no way to inline the `AddressSpace` parameter into `VariadicPack`. ORIGINAL_AUTHOR=martinvuyk <110240700+martinvuyk@users.noreply.github.com> PUBLIC_PR_LINK=modularml/mojo#3632 Co-authored-by: martinvuyk <110240700+martinvuyk@users.noreply.github.com> Closes modularml/mojo#3632 MODULAR_ORIG_COMMIT_REV_ID: 8a1c1a5a92c6910039158108ab47046583ae9586 --- stdlib/src/builtin/builtin_list.mojo | 26 ++++++------ stdlib/src/builtin/io.mojo | 15 ++++--- stdlib/src/sys/_assembly.mojo | 4 +- stdlib/src/sys/ffi.mojo | 63 +++++++++++++++++++--------- stdlib/src/sys/intrinsics.mojo | 21 +++++----- 5 files changed, 79 insertions(+), 50 deletions(-) diff --git a/stdlib/src/builtin/builtin_list.mojo b/stdlib/src/builtin/builtin_list.mojo index e74b93cad3..50e63d02d2 100644 --- a/stdlib/src/builtin/builtin_list.mojo +++ b/stdlib/src/builtin/builtin_list.mojo @@ -605,29 +605,30 @@ struct VariadicPack[ # C Pack Utilities # ===-------------------------------------------------------------------===# - # This is the element_types list lowered to `variadic` type for kgen. alias _kgen_element_types = rebind[ __mlir_type.`!kgen.variadic` ](Self.element_types) - - # Use variadic_ptr_map to construct the type list of the !kgen.pack that the - # !lit.ref.pack will lower to. It exposes the pointers introduced by the - # references. + """This is the element_types list lowered to `variadic` type for kgen. + """ alias _variadic_pointer_types = __mlir_attr[ `#kgen.param.expr: !kgen.variadic`, ] - - # This is the !kgen.pack type with pointer elements. + """Use variadic_ptr_map to construct the type list of the !kgen.pack that + the !lit.ref.pack will lower to. It exposes the pointers introduced by the + references. + """ alias _kgen_pack_with_pointer_type = __mlir_type[ `!kgen.pack<:variadic `, Self._variadic_pointer_types, `>` ] + """This is the !kgen.pack type with pointer elements.""" - # This rebinds `in_pack` to the equivalent `!kgen.pack` with kgen pointers. @doc_private @always_inline("nodebug") fn get_as_kgen_pack(self) -> Self._kgen_pack_with_pointer_type: + """This rebinds `in_pack` to the equivalent `!kgen.pack` with kgen + pointers.""" return rebind[Self._kgen_pack_with_pointer_type](self._value) alias _variadic_with_pointers_removed = __mlir_attr[ @@ -635,15 +636,16 @@ struct VariadicPack[ Self._variadic_pointer_types, `>: !kgen.variadic`, ] - - # This is the `!kgen.pack` type that happens if one loads all the elements - # of the pack. alias _loaded_kgen_pack_type = __mlir_type[ `!kgen.pack<:variadic `, Self._variadic_with_pointers_removed, `>` ] + """This is the `!kgen.pack` type that happens if one loads all the elements + of the pack. + """ - # This returns the stored KGEN pack after loading all of the elements. @doc_private @always_inline("nodebug") fn get_loaded_kgen_pack(self) -> Self._loaded_kgen_pack_type: + """This returns the stored KGEN pack after loading all of the elements. + """ return __mlir_op.`kgen.pack.load`(self.get_as_kgen_pack()) diff --git a/stdlib/src/builtin/io.mojo b/stdlib/src/builtin/io.mojo index fc363fc15c..a409c35c62 100644 --- a/stdlib/src/builtin/io.mojo +++ b/stdlib/src/builtin/io.mojo @@ -26,7 +26,7 @@ from sys import ( stdout, ) from sys._libc import dup, fclose, fdopen, fflush -from sys.ffi import OpaquePointer +from sys.ffi import OpaquePointer, c_char from builtin.dtype import _get_dtype_printf_format from builtin.file_descriptor import FileDescriptor @@ -165,11 +165,11 @@ fn _flush(file: FileDescriptor = stdout): @no_inline fn _printf[ fmt: StringLiteral, *types: AnyType -](*arguments: *types, file: FileDescriptor = stdout): +](*args: *types, file: FileDescriptor = stdout): # The argument pack will contain references for each value in the pack, # but we want to pass their values directly into the C printf call. Load # all the members of the pack. - var loaded_pack = arguments.get_loaded_kgen_pack() + var loaded_pack = args.get_loaded_kgen_pack() @parameter if is_nvidia_gpu(): @@ -181,6 +181,7 @@ fn _printf[ pass else: with _fdopen(file) as fd: + # FIXME: external_call should handle this _ = __mlir_op.`pop.external_call`[ func = "KGEN_CompilerRT_fprintf".value, variadicType = __mlir_attr[ @@ -201,7 +202,7 @@ fn _printf[ @no_inline fn _snprintf[ fmt: StringLiteral, *types: AnyType -](str: UnsafePointer[UInt8], size: Int, *arguments: *types) -> Int: +](str: UnsafePointer[UInt8], size: Int, *args: *types) -> Int: """Writes a format string into an output pointer. Parameters: @@ -211,16 +212,18 @@ fn _snprintf[ Args: str: A pointer into which the format string is written. size: At most, `size - 1` bytes are written into the output string. - arguments: Arguments interpolated into the format string. + args: Arguments interpolated into the format string. Returns: The number of bytes written into the output string. """ + # The argument pack will contain references for each value in the pack, # but we want to pass their values directly into the C snprintf call. Load # all the members of the pack. - var loaded_pack = arguments.get_loaded_kgen_pack() + var loaded_pack = args.get_loaded_kgen_pack() + # FIXME: external_call should handle this return int( __mlir_op.`pop.external_call`[ func = "snprintf".value, diff --git a/stdlib/src/sys/_assembly.mojo b/stdlib/src/sys/_assembly.mojo index 2074afde29..c48f518cd6 100644 --- a/stdlib/src/sys/_assembly.mojo +++ b/stdlib/src/sys/_assembly.mojo @@ -26,9 +26,9 @@ fn inlined_assembly[ *types: AnyType, constraints: StringLiteral, has_side_effect: Bool = True, -](*arguments: *types) -> result_type: +](*args: *types) -> result_type: """Generates assembly via inline assembly.""" - var loaded_pack = arguments.get_loaded_kgen_pack() + var loaded_pack = args.get_loaded_kgen_pack() @parameter if _mlirtype_is_eq[result_type, NoneType](): diff --git a/stdlib/src/sys/ffi.mojo b/stdlib/src/sys/ffi.mojo index 397bc099a3..429dd1e0f7 100644 --- a/stdlib/src/sys/ffi.mojo +++ b/stdlib/src/sys/ffi.mojo @@ -496,37 +496,60 @@ fn _get_global_or_null[name: StringLiteral]() -> OpaquePointer: @always_inline("nodebug") fn external_call[ - callee: StringLiteral, type: AnyTrivialRegType, *types: AnyType -](*arguments: *types) -> type: + callee: StringLiteral, + return_type: AnyTrivialRegType, + *types: AnyType, +](*args: *types) -> return_type: """Calls an external function. Args: - arguments: The arguments to pass to the external function. + args: The arguments to pass to the external function. Parameters: - callee: The name of the external function. - type: The return type. - types: The argument types. + callee: The name of the external function. + return_type: The return type. + types: The argument types. Returns: - The external call result. + The external call result. + """ + return external_call[callee, return_type](args) + + +@always_inline("nodebug") +fn external_call[ + callee: StringLiteral, + return_type: AnyTrivialRegType, +](args: VariadicPack[element_trait=AnyType]) -> return_type: + """Calls an external function. + + Parameters: + callee: The name of the external function. + return_type: The return type. + + Args: + args: The arguments to pass to the external function. + + Returns: + The external call result. """ # The argument pack will contain references for each value in the pack, # but we want to pass their values directly into the C printf call. Load # all the members of the pack. - var loaded_pack = arguments.get_loaded_kgen_pack() + var loaded_pack = args.get_loaded_kgen_pack() @parameter - if _mlirtype_is_eq[type, NoneType](): + if _mlirtype_is_eq[return_type, NoneType](): __mlir_op.`pop.external_call`[func = callee.value, _type=None]( loaded_pack ) - return rebind[type](None) + return rebind[return_type](None) else: - return __mlir_op.`pop.external_call`[func = callee.value, _type=type]( - loaded_pack - ) + return __mlir_op.`pop.external_call`[ + func = callee.value, + _type=return_type, + ](loaded_pack) # ===-----------------------------------------------------------------------===# @@ -536,18 +559,20 @@ fn external_call[ @always_inline("nodebug") fn _external_call_const[ - callee: StringLiteral, type: AnyTrivialRegType, *types: AnyType -](*arguments: *types) -> type: + callee: StringLiteral, + return_type: AnyTrivialRegType, + *types: AnyType, +](*args: *types) -> return_type: """Mark the external function call as having no observable effects to the program state. This allows the compiler to optimize away successive calls to the same function. Args: - arguments: The arguments to pass to the external function. + args: The arguments to pass to the external function. Parameters: callee: The name of the external function. - type: The return type. + return_type: The return type. types: The argument types. Returns: @@ -557,7 +582,7 @@ fn _external_call_const[ # The argument pack will contain references for each value in the pack, # but we want to pass their values directly into the C printf call. Load # all the members of the pack. - var loaded_pack = arguments.get_loaded_kgen_pack() + var loaded_pack = args.get_loaded_kgen_pack() return __mlir_op.`pop.external_call`[ func = callee.value, @@ -568,5 +593,5 @@ fn _external_call_const[ `argMem = none, `, `inaccessibleMem = none>`, ], - _type=type, + _type=return_type, ](loaded_pack) diff --git a/stdlib/src/sys/intrinsics.mojo b/stdlib/src/sys/intrinsics.mojo index 3eb81a505a..1a7242cc63 100644 --- a/stdlib/src/sys/intrinsics.mojo +++ b/stdlib/src/sys/intrinsics.mojo @@ -37,25 +37,24 @@ fn llvm_intrinsic[ type: AnyTrivialRegType, *types: AnyType, has_side_effect: Bool = True, -](*arguments: *types) -> type: - """Calls an LLVM intrinsic with no arguments. - - Calls an LLVM intrinsic with the name intrin and return type type. +](*args: *types) -> type: + """Calls an LLVM intrinsic with the name `intrin` and return type `type`. Parameters: - intrin: The name of the llvm intrinsic. - type: The return type of the intrinsic. - types: The argument types for the function. - has_side_effect: If `True` the intrinsic will have side effects, otherwise its pure. + intrin: The name of the llvm intrinsic. + type: The return type of the intrinsic. + types: The argument types for the function. + has_side_effect: If `True` the intrinsic will have side effects, + otherwise its pure. Args: - arguments: The arguments to the function. + args: The arguments to the function. Returns: - The result of calling the llvm intrinsic with no arguments. + The result of calling the llvm intrinsic with no arguments. """ - var loaded_pack = arguments.get_loaded_kgen_pack() + var loaded_pack = args.get_loaded_kgen_pack() @parameter if _mlirtype_is_eq[type, NoneType]():