**Changes**
Add tests in `test_frontend_tflite.py`.
Lower S`PACE_TO_BATCH_ND` / `BATCH_TO_SPACE_ND` through TOPI in
`tflite_frontend.py`.
Use tf.raw_ops.BatchToSpaceND in the test because tf.batch_to_space_nd
is not available in this TF build.
**Why the TFLite frontend changed**
The frontend was calling relax.op.nn.space_to_batch_nd /
relax.op.nn.batch_to_space_nd, which aren’t implemented in this
checkout. I updated the TFLite frontend to lower these ops via TOPI
packed calls so conversion works and the new tests can pass.
**Test:**
```
pytest test_frontend_tflite.py -k "test_space_to_batch_nd or test_batch_to_space_nd"
```
related to #18971
## Summary
This PR adds Relax TFLite frontend support for the following segment
operators from #19412:
- `SEGMENT_SUM`
- `UNSORTED_SEGMENT_MIN`
- `UNSORTED_SEGMENT_PROD`
These operators are lowered through `relax.op.scatter_nd` with the
corresponding reduction modes.
## Changes
### TFLite Frontend
1. Add TFLite converter mappings for segment operators:
- `SEGMENT_SUM` -> `scatter_nd(..., reduction="add")`
- `UNSORTED_SEGMENT_MIN` -> `scatter_nd(..., reduction="min")`
- `UNSORTED_SEGMENT_PROD` -> `scatter_nd(..., reduction="mul")`
2. Add shared segment lowering logic:
- Convert `segment_ids` into scatter indices via `expand_dims`.
- Build the output shape from `num_segments` or constant `segment_ids`.
- Initialize the scatter base tensor with the correct reduction
identity.
### Tests
Add TFLite frontend tests for:
- `test_segment_sum`
- `test_unsorted_segment_min`
- `test_unsorted_segment_prod`
Each test verifies the imported Relax IR lowers to `R.scatter_nd` with
the expected reduction mode and base tensor initialization.
## Testing
All targeted tests pass:
```bash
python -m pytest \
tests/python/relax/test_frontend_tflite.py::test_scatter_nd \
tests/python/relax/test_frontend_tflite.py::test_segment_sum \
tests/python/relax/test_frontend_tflite.py::test_unsorted_segment_min \
tests/python/relax/test_frontend_tflite.py::test_unsorted_segment_prod \
-q
```
## References
- Issue #19412: TFLite Relax frontend operator support tracking
- Related PR #19490: Adds SCATTER_ND support
This PR adds first-class `nn.ParameterList` and `nn.ParameterDict`
containers to the Relax frontend.
These containers provide PyTorch-like list/dict registration for raw
`nn.Parameter` objects while preserving Relax frontend semantics: values
must be explicit `nn.Parameter` instances, with no automatic
tensor-to-parameter conversion.
### Changes
- Add public `nn.ParameterList` and `nn.ParameterDict` exports.
- Support stable parameter names in traversal:
- `params.0`, `params.1`
- `params.foo`, `params.bar`
- Integrate the new containers with:
- `named_parameters()`
- `parameters()`
- `state_dict()`
- `load_state_dict()`
- `to(dtype=...)`
- `export_tvm()`
- `nn.Mutator`
- Add focused tests for basic container behavior, type validation,
nested traversal, export parameter names, state loading, dtype
conversion, and mutator naming.
---------
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
The Relax ops `relax.isnan`, `relax.isinf`, and `relax.isfinite` are
registered in the op registry and emit valid IR, but lack FLegalize
entries so LegalizeOps() leaves them unlowered and relax.build() crashes
at codegen with:
InternalError: CodeGenVM cannot handle this intrinsic now:
Op(relax.isnan)
Same crash on `isinf` and `isfinite`, on both LLVM and CUDA targets.
The TOPI implementations already exist (python/tvm/topi/math.py). The
fix is three register_legalize calls following the
_call_topi_without_attr pattern used by the other unary ops in the same
file.
Fixes#19452.
This PR adds support for three new operators in the Relax TFLite
frontend:` BROADCAST_TO`, `EMBEDDING_LOOKUP,` and `SELECT_V2.`
Passed all newly added unit tests using
```
pytest tests/python/relax/test_frontend_tflite.py -k "test_broadcast_to or test_embedding_lookup or test_select_v2"
```
reference #19412
This PR adds support for the `SCATTER_ND` operator in the Relax TFLite
frontend.
### Key Changes:
- Added handler `convert_scatter_nd` to parse `indices` and `updates`.
- Explicitly handles static vs dynamic shape tensor extraction via
`to_int_list` and `relax.op.tensor_to_shape`.
- Uses `relax.op.zeros` to initialize the base array based on the
`updates` precision dtype.
- Mapped `SCATTER_ND` to the corresponding `relax.op.scatter_nd()`
target.
- Registered the translator into `convert_map` and provided the matching
unit test in test_frontend_tflite.py.
### Testing:
Passed unit tests
```
pytest tests/python/relax/test_frontend_tflite.py::test_scatter_nd
```
Related to #19412
Hi Committers,
This PR is trying to fix issues
https://github.com/apache/tvm/issues/19437. Any suggestions would be
appreciated if you are available.
### Root Cause
The original CumSum converter always defaulted to axis=0 when the axis
input was a relax.Var (i.e., a runtime tensor), ignoring the actual
runtime value. This led to incorrect behavior and did not comply with
the ONNX specification.
### Solutions
Update CumSum._impl_v14 to:
- Check if the axis input is a Constant: require it to have exactly one
element, otherwise raise an error.
- If the axis input is a relax.Var, raise an error instead of always
defaulted to axis=0.
---------
Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
## Summary
This PR adds support for the TFLite `RANDOM_UNIFORM`,
`RANDOM_STANDARD_NORMAL`, and `MULTINOMIAL` operators in the Relax
TFLite frontend, covering the items I claimed in #19412.
`RANDOM_UNIFORM` and `RANDOM_STANDARD_NORMAL` are lowered to seeded
`tvm.contrib.random` calls with dynamic shape support. `MULTINOMIAL` is
lowered by composing existing Relax ops around
`relax.op.multinomial_from_uniform`.
## Changes
### Frontend
1. Add converter registrations for `RANDOM_UNIFORM`,
`RANDOM_STANDARD_NORMAL`, and `MULTINOMIAL`.
2. Add shared helpers to:
- parse TFLite `RandomOptions` seeds
- convert shape tensors to Relax shape expressions for dynamic-shape
random ops
3. Lower `RANDOM_UNIFORM` to `tvm.contrib.random.uniform`.
4. Lower `RANDOM_STANDARD_NORMAL` to `tvm.contrib.random.normal`.
5. Lower `MULTINOMIAL` by:
- applying `R.nn.softmax` to logits
- generating seeded uniform samples
- calling `R.multinomial_from_uniform`
- reshaping the result to `[batch_size, num_samples]`
### Runtime
1. Extend `src/runtime/contrib/random/random.cc` to accept seeded calls
for `tvm.contrib.random.uniform` and `tvm.contrib.random.normal`.
2. Preserve compatibility with the existing unseeded calling convention.
## Testing
All tests pass:
```bash
pytest
tests/python/relax/test_frontend_tflite.py::test_random_uniform_dynamic_shape
\
tests/python/relax/test_frontend_tflite.py::test_random_standard_normal_dynamic_shape
\
tests/python/relax/test_frontend_tflite.py::test_multinomial_dynamic_num_samples
-v
```
## References
- #19412
- Claimed items: MULTINOMIAL, RANDOM_STANDARD_NORMAL, RANDOM_UNIFORM
This PR adds TFLite frontend support for the BROADCAST_ARGS operator
which computes the broadcasted shape of two input shape vectors
Decomposes into existing Relax primitives instead of registering a new
op:
- relax.op.full + relax.op.concat align both inputs by left-padding with
1s
- relax.op.where + relax.op.maximum apply per-axis broadcast rule
Created tests cover equal-length and different-length input cases.
Validation:
python -m pytest tests/python/relax/test_frontend_tflite.py -k
broadcast_args
Addresses the BROADCAST_ARGS item under #19412.
This PR adds TFLite frontend support for the DILATE operator which
extends a tensor by inserting a padding value between existing elements
per axis according to the dilation strides.
Decomposes into existing Relax primitives instead of registering a new
op:
applied per axis:
- relax.op.reshape adds a size-1 stride-axis and merges it back after
padding
- relax.op.full builds a padding tensor with (stride - 1) values along
that axis
- relax.op.concat interleaves the padding between input elements
- relax.op.strided_slice trims the trailing pad to output size
Both static and dynamic dilations are supported.
Frontend tests use hand-rolled .tflite fixtures since DILATE has no
public TF Python emitter through tf.lite.TFLiteConverter, so the
standard verify(TestClass, Expected) pattern can't reach it. Extends
DENSIFY's fixture builders to handle BuiltinOptions2 and non-FLOAT32
tensors. _finish_tflite_model now writes the TFL3 file identifier so the
produced buffer is a valid input for tf.lite.Interpreter in the nightly
E2E path.
Validation:
python -m pytest tests/python/relax/test_frontend_tflite.py -k dilate -v
Addresses the DILATE item under #19412.
This PR adds the ATAN2 operator to the Relax TFLite frontend.
Introduces relax.op.atan2 as a new binary elementwise primitive (TOPI
broadcast op, Relax registration, legalization to topi.atan2, script
parser support) and registers ATAN2 in the TFLite convert_map. It reuses
the TIR primitive tvm::atan2 so this PR is the higher-layer plumbing.
Validation:
python -m pytest tests/python/relax/test_op_binary.py
python -m pytest tests/python/relax/test_frontend_tflite.py -k binary
Addresses the ATAN2 item under #19412.
## Motivation
The `value` attribute on ONNX `ConstantOfShape` is optional and defaults
to a zero float32 scalar of the requested shape. The converter calls
`get_numpy(attr.get("value", 0))`, which feeds the int default into
`get_numpy` → `onnx.numpy_helper.to_array(0)` → `int.HasField` and
crashes:
```
AttributeError: 'int' object has no attribute 'HasField'
```
Minimal repro:
```python
from onnx import TensorProto, helper
node = helper.make_node("ConstantOfShape", ["shape"], ["y"]) # no value attr
shape = helper.make_tensor("shape", TensorProto.INT64, [2], [2, 3])
graph = helper.make_graph([node], "g", [],
[helper.make_tensor_value_info("y", TensorProto.FLOAT, None)],
initializer=[shape])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
from tvm.relax.frontend.onnx import from_onnx
from_onnx(model) # AttributeError
```
ORT produces a `(2, 3)` float32 tensor of zeros, matching the spec
default.
## Fix
The existing converter was already structured to handle the absent-value
case: there is an `else` branch that sets `dtype = "float32"` and uses
`value = 0` in `np.full(shape, value, dtype)`, exactly the spec default.
The bug is just that `get_numpy(0)` raises before that branch is
reached.
Dispatch on attribute presence first:
```python
attr_value = attr.get("value")
value = get_numpy(attr_value) if attr_value is not None else 0
```
When `value` is supplied as a TensorProto, behavior is unchanged. When
absent, `value` becomes int `0`, the `isinstance(value, _np.ndarray)`
check falls through to `dtype = "float32"`, and `np.full(shape, 0,
"float32")` produces the spec default.
## Test plan
- [x] `pytest
tests/python/relax/test_frontend_onnx.py::test_constantofshape_default_value`
— new test covering the absent-value case passes.
- [x] Confirmed the new test fails on `upstream/main` without this patch
(same `AttributeError`).
## Summary
Bumps `3rdparty/tvm-ffi` from `1fed0ae` (0.1.10) to `3c35034`
(0.1.11rc2).
Two TVM-side fixes are needed for the bump:
- `include/tvm/ir/type.h`: drop redundant `span` field registration
from `TupleTypeNode`, `FuncTypeNode`, and `TensorMapTypeNode`.
tvm-ffi now warns when a child re-registers an ancestor field;
`span` is already registered on `TypeNode` with `SEqHashIgnore`.
- `python/tvm/relax/frontend/torch/dynamo.py`: short-circuit
`to_torch_tensor` for `torch.Tensor` items. [tvm-ffi
#517](https://github.com/apache/tvm-ffi/pull/517) added
recursive DLPack container conversion so `torch.Tensor` inputs
auto-promote `Array` elements to `torch.Tensor` on iteration; the
helper previously only handled `tvm.runtime.Tensor` / `tvm_ffi.Array`.
## Summary
`support.GetLibInfo` exposed ~30 build-time `TVM_INFO_*` strings (git
hash, LLVM/MLIR versions, every `USE_*` flag). Real callers reduce to
"is this feature enabled?" — better answered at runtime. `USE_CUDA=ON`
does not mean a CUDA device is loadable; runtime discovery is the actual
signal. Git versioning is already tracked via `tvm.__version__`.
Changes:
- Delete `src/support/libinfo.cc`, `cmake/modules/LibInfo.cmake`,
`tests/lint/check_cmake_options.py`, and the `check-cmake-options`
pre-commit hook.
- Delete `tvm.support.libinfo()` (no shim — callers migrate to runtime
discovery).
- Add `tvm.support.detect_active_modules()` which queries the FFI global
function registry for `ffi.Module.create.<kind>` registrations (cuda,
vulkan, opencl). `describe()` now prints active runtimes instead of
CMake build flags.
- Migrate 5 in-tree callers: `_get_targets()` uses `cudnn.exists()` /
`tvm.runtime.enabled()` for CUDNN and Hexagon; `_cmake_flag_enabled()`
is rewritten as a static map from cmake flag names to
`tvm.runtime.enabled()` or FFI-registry probes; `clml_sdk_version()`
uses the existing `relax.get_openclml_version` FFI global;
`test_clml_ops.py` uses the new helper.
After this PR: `src/support/` is header-only.
## Summary
Restructure TVMScript to be dialect-agnostic at the script-core layer
while letting each extension dialect (TIRX, Relax) own its own
per-dialect script subtree. IR is below script in the dependency
stack and is NOT a peer dialect — its script handlers stay in the
shared core.
This PR folds together two coupled refactors that were initially
opened as separate PRs (#19478 and the original #19479); they
share rename / relocation surface so they ship as one cohesive
change.
## What this PR does
### Per-dialect script subtree (originally #19479)
- Moves per-dialect printer + builder from
`src/script/{printer,ir_builder}/{tirx,relax}/` to
`src/{tirx,relax}/script/{printer,builder}/`.
- Tightens `src/script/*.cc` CMake glob to the dialect-free core.
- Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry
(`script.ir_builder.decl_function.<type-key>`); removes
cross-dialect includes from the shared core.
- Adds `tvm.script.register_dialect` API + `__getattr__` + a
`sys.meta_path` finder for Python-side dialect discovery.
In-tree dialects (tirx, relax) registered centrally in
`python/tvm/__init__.py`.
- Drops the obsolete static re-export shims at
`python/tvm/script/{parser,ir_builder}/{tirx,relax}/`.
### Dialect-agnostic printer config (originally #19478)
- Relocates `include/tvm/ir/script_printer.h` →
`include/tvm/script/printer/config.h` next to the rest of the
printer's public surface. The header is not IR-specific.
- Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for
clarity (the macro registers Script as the kRepr callback +
per-type vtable dispatch). Aligns with the `TVM_REGISTER_*`
family.
- Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`,
`relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor
of a generic `ffi::Map<String, Any> extra_config` keyed by
`"<dialect>.<knob>"`. Each call site reads via the templated
accessor `config->GetExtraConfig<T>("...", default)`.
- Promotes `std::string` config fields to `ffi::String`.
After this lands, the script-printer core knows nothing specific
about any dialect — new dialects plug in via the registry pattern
with zero core edits. Public Python API surface unchanged.
## Summary
This PR continues the TFLite frontend work tracked in #18971 for
`STRIDED_SLICE` and `SPLIT_V`.
Since the dynamic `FILL` / `SPLIT_V` partial-implementation work has
already been handled separately in #19433, this PR focuses on the
remaining pieces in this branch:
- fixing negative-stride `STRIDED_SLICE` conversion in the TFLite
frontend
- adding regression coverage for `STRIDED_SLICE` and static `SPLIT_V`
Relates to #18971.
## Changes
1. **`STRIDED_SLICE` negative-stride fix**
- Update the TFLite frontend `convert_strided_slice` handling of
`end_mask` when `stride < 0`.
- Use an exclusive lower bound compatible with Relax slicing semantics
so reverse slices like `x[::-1]` include index `0` correctly.
2. **TFLite frontend test coverage**
- Add `test_strided_slice_stride` to cover non-unit stride handling.
- Add `test_strided_slice_negative_stride` to cover reverse slicing with
negative strides.
- Add `test_split_v_static` to cover static `SPLIT_V` conversion.
3. **Scope clarification**
- Keep this PR focused on the remaining `STRIDED_SLICE` / `SPLIT_V` work
from #18971.
- Exclude the dynamic `FILL` / `SPLIT_V` changes that are already
addressed in #19433.
## Testing
```bash
python -m pytest -n 1
tests/python/relax/test_frontend_tflite.py::test_split_v_static -q
python -m pytest -n 1
tests/python/relax/test_frontend_tflite.py::test_strided_slice_stride -q
python -m pytest -n 1
tests/python/relax/test_frontend_tflite.py::test_strided_slice_negative_stride
-q
```
## Result
- The added STRIDED_SLICE and SPLIT_V tests passed locally.
- The negative-stride STRIDED_SLICE path now matches Relax slicing semantics.
## Why
This refactor reshapes each backend into a self-contained
`src/target/<X>/`
cluster (with optional `src/target/<X>/llvm/` for LLVM-dependent
codegen) and
introduces a per-backend fallback module that absorbs the cross-compile
role
cleanly — without `target/opt/` stubs, without `DeviceSourceModuleNode`,
and
without leaking a synthetic `kind()` to consumers.
## High-level principles
- **One directory per backend.** All codegen-side files for backend
`<X>` live
under `src/target/<X>/`. Optional `src/target/<X>/llvm/` subdir for
files
that require `USE_LLVM` at build time. Backend grouping wins over
build-dependency grouping (the latter being upstream's `target/llvm/`).
- **Plugin-only runtime modules.** `src/runtime/<X>/<X>_module.h` is
deleted.
The runtime's real `<X>ModuleNode` is reachable only via the FFI
registry
(`ffi.Module.create.<kind>`, `ffi.Module.load_from_bytes.<kind>`). No
C++ API surface other than the static registrations.
- **Per-backend fallback module for cross-compile.** Each `<X>` gets a
`<X>FallbackModuleNode` in `src/target/<X>/<X>_fallback_module.{h,cc}`.
Same `kind()` as the real module. Codegen-time only — never reachable
via
load. `GetFunction` errors with a backend-specific "runtime not linked"
message; `InspectSource` works.
- **Codegen-side wrapper does the fallback selection.** Codegen calls
`<X>ModuleCreateWithFallback(...)`, which tries
`ffi.Module.create.<kind>`
via the registry; on miss, falls through to `<X>FallbackModuleCreate`
(plain C++; reachable directly from the fallback header). When
`USE_<X>=ON`
is in effect, the registry hit returns the real module; when
`USE_<X>=OFF`, the fallback is what codegen gets. No CMake `if/else`
gating; fallback always compiled.
## Specific changes
### New per-backend directories (codegen + fallback)
- `src/target/cuda/` — `codegen_cuda.cc` + `intrin_rule_cuda.cc` +
fallback module pair + `llvm/codegen_nvptx.cc`
- `src/target/rocm/` — fallback module pair + `llvm/codegen_amdgpu.cc` +
`llvm/intrin_rule_rocm.cc`
- `src/target/hexagon/` — fallback module pair +
`llvm/codegen_hexagon.cc` + `llvm/intrin_rule_hexagon.cc`
- `src/target/metal/` — `codegen_metal.cc` + `intrin_rule_metal.cc` +
fallback module pair
- `src/target/vulkan/` — `build_vulkan.cc` + the rest of `target/spirv/`
absorbed + fallback module pair
- `src/target/opencl/` — `codegen_opencl.cc` + `intrin_rule_opencl.cc` +
fallback module pair
- `src/target/webgpu/` — `codegen_webgpu.cc` + fallback module pair (was
`WebGPUSourceModuleNode`, renamed)
### New fallback classes
`<X>FallbackModuleNode` for X in {`CUDA`, `ROCm`, `Hexagon`, `Metal`,
`Vulkan`, `OpenCL`, `WebGPU`}. Each:
- `kind()` matches the real backend
- Stores `(code or smap, fmt, fmap, source)` — no driver/runtime calls
- `GetFunction` errors with backend-specific "runtime not linked"
message
- `InspectSource` works fully
- `SaveToBytes` byte-identical to real
Hi Committers,
This PR is trying to fix issues
https://github.com/apache/tvm/issues/19419. Any suggestions would be
appreciated if you are available.
### Root Cause
- auto-detected CUDA might lacks `max_shared_memory_per_block` and it
would cause `KeyError`
### Solutions
- Add layered fall back strategy to handle missing attr
`max_shared_memory_per_block`
---------
Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
## Summary
This PR adds test coverage for the TFLite DENSIFY operator as requested
in issue #18971 , and fixes several related bugs in the TFLite frontend.
DENSIFY converts sparse weight tensors to dense format at **conversion
time** (not runtime). The dense weights become constants in the output
IR via the `prefetched_nodes` mechanism.
## Changes
### Bug Fixes
1. **`convert_op_to_relax`**: Check `ret is None` before
`normalize(ret)` to avoid crash when DENSIFY returns `None`.
2. **`get_tensor_expr`**: Add `is_prefetched()` check before
`get_tensor_value()` to handle DENSIFY outputs with empty buffers.
3. **`convert_fully_connected`**: Add prefetched weight handling
4. **`convert_transpose_conv`**: Add prefetched weight handling
### Tests
Four test cases covering different DENSIFY usage scenarios:
| Test | Downstream Op | Purpose |
|------|---------------|---------|
| `test_densify` | None | Basic DENSIFY to constant |
| `test_densify_with_add` | ADD | Prefetched as regular input |
| `test_densify_with_conv2d` | CONV2D | Network-level test (2D conv) |
| `test_densify_with_fully_connected` | FULLY_CONNECTED | Network-level
test (FC layer) |
**Note**: Sparse TFLite models are built manually using flatbuffers API
(TensorFlow does not provide an API for creating sparse models).
## Testing
All tests pass:
```bash
pytest tests/python/relax/test_frontend_tflite.py::test_densify \
tests/python/relax/test_frontend_tflite.py::test_densify_with_add \
tests/python/relax/test_frontend_tflite.py::test_densify_with_conv2d \
tests/python/relax/test_frontend_tflite.py::test_densify_with_fully_connected -v
```
- `test_densify` PASSED
- `test_densify_with_add` PASSED
- `test_densify_with_conv2d` PASSED
- `test_densify_with_fully_connected` PASSED
## References
Issue #18971 : TFLite operator test coverage tracking
Related: #19408 (MATRIX_DIAG, MATRIX_SET_DIAG, SPARSE_TO_DENSE tests)
Adds the missing TFLite NonMaxSuppressionV4 frontend handler. The
underlying relax.op.vision.non_max_suppression already covers V4's
behavior with soft_nms_sigma at the default 0.0 (hard-NMS path). The
handler bridges TFLite's tensor format to the Relax op, following the
same pattern as convert_nms_v5 (#19426) but without its soft-NMS
branching.
Tests cover conversion and IR structural assertions, run with
`pytest tests/python/relax/test_frontend_tflite.py -k nms_v4`. E2E
correctness runs on the nightly gate (CI_ENV_NIGHTLY).
Relates to #19412.
This PR adds TFLite frontend support for the BITCAST operator which
reinterprets a tensor's bytes as a different dtype without converting
the underlying data.
The handler lowers BITCAST to relax.op.memory.view which aliases the
input buffer with the new shape and dtype.
Frontend tests cover same-width (float32 -> int32, uint8 -> int8),
width-changing smaller (int32[3] -> int16[3, 2]), and width-changing
larger (int16[5, 2] -> int32[5]).
` python -m pytest tests/python/relax/test_frontend_tflite.py -k bitcast
-v `
Addresses the BITCAST item under #19412.
---------
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This PR fixes partial TFLite frontend support for dynamic `FILL` and
`SPLIT_V`.
Key changes:
- Allow `FILL` to accept runtime `dims` tensors by converting them with
`relax.op.tensor_to_shape` before calling `relax.op.full`.
- Allow `SPLIT_V` to accept runtime `size_splits` tensors by decomposing
the
op into `cumsum` and `dynamic_strided_slice`.
- Use the TFLite output tuple arity as the source of truth for dynamic
`SPLIT_V`, instead of relying on `size_splits` static shape information.
- Add TFLite frontend regression tests covering dynamic `FILL` import
and
dynamic `SPLIT_V` import/compile behavior.
This addresses the `FILL` and `SPLIT_V` items under the "Fix partial
implementations" section of #19412.
Validation:
```bash
python -m ruff check python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.py
```
```bash
python -m pytest tests/python/relax/test_frontend_tflite.py \
-k "split_v_dynamic or fill_dynamic_dims" -v
```
Result:
- All checks passed
## Summary
This PR completes the TFLite `NON_MAX_SUPPRESSION_V5` implementation in
Relax by adding support for `soft_nms_sigma != 0`.
It extends `relax.vision.non_max_suppression` with soft-NMS attributes,
updates the TFLite frontend to consume the soft-NMS outputs correctly,
and aligns the TOPI implementation with LiteRT's reference behavior.
Relates to #19412.
## Changes
1. **Relax / TOPI soft-NMS support**
- Extend `NonMaximumSuppressionAttrs` with `soft_nms_sigma` and
`score_threshold`.
- Thread the new attributes through Relax op registration, Python
wrapper, and legalization.
- Add soft-NMS handling to TOPI classic NMS so
`relax.vision.non_max_suppression` can represent the
`NON_MAX_SUPPRESSION_V5` behavior.
2. **TFLite frontend support for `NON_MAX_SUPPRESSION_V5`**
- Remove the previous `soft_nms_sigma != 0` unsupported-path guard in
the TFLite frontend.
- Forward `soft_nms_sigma` and `score_threshold` into
`relax.vision.non_max_suppression`.
- Handle the soft-NMS return path explicitly so the frontend reads
decayed scores from the processed NMS output instead of re-reading the
original score tensor.
3. **Soft-NMS correctness fixes**
- Fix the soft-NMS path so boxes whose scores fall below the threshold
after decay are invalidated consistently.
- Keep returned indices and decayed scores aligned in both the TOPI TIR
implementation and the NumPy reference implementation.
- Update the soft-NMS candidate selection logic to re-pick the current
best candidate after each decay step, matching LiteRT's
reference behavior.
- Align the Gaussian decay formula with LiteRT.
4. **Test coverage**
- Add Relax tests for soft-NMS struct-info inference and legalization.
- Add Relax E2E tests covering reordered outputs after score decay and
other soft-NMS follow-up cases.
- Add TFLite frontend tests for `NON_MAX_SUPPRESSION_V5` with
`soft_nms_sigma != 0`.
- Add IR checks to verify that `soft_nms_sigma` and `score_threshold`
are forwarded correctly.
## Testing
```bash
python -m pytest -n 1 tests/python/relax/test_op_vision.py -k "all_class_non_max_suppression or get_valid_counts or nms" -v
python -m pytest tests/python/relax/test_frontend_tflite.py -k "nms_v5" -v
```
## Result:
- Relax vision tests passed locally
- TFLite `NON_MAX_SUPPRESSION_V5` coverage added for both hard-NMS and
soft-NMS paths
Phase out the legacy spin-on-global-memory CUDA barrier machinery
(`tvm_global_barrier_state` / `__tvm_prepare_global_barrier` / the
`tvm_global_barrier_kinit()` builtin and the
`tirx.detect_global_barrier`
pass-config option). CUDA's native cooperative groups / grid sync
primitives cover the use case better; the bespoke implementation has
been dead in the active codegen pipelines.
This is a deletion-only refactor across 10 files (~−264 lines net):
- Public symbol constants in `include/tvm/runtime/device_api.h`
- TIR builtin op `tvm_global_barrier_kinit()`
(`include/tvm/tirx/builtin.h`,
`src/tirx/op/builtin.cc`)
- Pass config `tirx.detect_global_barrier` (`src/tirx/ir/transform.cc`)
- Entire kGlobal branch of `s_tir::ThreadSync` including
`InitGlobalBarrier`, `MakeGlobalBarrier`, and supporting state in
`ThreadSyncInserter`
- `CUDAPrepGlobalBarrier` runtime class + `CUDAModuleNode::GetGlobal()`
- `CodeGenCUDA::PrintStorageSync` "global" branch and the
`VisitStmt_(EvaluateNode*)` override + 3 member fields
- Two Python pipeline opt-in blocks (s_tir/pipeline.py + adreno mirror)
No Python or test references to these symbols. Build clean
(260 targets), CUDA codegen 50/50 passed, tirx-base + tirx-transform
621 passed, IRF cpptests 8/8 passed.
## Motivation
Historically TVM ships a single monolithic `libtvm.so` that bundles both
the
runtime and the compiler/LLVM-heavy code paths. Deployment scenarios
that only
need the runtime end up paying the full compiler footprint (LLVM-static
dominates
the binary size), and the layout makes it awkward to install the project
under a
single Python package directory the way
`tvm_ffi`/`libinfo.load_lib_ctypes`
expects.
This PR splits the single shared library into two:
- `libtvm_runtime.so` — runtime-only symbols (loaded `RTLD_GLOBAL`).
- `libtvm_compiler.so` — compiler / LLVM / codegen, links
`libtvm_runtime.so`
publicly (loaded `RTLD_LOCAL`).
## Target restructure
- New CMake target `tvm_compiler` replaces the old `tvm` SHARED target.
- `tvm_compiler` depends on `tvm_runtime` via `target_link_libraries(...
PUBLIC tvm_runtime)`,
so anything that linked the old `tvm` now picks up the runtime
transitively.
- `tvm_libinfo_objs` (build-info TU) moved from `tvm_runtime` into
`tvm_compiler`
— it is compiler-side metadata and the runtime no longer needs it.
- All `target_link_libraries` / `target_compile_*` /
`set_target_properties` /
`tvm_ffi_add_apple_dsymutil` callsites have been rewired.
- The separate `libtvm_allvisible.so` target is **removed** (was only
consumed
by cpptests). Cpptests with private-symbol deps are deleted; remaining
cpptests now link directly against `libtvm_compiler.so` /
`libtvm_runtime.so`. `src/support/hexdump.cc` is folded into the header.
- `BUILD_DUMMY_LIBTVM` and the `BUILD_FOR_HEXAGON + USE_HEXAGON_GTEST`
cpp-test wiring are removed.
## Output and install layout
- All artifacts now go to `build/lib/` (was `build/`):
- `build/lib/libtvm_runtime.so`
- `build/lib/libtvm_compiler.so`
- Install layout is now `<package>/lib/` so
`tvm_ffi.libinfo.load_lib_ctypes`
with `package="tvm"` finds the libs in the wheel.
- CI Jenkins stash paths and `apps/hexagon_*` paths updated to the new
`build/lib/...` location.
## Python loader change
`python/tvm/base.py` now resolves the libs directly via a small
`package_lib_paths()` helper in `python/tvm/libinfo.py` (anchored on
`python/tvm/__file__`, returning the wheel `lib/`,
`<worktree>/build/lib`, and
`<worktree>/lib` candidates). Module-level `_LIB_RUNTIME`, `_LIB`, and
`_RUNTIME_ONLY` are set inline at import time:
- `libtvm_runtime.{so,dylib,dll}` loaded `RTLD_GLOBAL`.
- `libtvm_compiler.{so,dylib,dll}` loaded `RTLD_LOCAL`.
- `TVM_USE_RUNTIME_LIB` (parsed strictly: `1`/`true`/`yes`) selects
runtime-only at the loader level.
- When the compiler lib is absent, `_RUNTIME_ONLY` is set to True
automatically and `_LIB is _LIB_RUNTIME`.
## Non-obvious build-integration fixes
Three issues surfaced once both libs are loaded into the same process
and are
worth calling out:
1. **`fpA_intB_gemm` double-registration.** `fpA_intB_gemm_tvm` is an
OBJECT
library that registers a global `fastertransformer.gemm_fp16_int` at
static
init. Linking it into both `tvm_runtime` and `tvm_compiler` made the
registration run twice and trip the duplicate-registration check. Fix:
link
it (and the other runtime-only externals — `flash_attn`, NCCL, NVSHMEM,
RCCL) only into `tvm_runtime`. `tvm_compiler` picks them up via the
PUBLIC
`tvm_runtime` link.
2. **`-Wl,--no-as-needed` for minrpc.** `python/tvm/rpc/minrpc.py`
defaults
to `runtime="libtvm_runtime"` and passes `-Wl,--no-as-needed` so the
runtime static initializers actually run in the spawned minrpc binary
(without it, the linker drops the lib because no symbol is referenced
directly from the minrpc TU). minrpc does **not** link
`libtvm_compiler.so`.
3. **`testing.GetShape{Elem,Size}` moved to runtime.** Those two test
helpers
(the only `testing.*` symbols the minrpc test exercises) were registered
in
`src/support/ffi_testing.cc` (compiler-side). They are now registered in
`src/runtime/rpc/testing.cc` under `rpc.testing.GetShape{Elem,Size}` so
the minrpc server binary — runtime-only — can resolve them.
## Deprecations and breaking changes
- `BUILD_DUMMY_LIBTVM` is **removed** (option, libinfo entry, and CMake
wiring). Downstream consumers that built the dummy variant should link
`libtvm_runtime.so` directly.
- **Breaking change for downstream consumers** that read `libtvm.so` by
name:
there is no longer a `libtvm.so`. Replace with `libtvm_compiler.so`
(full)
or `libtvm_runtime.so` (runtime-only). The Vulkan device comment and a
few
test/CI comments have been updated accordingly.
- `libtvm_allvisible.so` is **removed**. Cpptests that depended on
private
out-of-line symbols have been deleted; the remaining cpp-test contract
is
documented as "public API or private header-only API only" (see
`tests/cpp/`).
- `tests/cpp-runtime/` (Hexagon + OpenCL backend tests) is **removed**
until
TVM moves to a plugin-mode backend architecture where each backend can
ship its own test harness with its own visibility scope.
## Tested
- `ninja` build: `build/lib/libtvm_runtime.so`,
`build/lib/libtvm_compiler.so`;
no `build/libtvm.so`, no `build/lib/libtvm_allvisible.so`.
`ldd build/lib/libtvm_compiler.so` links `libtvm_runtime.so`,
`libtvm_ffi.so`, `libfpA_intB_gemm.so`, `libflash_attn.so`.
- `ldd build/cpptest`: only `libtvm_compiler.so` + `libtvm_runtime.so` +
`libtvm_ffi.so` (no `libtvm_allvisible.so`).
- `./build/cpptest`: 144 / 144 tests pass across 29 suites.
- Smoke imports: full and `TVM_USE_RUNTIME_LIB=1` — both pass.
`TVM_USE_RUNTIME_LIB=0` correctly disables runtime-only mode (strict
parse).
- `tests/python/all-platform-minimal-test`: 75 passed, 77 skipped.
- `tests/python/runtime/`: 81 passed, 2 skipped (incl.
`test_rpc_return_remote_object` exercising the minrpc executable
end-to-end
via `rpc.testing.GetShape{Elem,Size}`).
- `tests/python/relax/test_vm_*.py`: 150 passed, 3 deselected
(`test_vm_multi_device.py` requires 3+ GPUs; host has 2 — env, not
regression),
2 xfailed.
- `tests/python/tirx-base/`: 273 passed, 2 skipped.
- `pre-commit` on edited files: green.
Closes#19443.
This pr refactors the BYOC tutorial for the example NPU backend so the
full pipeline (register → partition → codegen → VM execute) actually
runs and visibly demonstrates fusion.
Also picks up several latent bugs in the example backend that the
original tutorial was implicitly papering over.
## Summary
TVM-side cleanup that drops the `python/tvm/runtime/object.py` shim and
routes `tvm.runtime.Object` directly to `tvm_ffi.Object`. The
`tvm.runtime.Object` re-export is preserved (now a re-export of
`tvm_ffi.Object`) so external callers keep working.
The load-bearing `__object_repr__` install — which wires TVM IR objects
up to the rich C++ `ReprPrinter` registered through
`init_ffi_api("node", ...)` — moves into
`python/tvm/runtime/_ffi_node_api.py`.
That module is already imported as a side-effect-only module from
`python/tvm/runtime/__init__.py`, so the override fires at the right
time (after `init_ffi_api` registers the C++ printer).
`_ffi_node_api.AsRepr` itself is **kept**: `tvm_ffi`'s default repr is
primitive (`ClassName(ptr)`); TVM IR objects need the rich printer
registered via `init_ffi_api("node", ...)`. `AsRepr` is what bridges
that printer back into Python `repr(obj)` and is also the runtime-only
fallback when `libtvm.so` is unavailable.
The 7 in-tree importers of the deleted shim (plus one straggler in
`runtime/disco/session.py`) are switched to either
`from tvm.runtime import Object` or `from tvm_ffi import Object`,
depending on which pattern the file already uses.
## Test plan
- [x] `python -c "import tvm; print(repr(tvm.IRModule({})))"` produces
TVMScript-style output (rich repr preserved).
- [x] `pytest tests/python/all-platform-minimal-test/ -x` — 75 passed,
77 skipped (matches baseline).
- [x] `pytest tests/python/tirx-base/ -x` — 273 passed, 2 skipped.
- [x] `pre-commit run --files <changed files>` — all hooks pass.
- [ ] CI green.
This PR extends `_attention_sequence_prefill_with_mask` to support a
second mask regime for decoder-style embedding workloads.
### Summary
- Keep the existing right-padded bidirectional behavior as
`mask_mode="padded"`.
- Add `mask_mode="causal_padded_left"` for left-padded causal sequence
prefill.
- Add a `softmax_update_causal_padded_left` macro for the online softmax
mask.
- Add tests for causal left-padding with zero, full, mixed, and GQA
valid lengths.
### Motivation
This is a TVM-side kernel dependency for the first-class embedding
serving work tracked in mlc-ai/mlc-llm#3451.
The existing masked sequence prefill kernel supports encoder-style
batches where real tokens occupy the valid prefix `[0, valid_len)` and
padding is on the right.
Decoder-style embedding batches, such as the decoder-only embedding
path, commonly left-pad variable-length inputs so the final real token /
EOS lands at the same final column across the batch. This allows
last-token pooling to read `output[:, -1, :]`, while still requiring
causal masking within each valid suffix.
For each batch row:
- `mask_mode="padded"`: real tokens are `[0, valid_len)`.
- `mask_mode="causal_padded_left"`: real tokens are `[seq_len -
valid_len, seq_len)`, with `col <= row`.
### Testing
- `git diff --check`
- Attempted:
`python -m pytest -q
tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py -k
'causal_padded_left or valid_len_mixed'`
This commit adds frontend support for the TFLite `CUMSUM` operator by
lowering it to `relax.op.cumsum`.
Specifically, it handles:
- Extracting the `axis` parameter from a constant tensor and converting
it to an integer.
- Parsing the `exclusive` flag from `CumsumOptions` via FlatBuffers.
- Deriving the target `dtype` from the output tensor.
- Raising a `NotImplementedError` for the `reverse` flag as it is not
yet supported by the Relax op.
Tracked in apache#19412.
This patch aligns the dtype of the `apply_rope` flag used by
`llama_rope_with_position_map` with the host-side value passed through
Relax call_tir.
Previously the PrimFunc declared `apply_rope` as `T.int32`, while the
caller-side scalar value is represented as an int64 Relax PrimValue /
ShapeExpr value. This caused Relax well-formed analysis to reject the IR
with:
Argument N type mismatch: expected R.Prim("int32"), given
R.Prim(value=1)
The mismatch can be reproduced through downstream `nn.Module.export_tvm`
paths such as MLC-LLM `convert_weight` / `compile`.
This change updates:
- `llama_rope_with_position_map`: `apply_rope: T.int32` -> `T.int64`
- `PagedKVCache`: pass the split-rotary flag as `int64_t`
The C++ Analyzer::Bind() already supports allow_override, but the FFI
bridge always used the default (false). This change threads the optional
argument through the FFI layer and the Python wrapper so callers can
rebind variables without triggering an error.
Supersedes #18247. Per maintainer guidance, resubmitting as a fresh PR
due to CI workflow changes affecting old PRs.
## Summary
This PR adds an example NPU BYOC backend for Relax, including end-to-end
integration points:
- pattern registration
(`python/tvm/relax/backend/contrib/example_npu/patterns.py`)
- backend registration
(`python/tvm/relax/backend/contrib/example_npu/__init__.py`)
- codegen entrypoint
(`src/relax/backend/contrib/example_npu/codegen.cc`)
- runtime module
(`src/runtime/contrib/example_npu/example_npu_runtime.cc`)
- CMake integration (`cmake/modules/contrib/ExampleNPU.cmake`,
`CMakeLists.txt`, `cmake/modules/LibInfo.cmake`,
`src/support/libinfo.cc`)
- tutorial/docs (`docs/how_to/tutorials/byoc_npu_example.py`, README
under contrib path)
- tests (`tests/python/contrib/test_example_npu.py`)
- CI build config enablement (`tests/scripts/task_config_build_cpu.sh`)
## Review feedback addressed from #18247
- Test location under `tests/python/contrib/`
- README includes explicit enable instructions for
`USE_EXAMPLE_NPU_CODEGEN` and `USE_EXAMPLE_NPU_RUNTIME`
- README quick-start uses inline `MatmulReLU` (no import from test
module)
- Added CMake source wiring and feature flags for runtime/codegen
- Added docs tutorial under `docs/how_to/tutorials/` (not only README)
- Reorganized motivation/context section near top of README
- Extended pattern coverage to include `example_npu.softmax` with tests
## Validation
Local checks run:
- `pre-commit` on touched files (pass)
- `PYTHONPATH=python python -m pytest -q
tests/python/contrib/test_example_npu.py` (pass)
## Notes
This backend is an example/tutorial implementation (CPU-emulated)
intended to document modern NPU-oriented BYOC integration patterns and
provide a reference path for future hardware-specific backends.
## Problem
#19413 registered `REDUCE_ANY` / `REDUCE_ALL` as `_convert_reduce` with
`relax.op.max` / `relax.op.min`. These TFLite ops are bool-only (per TFL
op schema: `TFL_ReduceAnyOp` / `TFL_ReduceAllOp` take and return
`TFL_BoolTensor`), and `relax.op.max` / `relax.op.min` are not defined
on bool, so any real model using these ops fails at compile time with:
```
Cannot decide min_value for type bool
Cannot decide max_value for type bool
```
The existing structural-equality test passed because it never attempted
to compile the converted module (E2E is gated on `CI_ENV_NIGHTLY`).
## Fix
Introduce a dedicated `_convert_reduce_bool` handler that casts the
input to int8, reduces with max/min, and casts back to bool. Update the
test to compile the expected module so this lowering is exercised
without `CI_ENV_NIGHTLY`.
## Testing
Verified compile + VM-run (TF converter → Relax → LLVM) across the full
shape / axes / keepdims matrix from `test_reduction_bool_ops`: 12 cases,
all PASS.
Follow-up to #19413.
## Summary
Adds TFLite frontend support for `REDUCE_ANY` and `REDUCE_ALL` (B-group
item in #19412).
## Lowering
`REDUCE_ANY` / `REDUCE_ALL` take bool tensors and compute logical OR /
AND along the given axes. Max / min on a bool tensor produces the same
result, so no new Relax op is required:
- `REDUCE_ANY` → `relax.op.max`
- `REDUCE_ALL` → `relax.op.min`
Both reuse the existing `_convert_reduce` handler (shared with
`REDUCE_MAX`, `REDUCE_MIN`, `REDUCE_PROD`, `MEAN`, `SUM`), with entries
added alphabetically to `convert_map`.
## Testing
Added `test_reduction_bool_ops` in
`tests/python/relax/test_frontend_tflite.py`, parametrized over the same
shape / axes / keepdims matrix as the existing `test_reduction_ops` (24
combinations total). Verified locally by running the test function
directly across all parametrizations.
Refs #19412.
This PR partially implements test coverage requested in issue #18971 for
Relax TFLite frontend operator tests.
## Bug Fix
The TFLite frontend converters for `MATRIX_DIAG`, `MATRIX_SET_DIAG`, and
`SPARSE_TO_DENSE` were broken due to calling non-existent Relax ops:
- `relax.op.matrix_set_diag` - never registered in Relax
- `relax.op.sparse_to_dense` - never registered in Relax
These ops only exist as TOPI packed functions (`topi.matrix_set_diag`,
`topi.sparse_to_dense`).
**Fix:** Replace direct op calls with `call_dps_packed` to invoke the
TOPI packed functions:
- `convert_matrix_diag`: zeros + call_dps_packed("topi.matrix_set_diag",
...)
- `convert_matrix_set_diag`: call_dps_packed("topi.matrix_set_diag",
...)
- `convert_sparse_to_dense`: call_dps_packed("topi.sparse_to_dense",
...)
Refs: #18971
This PR fixes RPC tensor cleanup for tensors returned from remote calls.
When a remote function returns a `Tensor`, the RPC protocol sends both:
- the remote backing data pointer
- the remote tensor object handle used for deletion
Previously, `TensorFromRemoteOpaqueHandle` stored only the data pointer
and called
`FreeHandle(space_.data)` during local tensor destruction. That is
incorrect:
`FreeHandle` is meant for remote object handles, not raw data-space
pointers.
This could lead to invalid cleanup behavior and crashes during teardown
in RPC workflows, including the cross-compilation + RPC tutorial
scenario reported in #18923.
This change:
- stores the remote tensor object handle in `RemoteSpace`
- calls `FreeHandle(remote_tensor_handle)` during tensor destruction
- keeps cleanup fault-tolerant if the remote connection is already
closed
Pure refactor — does not change any generated TIR / kernel behavior.
Dedupe the tiled prefill kernels in kv_cache.py by extracting the shared
online-softmax pieces as T.macro helpers (init_states, compute_s_gemm,
softmax_update_{causal,valid_length}, compute_o_gemm,
advance_tile_batch, paged_store_output_lse), plus Python helpers for the
common buffer allocations (softmax state, MHA/MLA Q/K/V/O, tile-walk
scalars).
Split the kernel factories out of kv_cache.py into private sibling
modules: _kernel_common.py (shared helpers + macros + schedule),
_page_kernels.py (append/debug/copy/compact), _prefill_kernels.py
(paged/ragged/MLA/dense/masked-sequence), _decode_kernels.py (decode +
state merge). kv_cache.py now holds only the PagedKVCache classes and
re-exports every moved symbol so existing imports keep working.
tree_attn.py also switches to the shared helpers.
kv_cache.py drops from 2815 to 668 lines; the package is ~2.4k lines
smaller overall. No test files modified; GPU tests pass unchanged (72
passed, 4 pre-existing skips).
Initial goal was to add SHAPE and RANGE tests, solving part of #18971
This PR achieves that and includes the minimum necessary frontend fixes
discovered during implementation so those tests reflect real supported
behavior instead of xfail/workarounds.
so this PR includes both:
**1. New SHAPE/RANGE tests
2. Targeted frontend fixes required to make those tests pass correctly**
## Why These Changes Were Needed
- SHAPE conversion previously produced symbolic shape info instead of a
tensor output aligned with TFLite SHAPE semantics.
- RANGE conversion passed tensor expressions into arange instead of
scalar values for constant scalar bounds.
- Zero-input TFLite subgraphs (valid for constant-only models such as
RANGE without inputs) were blocked by a strict assertion.
- Model output collection was brittle for constant/prefetched outputs
and could fail when output expressions were not already in the expr
table.
- As a result, i could not add meaningful SHAPE/RANGE coverage without
fixing frontend behavior.
## **Modifications**
### **Frontend Changes** (In tflite_frontend.py):
- Updated convert_shape: SHAPE now materializes shape output as a tensor
using shape_to_tensor(shape_of(...))
- Applies output dtype casting based on ShapeOptions OutType
(int32/int64)
- Updated convert_range: Extracts scalar values for start/limit/delta
from scalar constants
- Calls arange with scalar-like values
- Keeps dynamic scalar RANGE explicit as unsupported (raises
OpNotImplemented with clear message)
- Updated _input_type: Removed assumption that every subgraph must have
at least one input
- Supports valid zero-input subgraphs
- Updated from_tflite output assembly: Resolves outputs via tensor
wrappers and get_tensor_expr instead of direct expr-table lookup by name
---
**Main functional changes are localized to SHAPE/RANGE conversion and
model output/input handling.**
---------
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This PR adds unit test coverage for `PRELU` activation
in the Relax TFLite frontend, as part of
https://github.com/apache/tvm/issues/18971
- Added unit test for `PRELU` and
Enabled converter to handle alpha broadcasting more cleanly across
constant and expression-backed alpha inputs.
Adds `_attention_sequence_prefill_with_mask` in
`python/tvm/relax/frontend/nn/llm/kv_cache.py` — a masked variant of the
existing sequence prefill kernel that supports right-padded encoder
batches with per-sample `valid_lens`.
The existing `_attention_sequence_prefill` assumes all positions in `[0,
seq_len)` are valid, which breaks for padded encoder inputs where each
batch element has a different valid prefix length. This helper adds the
masking semantics needed for correctness:
- accepts a per-batch `valid_lens` input
- ignores padded query rows and padded key/value positions
- excludes padded `(row, col)` pairs from the online softmax update
It reuses the existing prefill kernel config and schedule — no new
tuning knobs, no target-specific changes, no performance claims.
Correctness only.
## Motivation: encoder batch prefill for downstream consumers
This is the TVM-side primitive needed to support **encoder batch
prefill** in downstream projects like `mlc-llm`, where padded encoder
batches with `valid_lens` need to be lowered without materializing an
explicit broadcast attention mask on the host.
The helper is generic and useful for any encoder-style sequence prefill
consumer with per-sample valid lengths.
Bana