NVIDIA / NVIDIA/cutlass

[BUG] pycute: logical_product with an integer tiler computes logical_divide

Open
#3,478 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

Describe the bug

In python/pycute/layout.py, logical_product dispatches an integer tiler to logical_divide, so it returns a divide result under a product name. The Layout-wrapped form is unaffected, so the two spellings of the same call disagree.

python/pycute/layout.py:312-316:

def logical_product(layoutA, layoutB):
  if layoutB is None:
    return layoutA
  elif is_int(layoutB):
    return logical_divide(layoutA, Layout(layoutB))   # <-- should be logical_product

The identical line at python/pycute/layout.py:301 is correct: it belongs to logical_divide directly above.

The C++ reference does the right thing at include/cute/layout.hpp:1670:

} else if constexpr (is_integral<Tiler>::value) {
    return logical_product(block, make_layout(tiler));
}

Steps/Code to reproduce bug

cd python
python -c "
from pycute import *
A = Layout((2,5),(5,1))
print('product Layout(4):', logical_product(A, Layout(4)))
print('product int 4    :', logical_product(A, 4))
print('divide  int 4    :', logical_divide(A, 4))
"

Output:

product Layout(4): ((2, 5), 4):((5, 1), 10)     <- correct, matches C++
product int 4    : ((2, 2), 3):((5, 1), 2)      <- wrong
divide  int 4    : ((2, 2), 3):((5, 1), 2)      <- identical to the line above

The size is the clearest tell: size(A) is 10, the product should be 40, and the returned layout has size 12.

Expected behavior

logical_product(A, 4) should equal logical_product(A, Layout(4)).

Scope

Everything reaching the integer branch is affected:

Call Affected
logical_product(A, 4) yes
logical_product(A, (4,2)) yes, the tuple branch recurses per mode into the int branch
zipped_product(A, 4) yes, hier_unzip(logical_product, ...) bottoms out there
tiled_product(A, 4) yes, built on zipped_product
logical_product(A, Layout(4)) no

operators/cutlass/operators/fusion/pycute/layout.py:339 is a copy of the same file and has the same line.

There is currently no test for logical_product or logical_divide in test/python/pycute/.

Environment details

  • Environment location: Bare-metal
  • Pure Python, no GPU, CUDA or build needed
  • Reproduced on main at the time of filing

Additional context

Note this is python/pycute, the Python reference implementation, not the CuTe DSL. The issue-form dropdown has no option for it.

I have a fix and a regression test ready and will open a PR referencing this issue.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at python/pycute/layout.py:312-316 and compare the integer branch with the C++ reference at include/cute/layout.hpp:1670; check the duplicated copy at operators/cutlass/operators/fusion/pycute/layout.py:339. Add a regression test under test/python/pycute/ covering integer and Layout tilers, and confirm the product spellings agree, including the expected size.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
tooling
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.