Skip to main content

Mojo function

matmul

matmul[transpose_a: Bool = False, transpose_b: Bool = False, b_packed: Bool = False, elementwise_lambda_fn: Optional[elementwise_epilogue_type] = None, elementwise_compute_lambda_fn: Optional[elementwise_compute_lambda_type] = None, saturated_vnni: Bool = False, single_thread_blocking_override: Bool = False, _trace_description: StringSlice[StaticConstantOrigin] = "", target: StringSlice[StaticConstantOrigin] = "cpu"](c: LayoutTensor[c.dtype, c.layout, c.origin, element_layout=c.element_layout, layout_int_type=c.layout_int_type, linear_idx_type=c.linear_idx_type, masked=c.masked, alignment=c.alignment], a: LayoutTensor[a.dtype, a.layout, a.origin, element_layout=a.element_layout, layout_int_type=a.layout_int_type, linear_idx_type=a.linear_idx_type, masked=a.masked, alignment=a.alignment], b: LayoutTensor[b.dtype, b.layout, b.origin, element_layout=b.element_layout, layout_int_type=b.layout_int_type, linear_idx_type=b.linear_idx_type, masked=b.masked, alignment=b.alignment], ctx: Optional[DeviceContext])

matmul[transpose_a: Bool = False, transpose_b: Bool = False, elementwise_lambda_fn: Optional[elementwise_epilogue_type] = None, elementwise_compute_lambda_fn: Optional[elementwise_compute_lambda_type] = None, _trace_description: StringSlice[StaticConstantOrigin] = "", target: StringSlice[StaticConstantOrigin] = "cpu"](c: TileTensor[c.dtype, c.LayoutType, c.origin, address_space=c.address_space, linear_idx_type=c.linear_idx_type, element_size=c.element_size], a: TileTensor[a.dtype, a.LayoutType, a.origin, address_space=a.address_space, linear_idx_type=a.linear_idx_type, element_size=a.element_size], b: TileTensor[b.dtype, b.LayoutType, b.origin, address_space=b.address_space, linear_idx_type=b.linear_idx_type, element_size=b.element_size], ctx: Optional[DeviceContext])

TileTensor overload. Routes directly to GPU dispatch, avoiding intermediate NDBuffer construction.

matmul[transpose_a: Bool = False, transpose_b: Bool = False, b_packed: Bool = False, elementwise_lambda_fn: Optional[elementwise_epilogue_type] = None, elementwise_compute_lambda_fn: Optional[elementwise_compute_lambda_type] = None, saturated_vnni: Bool = False, single_thread_blocking_override: Bool = False, _trace_description: StringSlice[StaticConstantOrigin] = "", target: StringSlice[StaticConstantOrigin] = "cpu"](c: TileTensor[c.dtype, c.LayoutType, c.origin, linear_idx_type=c.linear_idx_type, element_size=c.element_size], a: TileTensor[a.dtype, a.LayoutType, a.origin, linear_idx_type=a.linear_idx_type, element_size=a.element_size], b: TileTensor[b.dtype, b.LayoutType, b.origin, linear_idx_type=b.linear_idx_type, element_size=b.element_size], ctx: DeviceContextPtr = DeviceContextPtr())

TileTensor overload of matmul. Converts to NDBuffer and delegates.

matmul[transpose_a: Bool = False, transpose_b: Bool = False, b_packed: Bool = False, elementwise_lambda_fn: Optional[elementwise_epilogue_type] = None, elementwise_compute_lambda_fn: Optional[elementwise_compute_lambda_type] = None, saturated_vnni: Bool = False, single_thread_blocking_override: Bool = False, _trace_description: StringSlice[StaticConstantOrigin] = "", target: StringSlice[StaticConstantOrigin] = "cpu"](c: TileTensor[c.dtype, c.LayoutType, c.origin, linear_idx_type=c.linear_idx_type, element_size=c.element_size], a: TileTensor[a.dtype, a.LayoutType, a.origin, linear_idx_type=a.linear_idx_type, element_size=a.element_size], b: TileTensor[b.dtype, b.LayoutType, b.origin, linear_idx_type=b.linear_idx_type, element_size=b.element_size], ctx: Optional[DeviceContext])

TileTensor overload of matmul. Converts to NDBuffer and delegates.

matmul[transpose_a: Bool = False, transpose_b: Bool = False, b_packed: Bool = False, elementwise_lambda_fn: Optional[elementwise_epilogue_type] = None, elementwise_compute_lambda_fn: Optional[elementwise_compute_lambda_type] = None, saturated_vnni: Bool = False, single_thread_blocking_override: Bool = False, _trace_description: StringSlice[StaticConstantOrigin] = "", target: StringSlice[StaticConstantOrigin] = "cpu"](c: NDBuffer[c.dtype, c.origin, c.shape, DimList.create_unknown[2]()], a: NDBuffer[a.dtype, a.origin, a.shape, DimList.create_unknown[2]()], b: NDBuffer[b.dtype, b.origin, b.shape, DimList.create_unknown[2]()], ctx: DeviceContextPtr = DeviceContextPtr())

matmul[transpose_a: Bool = False, transpose_b: Bool = False, b_packed: Bool = False, elementwise_lambda_fn: Optional[elementwise_epilogue_type] = None, elementwise_compute_lambda_fn: Optional[elementwise_compute_lambda_type] = None, saturated_vnni: Bool = False, single_thread_blocking_override: Bool = False, _trace_description: StringSlice[StaticConstantOrigin] = "", target: StringSlice[StaticConstantOrigin] = "cpu"](c: NDBuffer[c.dtype, c.origin, c.shape, DimList.create_unknown[2]()], a: NDBuffer[a.dtype, a.origin, a.shape, DimList.create_unknown[2]()], b: NDBuffer[b.dtype, b.origin, b.shape, DimList.create_unknown[2]()], ctx: Optional[DeviceContext])

Was this page helpful?