Mapping

Mapping provides a way to explicitly control the relationship between argument dimensions and kernel dimensions in SlangPy. This extends the broadcasting rules discussed earlier, giving you more precise control over vectorization.

Note

Mapping Uses ML Convention For Dimensions

Mapping dimension indices like (0, 1, 2) refer to numpy array dimensions using the ML convention, where the leftmost dimension (index 0 in this case) has the largest stride, and the rightmost dimension (index 2 in this case) has the smallest stride. See Index Representation Conventions for details on SlangPy’s different index representation conventions.

A Simple Example

Consider this call to an add function that adds two floats:

a = np.random.rand(10, 3, 4)
b = np.random.rand(10, 3, 4)
result = mymodule.add(a, b, _result='numpy')

In this case:

  • a and b are arguments to the add kernel, each with the shape (10, 3, 4).

  • The kernel is dispatched with an overall shape of (10, 3, 4).

  • Each thread indexed by [i, j, k] processes a[i, j, k] and b[i, j, k], writing the result to result[i, j, k].

This represents a straightforward 1-to-1 mapping between argument dimensions and kernel dimensions.

Re-Mapping Dimensions

The map function allows you to modify how argument dimensions correspond to kernel dimensions. For example, the earlier code could be rewritten as:

a = np.random.rand(10, 3, 4)
b = np.random.rand(10, 3, 4)
result = mymodule.add.map((0, 1, 2), (0, 1, 2))(a, b, _result='numpy')

Here, the tuples passed to map explicitly define the mapping: dimension 0 maps to 0, dimension 1 to 1, and dimension 2 to 2 for both a and b. This is the default behavior in SlangPy.

Alternatively, you can use named parameters for clarity:

# Assuming the add function has the signature add(float3 a, float3 b)
a = np.random.rand(10, 3, 4)
b = np.random.rand(10, 3, 4)
result = mymodule.add.map(a=(0, 1, 2), b=(0, 1, 2))(a=a, b=b, _result='numpy')

Mapping arguments with different dimensionalities

Unlike NumPy, SlangPy does not auto-pad dimensions by default. If this behavior is needed, map lets you explicitly define how smaller inputs are aligned with the kernel:

a = np.random.rand(8, 8).astype(np.float32)
b = np.random.rand(8).astype(np.float32)

# This will fail in SlangPy, as `b` is not automatically extended:
result = mymodule.add(a, b, _result='numpy')

# Use explicit mapping instead:
# Equivalent to padding `b` as NumPy would
result = mymodule.add.map(a=(0, 1), b=(1,))(a=a, b=b, _result='numpy')

# Alternatively, you can omit mapping for a as it defaults to 1-to-1:
result = mymodule.add.map(b=(1,))(a=a, b=b, _result='numpy')

Mapping arguments to different dimensions

Another use case is performing some operation in which you wish to broadcast all the elements of one argument across the other. The simplest is the mathematical outer-product:

# Assuming the multiply function has the signature multiply(float a, float b)
a = np.random.rand(10).astype(np.float32)
b = np.random.rand(20).astype(np.float32)

# Map dimensions:
# - a maps to dimension 0 (size 10)
# - b maps to dimension 1 (size 20)
# Resulting kernel and output shape: (10, 20)
result = mymodule.multiply.map(a=(0,), b=(1,))(a=a, b=b, _result='numpy')

Mapping to re-order dimensions

Re-ordering argument dimensions is straightforward with map. For example, to transpose a matrix:

# Assuming the copy function has the signature float copy(float val)
a = np.random.rand(10, 20).astype(np.float32)

# Swap rows and columns:
result = mymodule.copy.map(val=(1, 0))(val=a, _result='numpy')

Mapping to resolve ambiguities

map can resolve ambiguities that would otherwise prevent SlangPy from vectorizing. For example:

# A generic function from the 'nested' section:
void copy_generic<T>(T src, out T dest) {
    dest = src;
}

# Explicitly map dimensions to remove ambiguity:
src = np.random.rand(100).astype(np.float32)
dest = np.zeros_like(src)
result = module.copy_generic.map(src=(0,), dest=(0,))(src=src, dest=dest)

Slangpy now knows:

  • src and dest should map 1 dimension

  • src and dest are both 1D arrays of float

Thus it can infer that you want to pass float into copy_generic and generates the correct kernel.

Mapping Types

map can also define argument types directly, which may improve readability for simple cases:

src = np.random.rand(100)
dest = np.zeros_like(src)

# Map argument types explicitly:
result = module.copy_generic.map(src='float', dest='float')(src=src, dest=dest)

Where in the previous example SlangPy inferred type from dimensionality, it now knows:

  • src and dest should map to float

  • src and dest are both 1D arrays of float

Thus it can infer that you want a 1D kernel.

Summary

The map function in SlangPy provides powerful tools for customizing how arguments align with kernel dimensions. This capability allows you to:

  • Precisely control dimension mappings for arguments, enabling efficient vectorization of complex operations.

  • Handle cases where arguments have different dimensionalities by explicitly aligning dimensions, avoiding the need for auto-padding.

  • Perform operations like broadcasting (e.g., outer products) and reordering dimensions (e.g., matrix transposition) with ease.

  • Resolve ambiguities in generic functions, ensuring correct kernel generation and execution.

These features make map particularly useful for machine learning algorithms, where operations often involve multi-dimensional data with varying shapes and alignment requirements. By enabling fine-grained control over dimension mappings, SlangPy helps optimize operations like tensor manipulations, matrix multiplications, and custom kernels, which are foundational to modern ML workflows.