diff --git a/stdlib/src/builtin/builtin_list.mojo b/stdlib/src/builtin/builtin_list.mojo index 80e4b61185..81eba3f99a 100644 --- a/stdlib/src/builtin/builtin_list.mojo +++ b/stdlib/src/builtin/builtin_list.mojo @@ -634,29 +634,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[ @@ -664,15 +665,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 c2d3f82a33..15f0fd6474 100644 --- a/stdlib/src/builtin/io.mojo +++ b/stdlib/src/builtin/io.mojo @@ -23,7 +23,7 @@ from sys import ( _libc as libc, ) from sys._libc import dup, fclose, fdopen, fflush -from sys.ffi import OpaquePointer +from sys.ffi import OpaquePointer, c_char from utils import Span, write_buffered, write_args from collections import InlineArray @@ -159,19 +159,16 @@ fn _flush(file: FileDescriptor = stdout): @no_inline fn _printf[ fmt: StringLiteral, *types: AnyType -](*arguments: *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() +](*args: *types, file: FileDescriptor = stdout): + var loaded_pack = args.get_loaded_kgen_pack() + var f = fmt.unsafe_ptr().bitcast[c_char]() @parameter if is_nvidia_gpu(): - _ = external_call["vprintf", Int32]( - fmt.unsafe_cstr_ptr(), Pointer.address_of(loaded_pack) - ) + _ = external_call["vprintf", Int32](f, Pointer.address_of(loaded_pack)) 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[ @@ -181,7 +178,7 @@ fn _printf[ `) -> !pop.scalar`, ], _type=Int32, - ](fd, fmt.unsafe_cstr_ptr(), loaded_pack) + ](fd, f, loaded_pack) # ===----------------------------------------------------------------------=== # @@ -192,7 +189,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: @@ -202,29 +199,26 @@ 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() - - return int( - __mlir_op.`pop.external_call`[ - func = "snprintf".value, - variadicType = __mlir_attr[ - `(`, - `!kgen.pointer>,`, - `!pop.scalar, `, - `!kgen.pointer>`, - `) -> !pop.scalar`, - ], - _type=Int32, - ](str, size, fmt.unsafe_cstr_ptr(), loaded_pack) - ) + + # FIXME: external_call should handle this + var f = fmt.unsafe_ptr().bitcast[c_char]() + var num = __mlir_op.`pop.external_call`[ + func = "snprintf".value, + variadicType = __mlir_attr[ + `(`, + `!kgen.pointer>,`, + `!pop.scalar, `, + `!kgen.pointer>`, + `) -> !pop.scalar`, + ], + _type=Int32, + ](str.bitcast[c_char](), size, f, args.get_loaded_kgen_pack()) + return int(num) # ===----------------------------------------------------------------------=== # diff --git a/stdlib/src/sys/_assembly.mojo b/stdlib/src/sys/_assembly.mojo index 6cb2123393..c393f14999 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 cdfbcfb18f..2ac967f341 100644 --- a/stdlib/src/sys/ffi.mojo +++ b/stdlib/src/sys/ffi.mojo @@ -467,39 +467,52 @@ fn _get_global_or_null[name: StringLiteral]() -> OpaquePointer: # ===----------------------------------------------------------------------===# -@always_inline("nodebug") +@always_inline fn external_call[ - callee: StringLiteral, type: AnyTrivialRegType, *types: AnyType -](*arguments: *types) -> type: + callee: StringLiteral, return_type: AnyTrivialRegType, *T: AnyType +](*args: *T) -> 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. + T: The argument types. Returns: - The external call result. + The external call result. """ + return external_call[callee, return_type](args) - # 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() + +@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. + """ + var p = args.get_loaded_kgen_pack() @parameter - if _mlirtype_is_eq[type, NoneType](): - __mlir_op.`pop.external_call`[func = callee.value, _type=None]( - loaded_pack - ) - return rebind[type](None) + if _mlirtype_is_eq[return_type, NoneType](): + __mlir_op.`pop.external_call`[func = callee.value, _type=None](p) + 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 + ](p) # ===----------------------------------------------------------------------===# @@ -509,29 +522,24 @@ 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: 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() - return __mlir_op.`pop.external_call`[ func = callee.value, resAttrs = __mlir_attr.`[{llvm.noundef}]`, @@ -541,5 +549,5 @@ fn _external_call_const[ `argMem = none, `, `inaccessibleMem = none>`, ], - _type=type, - ](loaded_pack) + _type=return_type, + ](args.get_loaded_kgen_pack()) diff --git a/stdlib/src/sys/intrinsics.mojo b/stdlib/src/sys/intrinsics.mojo index 38f6ae8e5b..b22848ec48 100644 --- a/stdlib/src/sys/intrinsics.mojo +++ b/stdlib/src/sys/intrinsics.mojo @@ -36,25 +36,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]():