buffer

Lower the IRx buffer/view substrate as plain structs and explicit runtime helper calls without adding a scientific array API.

Classes:

• BufferVisitorMixin –

BufferVisitorMixin

Bases: VisitorMixinBase

Methods:

• lower_buffer_element_pointer –
• visit –

lower_buffer_element_pointer

lower_buffer_element_pointer(
    view: Value,
    indices: list[Value],
    element_type: DataType,
    *,
    index_nodes: list[AST],
    bounds_policy: BufferIndexBoundsPolicy = DEFAULT,
) -> Value

Source code in src/irx/builder/lowering/buffer.py

def lower_buffer_element_pointer(
    self,
    view: ir.Value,
    indices: list[ir.Value],
    element_type: astx.DataType,
    *,
    index_nodes: list[astx.AST],
    bounds_policy: BufferIndexBoundsPolicy = (
        BufferIndexBoundsPolicy.DEFAULT
    ),
) -> ir.Value:
    """
    title: Lower a buffer view indexed access to a typed element pointer.
    parameters:
      view:
        type: ir.Value
      indices:
        type: list[ir.Value]
      element_type:
        type: astx.DataType
      index_nodes:
        type: list[astx.AST]
      bounds_policy:
        type: BufferIndexBoundsPolicy
    returns:
      type: ir.Value
    """
    _ = bounds_policy
    if view.type != self._llvm.BUFFER_VIEW_TYPE:
        raise Exception("buffer view indexing requires a BufferViewType")
    if len(indices) != len(index_nodes):
        raise Exception("buffer view index lowering arity mismatch")

    element_llvm_type = self._llvm_type_for_ast_type(element_type)
    if element_llvm_type is None:
        raise Exception(
            "buffer view indexing has unsupported element type"
        )

    data = self._extract_buffer_view_data(view)
    strides = self._extract_buffer_view_strides(view)
    total_offset = self._extract_buffer_view_offset_bytes(view)

    for axis, (index, index_node) in enumerate(
        zip(indices, index_nodes, strict=True)
    ):
        index64 = self._normalize_buffer_index_value(index, index_node)
        stride_ptr = self._llvm.ir_builder.gep(
            strides,
            [ir.Constant(self._llvm.INT64_TYPE, axis)],
            name=f"irx_buffer_index_stride_ptr_{axis}",
        )
        stride = self._llvm.ir_builder.load(
            stride_ptr,
            name=f"irx_buffer_index_stride_{axis}",
        )
        scaled_index = self._llvm.ir_builder.mul(
            index64,
            stride,
            name=f"irx_buffer_index_scaled_{axis}",
        )
        total_offset = self._llvm.ir_builder.add(
            total_offset,
            scaled_index,
            name=f"irx_buffer_index_offset_{axis}",
        )

    byte_ptr = self._llvm.ir_builder.gep(
        data,
        [total_offset],
        name="irx_buffer_index_byte_ptr",
    )
    element_ptr_type = element_llvm_type.as_pointer()
    if byte_ptr.type == element_ptr_type:
        return byte_ptr
    return self._llvm.ir_builder.bitcast(
        byte_ptr,
        element_ptr_type,
        name="irx_buffer_index_element_ptr",
    )

visit

visit(node: BufferViewRelease) -> None

Source code in src/irx/builder/lowering/buffer.py

@VisitorCore.visit.dispatch
def visit(self, node: astx.BufferViewRelease) -> None:
    """
    title: Visit BufferViewRelease nodes.
    parameters:
      node:
        type: astx.BufferViewRelease
    """
    release = self.require_runtime_symbol(
        "buffer",
        "irx_buffer_view_release",
    )
    view_ptr = self._buffer_view_pointer_for_call(
        node.view,
        name="irx_buffer_release_view",
    )
    result = self._llvm.ir_builder.call(
        release,
        [view_ptr],
        name="irx_buffer_release_status",
    )
    self.result_stack.append(result)