chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,20 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
from .default import *
|
||||
from .utils import *
|
||||
from .with_bias_scale import *
|
||||
@@ -0,0 +1,89 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Implementation of default unary operator dispatches."""
|
||||
|
||||
from tvm.tirx import FloatImm, PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
|
||||
from tvm.tirx.operator.tile_primitive.common import MapOpType
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import init_analyzer
|
||||
from ..instruction_generator import InstructionGenerator
|
||||
from .utils import (
|
||||
const_input_ops,
|
||||
generate_unary_func,
|
||||
non_activation_unary_map_ops,
|
||||
try_find_inst_unary,
|
||||
)
|
||||
|
||||
|
||||
def unary_trn(op: TilePrimitiveCall, unary_op: MapOpType, sctx: DispatchContext) -> PrimFunc | None:
|
||||
"""Schedule unary operation on Trainium."""
|
||||
# Check execution environment
|
||||
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
|
||||
fail("requires Trainium target and thread exec_scope")
|
||||
|
||||
# Extract operation arguments
|
||||
dst_buffer_region, _src = op.args
|
||||
|
||||
# Handle constant or buffer source
|
||||
if isinstance(_src, FloatImm):
|
||||
if unary_op not in const_input_ops:
|
||||
assert False, f"Unsupported unary operation {unary_op} taking const as input"
|
||||
CONST = _src
|
||||
src_buffer_region = None
|
||||
else:
|
||||
CONST = None
|
||||
src_buffer_region = _src
|
||||
|
||||
# Initialize analyzer and validate operation type
|
||||
analyzer = init_analyzer(sctx)
|
||||
assert unary_op in non_activation_unary_map_ops, f"Unsupported unary operation {unary_op}"
|
||||
|
||||
inst_gen = InstructionGenerator([dst_buffer_region, _src], analyzer)
|
||||
# Find instruction parameters
|
||||
if CONST is None:
|
||||
inst_repr = try_find_inst_unary(dst_buffer_region, src_buffer_region, analyzer, inst_gen)
|
||||
else:
|
||||
inst_repr = try_find_inst_unary(dst_buffer_region, dst_buffer_region, analyzer, inst_gen)
|
||||
# Generate and return the implementation function
|
||||
return generate_unary_func(
|
||||
dst_buffer_region,
|
||||
_src,
|
||||
inst_gen,
|
||||
inst_repr,
|
||||
unary_op,
|
||||
None, # No bias
|
||||
None, # No scale
|
||||
analyzer,
|
||||
op.workspace,
|
||||
op.config,
|
||||
sctx,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration: bind each default unary op name to its TRN schedule candidates.
|
||||
# ---------------------------------------------------------------------------
|
||||
from tvm.tirx.operator.tile_primitive import register_dispatch # noqa: E402
|
||||
|
||||
for _op_name, _op_type in {"reciprocal": MapOpType.RECIPROCAL, "memset": MapOpType.FILL}.items():
|
||||
|
||||
@register_dispatch(_op_name, "trn", variant="unary", priority=0)
|
||||
def _unary_dispatch(op, sctx, _ty=_op_type):
|
||||
return unary_trn(op, _ty, sctx)
|
||||
@@ -0,0 +1,190 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Shared helpers, op tables, and validation functions for unary operator dispatches."""
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.backend.trn.layout import is_trainium_layout
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import BufferRegion, FloatImm
|
||||
from tvm.tirx.operator.tile_primitive.common import MapOpType
|
||||
|
||||
from ..common import nki_dim
|
||||
from ..dim_utils import get_ewise_dim_map
|
||||
from ..instruction_generator import InstructionGenerator
|
||||
from ..workspace_utils import check_workspace_buffer
|
||||
|
||||
# Operation type classifications
|
||||
non_activation_unary_map_ops = [MapOpType.RECIPROCAL, MapOpType.FILL]
|
||||
activation_map_ops = [MapOpType.SQRT, MapOpType.EXP]
|
||||
|
||||
# Operation code table for instructions
|
||||
opcode_table = {MapOpType.SQRT: "sqrt", MapOpType.EXP: "exp"}
|
||||
|
||||
# Operations that take constants as input
|
||||
const_input_ops = [MapOpType.FILL]
|
||||
|
||||
|
||||
def try_find_inst_unary(
|
||||
dst_buffer_region: BufferRegion,
|
||||
src_buffer_region: BufferRegion,
|
||||
analyzer: Analyzer,
|
||||
inst_gen: InstructionGenerator,
|
||||
allowed_f_dim_dst: tuple[int] | None = None,
|
||||
allowed_f_dim_src: tuple[int] | None = None,
|
||||
):
|
||||
"""Find instruction parameters for a unary operation."""
|
||||
dst = dst_buffer_region.buffer
|
||||
src = src_buffer_region.buffer
|
||||
|
||||
# Validate buffer layouts and scopes
|
||||
valid_layout_scope = all(
|
||||
[
|
||||
src.layout and dst.layout,
|
||||
src.scope() in ("trn.sbuf", "trn.psum"),
|
||||
dst.scope() == "trn.sbuf",
|
||||
is_trainium_layout(src.layout),
|
||||
is_trainium_layout(dst.layout),
|
||||
]
|
||||
)
|
||||
|
||||
if not valid_layout_scope:
|
||||
assert False, (
|
||||
f"scope or layout mismatch, src: {src_buffer_region}, dst: {dst_buffer_region}"
|
||||
)
|
||||
|
||||
# Extract and validate dimensions
|
||||
dst_region = dst_buffer_region.region
|
||||
src_region = src_buffer_region.region
|
||||
|
||||
dst_extent = [r.extent for r in dst_region]
|
||||
src_extent = [r.extent for r in src_region]
|
||||
|
||||
dst_extent_nonunit = [e for e in dst_extent if e != 1]
|
||||
src_extent_nonunit = [e for e in src_extent if e != 1]
|
||||
|
||||
# Verify dimensions match
|
||||
dims_match = len(src_extent_nonunit) == len(dst_extent_nonunit) and all(
|
||||
analyzer.can_prove_equal(s, d) for s, d in zip(src_extent_nonunit, dst_extent_nonunit)
|
||||
)
|
||||
|
||||
if not dims_match:
|
||||
assert False, (
|
||||
f"shape or dimension mismatch, src: {src_buffer_region}, dst: {dst_buffer_region}"
|
||||
)
|
||||
dim_map = get_ewise_dim_map(src_buffer_region, dst_buffer_region, analyzer)
|
||||
inst_gen.link_buffer_regions(src_buffer_region, dst_buffer_region, dim_map)
|
||||
# Find optimal instruction parameters
|
||||
inst_repr = inst_gen.find_max_inst_size_from_one_region(dst_buffer_region, allowed_f_dim_dst)
|
||||
inst_repr = inst_gen.fit_inst_tile_to_region(inst_repr, src_buffer_region, allowed_f_dim_src)
|
||||
return inst_repr
|
||||
|
||||
|
||||
def get_const_bias_tensor(bias, shape, dtype, workspace, sctx):
|
||||
"""Create or retrieve a constant bias tensor."""
|
||||
if "const_bias" not in workspace:
|
||||
assert sctx.alloc_only, (
|
||||
"Constant bias tensor must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
|
||||
)
|
||||
# Create new bias buffer
|
||||
bias_buffer = T.buffer(shape, dtype, scope="trn.sbuf", buffer_name="const_bias")
|
||||
sctx.add_alloc_buffer(bias_buffer)
|
||||
|
||||
@T.prim_func
|
||||
def const_bias_init():
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(0, shape[0], annotations={nki_dim: "P"}):
|
||||
for f_loop in T.serial(0, shape[1], annotations={nki_dim: "F"}):
|
||||
T.evaluate(T.nki.memset(bias_buffer[p_loop, f_loop], bias))
|
||||
T.tvm_kernel_replace_point()
|
||||
|
||||
sctx.add_init_stmt(const_bias_init.body)
|
||||
else:
|
||||
# Use existing bias buffer
|
||||
bias_buffer = workspace["const_bias"]
|
||||
check_workspace_buffer(bias_buffer, shape, "trn.sbuf")
|
||||
|
||||
return bias_buffer
|
||||
|
||||
|
||||
def generate_unary_func(
|
||||
dst_buffer_region,
|
||||
_src,
|
||||
inst_gen: InstructionGenerator,
|
||||
inst_repr,
|
||||
unary_op,
|
||||
bias,
|
||||
scale,
|
||||
analyzer,
|
||||
workspace,
|
||||
config,
|
||||
sctx,
|
||||
):
|
||||
"""Generate a function that implements a unary operation."""
|
||||
# Prepare parameters
|
||||
p_size = dst_buffer_region.buffer.layout.size("P")
|
||||
|
||||
# Apply instruction size limits if specified
|
||||
inst_size_limit = config.get("max_inst_size", 512)
|
||||
inst_repr.bound_inst_size(inst_size_limit, analyzer)
|
||||
|
||||
f_var = T.Var("F", "int32")
|
||||
p_var = T.Var("P", "int32")
|
||||
b_var = T.Var("B", "int32")
|
||||
inst_gen.bind_inst_iter(dst_buffer_region, f_var, inst_repr.size, inst_repr.stride, True)
|
||||
inst_gen.bind_inst_iter(dst_buffer_region, p_var, p_size, 1, False)
|
||||
b_extent = inst_gen.fill_in_block_dim(dst_buffer_region, b_var)
|
||||
|
||||
# Get operation code if available
|
||||
opcode = opcode_table.get(unary_op, None)
|
||||
|
||||
# Extract buffers
|
||||
dst = dst_buffer_region.buffer
|
||||
src = _src.buffer if isinstance(_src, BufferRegion) else None
|
||||
|
||||
# Handle bias tensor
|
||||
if isinstance(bias, FloatImm | float):
|
||||
bias_buffer = get_const_bias_tensor(
|
||||
bias, (p_size, inst_repr.size), dst.dtype, workspace, sctx
|
||||
)
|
||||
elif isinstance(bias, BufferRegion):
|
||||
bias_buffer = bias.buffer
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def impl():
|
||||
for b_loop in T.serial(0, b_extent):
|
||||
with T.attr(0, "tensorized_nki_instruction", 1):
|
||||
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
|
||||
for f_loop in T.serial(0, inst_repr.size, annotations={nki_dim: "F"}):
|
||||
inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, b_var: b_loop})
|
||||
dst_indices = T.meta_var(inst_gen.generate_indices(dst_buffer_region))
|
||||
if inst_gen.make_guard(dst_buffer_region):
|
||||
if unary_op == MapOpType.FILL:
|
||||
T.evaluate(T.nki.memset(dst[tuple(dst_indices)], _src))
|
||||
else:
|
||||
src_indices = T.meta_var(inst_gen.generate_indices(_src))
|
||||
if unary_op == MapOpType.RECIPROCAL:
|
||||
T.evaluate(T.nki.reciprocal(dst[tuple(dst_indices)], src[tuple(src_indices)])) # noqa: E501
|
||||
elif isinstance(bias, BufferRegion):
|
||||
bias_indices = T.meta_var(inst_gen.generate_indices(bias))
|
||||
T.evaluate(T.nki.activation(dst[tuple(dst_indices)], src[tuple(src_indices)], opcode, scale=scale, bias=bias_buffer[tuple(bias_indices)])) # noqa: E501
|
||||
else:
|
||||
T.evaluate(T.nki.activation(dst[tuple(dst_indices)], src[tuple(src_indices)], opcode, scale=scale, bias=bias_buffer[p_loop, f_loop])) # noqa: E501
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
@@ -0,0 +1,87 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Implementation of unary with bias and scale operator dispatches."""
|
||||
|
||||
from tvm.tirx import BufferRegion, PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
|
||||
from tvm.tirx.operator.tile_primitive.common import MapOpType
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..binary import try_find_inst_nary
|
||||
from ..common import init_analyzer
|
||||
from ..instruction_generator import InstructionGenerator
|
||||
from .utils import activation_map_ops, generate_unary_func, try_find_inst_unary
|
||||
|
||||
|
||||
def unary_with_bias_scale_trn(
|
||||
op: TilePrimitiveCall, unary_op: MapOpType = MapOpType.SQRT, sctx: DispatchContext = None
|
||||
) -> PrimFunc | None:
|
||||
"""Schedule unary operation with bias and scale on Trainium."""
|
||||
# Check execution environment
|
||||
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
|
||||
fail("requires Trainium target and thread exec_scope")
|
||||
|
||||
# Extract operation arguments with defaults
|
||||
dst_buffer_region, src_buffer_region, _bias, scale = op.args
|
||||
scale = 1.0 if scale is None else scale
|
||||
_bias = 0.0 if _bias is None else _bias
|
||||
|
||||
# Initialize analyzer and validate operation type
|
||||
analyzer = init_analyzer(sctx)
|
||||
assert unary_op in activation_map_ops, f"Unsupported activation operation {unary_op}"
|
||||
|
||||
# Find instruction parameters
|
||||
inst_gen = InstructionGenerator([dst_buffer_region, src_buffer_region, _bias], analyzer)
|
||||
if isinstance(_bias, BufferRegion):
|
||||
inst_repr, _, _ = try_find_inst_nary(
|
||||
dst_buffer_region,
|
||||
[src_buffer_region, _bias],
|
||||
analyzer,
|
||||
inst_gen,
|
||||
allow_first_op_tensortensor=False,
|
||||
)
|
||||
else:
|
||||
# Handle scalar bias
|
||||
inst_repr = try_find_inst_unary(dst_buffer_region, src_buffer_region, analyzer, inst_gen)
|
||||
|
||||
# Generate and return the implementation function
|
||||
return generate_unary_func(
|
||||
dst_buffer_region,
|
||||
src_buffer_region,
|
||||
inst_gen,
|
||||
inst_repr,
|
||||
unary_op,
|
||||
_bias,
|
||||
scale,
|
||||
analyzer,
|
||||
op.workspace,
|
||||
op.config,
|
||||
sctx,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration: bind each unary_with_bias_scale op name to its TRN schedule candidates.
|
||||
# ---------------------------------------------------------------------------
|
||||
from tvm.tirx.operator.tile_primitive import register_dispatch # noqa: E402
|
||||
|
||||
for _op_name, _op_type in {"sqrt": MapOpType.SQRT, "exp": MapOpType.EXP}.items():
|
||||
|
||||
@register_dispatch(_op_name, "trn", variant="unary_with_bias_scale", priority=0)
|
||||
def _unary_bs_dispatch(op, sctx, _ty=_op_type):
|
||||
return unary_with_bias_scale_trn(op, _ty, sctx)
|
||||
Reference in New Issue
Block a user