[BUG] [CuTeDSL] cute.composition throws when rhs has a dynamic shape

[lw]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
Invoking cute.composition on a RHS layout that contains one mode with a dynamic size causes an exception to be raised. At first I thought that this was by design, because if cute.composition were to verify that RHS's codomain falls within LHS's domain, it needs to know RHS's domain. But apparently such a check isn't actually performed (not even for fully-static layouts).

Steps/Code to reproduce bug

import torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack


@cute.jit
def repro(tensor: cute.Tensor):
    lhs = cute.make_layout((4096, 8192), stride=(1, 4096))
    rhs = cute.make_layout(
        (32, tensor.shape[1] // 8, 8, 128),
        stride=(
            cute.E(0),
            8 * cute.E(1),
            cute.E(1),
            32 * cute.E(0),
        ),
    )
    print(cute.composition(lhs, rhs))


tensor = from_dlpack(torch.empty((4096, 8192), device="cuda"))
tensor = tensor.mark_compact_shape_dynamic(mode=1)
repro(tensor)

Raises:

cutlass._mlir._mlir_libs._site_initialize.<locals>.MLIRError: Operation creation failed:
error: unknown: unable to compute the following composition: '!cute.layout<"(4096,8192):(1,4096)">' o '!cute.layout<"(32,?,8,128):(1@0,8@1,1@1,32@0)">'

Environment details (please complete the following information):
CuTeDSL version 4.7.0

[ccecka]

This does appear to be a bug in CuTeDSL.

It functions in PyCuTe

>>> from pycute import *
>>> import sympy
>>> N = sympy.symbols("N", positive=True, integer=True)
>>> composition(Layout((4096,8192), (1,4096)), Layout((32,N,8,128), (E(0), 8*E(1), E(1), 32*E(0))))
Layout((32, N, 8, 128), (1, 32768, 4096, 32))