[External] [stdlib] Doc improvements and cleanups around low-level `V…

…ariadicPack` use (#48996) [External] [stdlib] Doc improvements and cleanups around low-level VariadicPack use #48996 This includes changes from modularml#3632, minus parts which ended up conflicting with internal changes doing the same refactoring. Context: modularml#3632 (comment) 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 <[email protected]> PUBLIC_PR_LINK=modularml#3632 Co-authored-by: martinvuyk <[email protected]> Closes modularml#3632 MODULAR_ORIG_COMMIT_REV_ID: 8a1c1a5a92c6910039158108ab47046583ae9586
msaelices · Dec 21, 2024 · b845157 · b845157
1 parent c84c329
commit b845157
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 50 deletions.
diff --git a/stdlib/src/builtin/builtin_list.mojo b/stdlib/src/builtin/builtin_list.mojo
@@ -605,45 +605,47 @@ struct VariadicPack[
     # C Pack Utilities
     # ===-------------------------------------------------------------------===#
 
-    # This is the element_types list lowered to `variadic<type>` type for kgen.
     alias _kgen_element_types = rebind[
         __mlir_type.`!kgen.variadic<!kgen.type>`
     ](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>` type for kgen.
+    """
     alias _variadic_pointer_types = __mlir_attr[
         `#kgen.param.expr<variadic_ptr_map, `,
         Self._kgen_element_types,
         `, 0: index>: !kgen.variadic<!kgen.type>`,
     ]
-
-    # 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<type> `, 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[
         `#kgen.param.expr<variadic_ptrremove_map, `,
         Self._variadic_pointer_types,
         `>: !kgen.variadic<!kgen.type>`,
     ]
-
-    # 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<type> `, 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
@@ -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
@@ -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
@@ -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
@@ -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]():