nki.isa.nc_match_replace8#

nki.isa.nc_match_replace8(dst, data, vals, imm, dst_idx=None, name=None)[source]#

Replace first occurrence of each value in vals with imm in data using the Vector engine and return the replaced tensor. If dst_idx tile is provided, the indices of the matched values are written to dst_idx.

This instruction reads the input data, replaces the first occurrence of each of the given values (from vals tensor) with the specified immediate constant and, optionally, output indices of matched values to dst_idx. When performing the operation, the free dimensions of both data and vals are flattened. However, these dimensions are preserved in the replaced output tensor and in dst_idx respectively. The partition dimension defines the parallelization boundary. Match, replace, and index generation operations execute independently within each partition.

The data tensor can be up to 5-dimensional, while the vals tensor can be up to 3-dimensional. The vals tensor must have exactly 8 elements per partition. The data tensor must have no more than 16,384 elements per partition. The replaced output will have the same shape as the input data tensor. data and vals must have the same number of partitions. Both input tensors can come from SBUF or PSUM.

Behavior is undefined if vals tensor contains values that are not in the data tensor.

If provided, a mask is applied to the data tensor.

NumPy equivalent:

# Let's assume we work with NumPy, and ``data``, ``vals`` are 2-dimensional arrays
# (with shape[0] being the partition axis) and imm is a constant float32 value.

import numpy as np

# Get original shapes
data_shape = data.shape
vals_shape = vals.shape

# Reshape to 2D while preserving first dimension
data_2d = data.reshape(data_shape[0], -1)
vals_2d = vals.reshape(vals_shape[0], -1)

# Initialize output array for indices
indices = np.zeros(vals_2d.shape, dtype=np.uint32)

for i in range(data_2d.shape[0]):
  for j in range(vals_2d.shape[1]):
    val = vals_2d[i, j]
    # Find first occurrence of val in data_2d[i, :]
    matches = np.where(data_2d[i, :] == val)[0]
    if matches.size > 0:
      indices[i, j] = matches[0]  # Take first match
      data_2d[i, matches[0]] = imm

output = data_2d.reshape(data.shape)
indices = indices.reshape(vals.shape) # Computed only if ``dst_idx`` is specified
Parameters:

  • dst – the modified data tensor

  • data – the data tensor to modify

  • dst_idx – (optional) the destination tile to write flattened indices of matched values

  • vals – tensor containing the 8 values per partition to replace

  • imm – float32 constant to replace matched values with