Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughMACA dense, preshuffled, and sparse tensor-core emitters now default to ChangesMACA MMA dtype handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tilelang/maca/intrinsics/macro/mma_macro_generator.py (1)
145-163: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRefactor
_initialize_mma_prefixto reuse_normalize_dtype.The new
_normalize_dtypehelper was created to sharedtype('...')unwrapping, but_initialize_mma_prefix(lines 172–174) still has its own inline copy of the same logic. Consolidating this eliminates a DRY violation and ensures consistent normalization across all dtype parsing paths.♻️ Proposed refactor for _initialize_mma_prefix
def _initialize_mma_prefix(self, k_dim=16): in_dtype = self.a_dtype M_DIM, N_DIM = self.M_DIM, self.N_DIM - in_dtype_key = str(in_dtype) - if in_dtype_key.startswith("dtype('") and in_dtype_key.endswith("')"): - in_dtype_key = in_dtype_key[7:-2] + in_dtype_key = self._normalize_dtype(in_dtype) in_dtype_map = {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/maca/intrinsics/macro/mma_macro_generator.py` around lines 145 - 163, Update _initialize_mma_prefix to call the existing _normalize_dtype helper for each dtype instead of duplicating the inline dtype('...') unwrapping logic. Preserve the current prefix-selection behavior while ensuring all dtype parsing uses the shared normalization path.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tilelang/maca/intrinsics/macro/mma_macro_generator.py`:
- Around line 158-163: The MACA GEMM path must not derive the MMA accumulator
dtype from the float16 output buffer. Update GemmMMA’s TensorCoreIntrinEmitter
setup to pass float32 explicitly (or otherwise separate accumulator and output
dtypes), while preserving C.dtype as the output dtype and leaving the
_initialize_abbrev validation intact.
---
Nitpick comments:
In `@tilelang/maca/intrinsics/macro/mma_macro_generator.py`:
- Around line 145-163: Update _initialize_mma_prefix to call the existing
_normalize_dtype helper for each dtype instead of duplicating the inline
dtype('...') unwrapping logic. Preserve the current prefix-selection behavior
while ensuring all dtype parsing uses the shared normalization path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 81dd570d-2d5c-426e-baa3-4d2af1c21a3e
📒 Files selected for processing (1)
tilelang/maca/intrinsics/macro/mma_macro_generator.py
a105858 to
5832061
Compare
✅ Action performedFull review finished. Your plan includes PR reviews subject to rate limits. More reviews will be available in 3 minutes. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tilelang/maca/intrinsics/macro/mma_sp_macro_generator.py (1)
150-157: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueNormalize
in_dtypehere tootilelang/maca/intrinsics/macro/mma_sp_macro_generator.py:159still usesself.a_dtypedirectly for the prefix lookup, while the dense emitter normalizes before the same step.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/maca/intrinsics/macro/mma_sp_macro_generator.py` around lines 150 - 157, Update the prefix lookup in the MMA macro generator to use the normalized input dtype, matching the normalization already applied in _initialize_abbrev and the dense emitter. Replace direct use of self.a_dtype in the relevant lookup while preserving the existing abbreviation and accumulation validation behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@testing/python/maca/test_maca_mma_accum_dtype.py`:
- Around line 50-56: Make test_int32_accumulation_is_still_accepted independent
of the host environment by supplying a fixed MACA target or applying the
existing MACA-target skip/fixture before _emitter initialization. Ensure
determine_target(return_object=True) cannot fail when no MACA target is
installed, while preserving the int32 accumulation assertion.
---
Nitpick comments:
In `@tilelang/maca/intrinsics/macro/mma_sp_macro_generator.py`:
- Around line 150-157: Update the prefix lookup in the MMA macro generator to
use the normalized input dtype, matching the normalization already applied in
_initialize_abbrev and the dense emitter. Replace direct use of self.a_dtype in
the relevant lookup while preserving the existing abbreviation and accumulation
validation behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 1e09e06a-b5a3-4b9f-b1cf-ae5882b1cc9b
📒 Files selected for processing (3)
testing/python/maca/test_maca_mma_accum_dtype.pytilelang/maca/intrinsics/macro/mma_macro_generator.pytilelang/maca/intrinsics/macro/mma_sp_macro_generator.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tilelang/maca/intrinsics/macro/mma_macro_generator.py
| @pytest.mark.parametrize("emitter_cls", DENSE_EMITTERS) | ||
| def test_int32_accumulation_is_still_accepted(emitter_cls): | ||
| # `__builtin_mxc_mma_16x16x16i8` accumulates in `int32x4`, so rejecting fp16 must not | ||
| # spill over onto the integer shapes. | ||
| emitter = _emitter(emitter_cls, a_dtype=T.int8, b_dtype=T.int8, accum_dtype=T.int32) | ||
|
|
||
| assert emitter.accum_dtype_abbrv == "int32" |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate determine_target definition and check if it has a default fallback
rg -n --type=python 'def determine_target' -A 20Repository: tile-ai/tilelang-metax
Length of output: 191
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find the relevant definitions and call sites.
git ls-files | rg 'maca_mma_accum_dtype\.py|mma_macro_generator\.py|determine_target|target.*mcpu|_initialize_k_dim'
# Map likely files if present.
for f in $(git ls-files | rg 'maca_mma_accum_dtype\.py|mma_macro_generator\.py'); do
echo "=== $f ==="
wc -l "$f"
done
# Read the review-target test file if it exists.
test_file=$(git ls-files | rg 'testing/python/maca/test_maca_mma_accum_dtype\.py' | head -n 1 || true)
if [ -n "${test_file:-}" ]; then
echo "=== $test_file ==="
cat -n "$test_file" | sed -n '1,220p'
fi
# Read the emitter implementation around target selection if it exists.
impl_file=$(git ls-files | rg 'mma_macro_generator\.py$' | head -n 1 || true)
if [ -n "${impl_file:-}" ]; then
echo "=== $impl_file ==="
cat -n "$impl_file" | sed -n '1,220p'
fi
# Search for determine_target implementation and any fallback/default target logic.
rg -n --glob '*.py' 'def determine_target|determine_target\(|return_object=True|target\.attrs\["mcpu"\]|attrs\["mcpu"\]' .Repository: tile-ai/tilelang-metax
Length of output: 22022
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read determine_target and the MACA emitter logic around target access.
target_file="tilelang/backend/target.py"
maca_file="tilelang/maca/intrinsics/macro/mma_macro_generator.py"
echo "=== $target_file ==="
cat -n "$target_file" | sed -n '122,260p'
echo "=== $maca_file ==="
cat -n "$maca_file" | sed -n '100,170p'Repository: tile-ai/tilelang-metax
Length of output: 4642
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target_file="tilelang/backend/target.py"
echo "=== auto_detect_target / _finalize_target / related helpers ==="
rg -n --glob '*.py' 'def auto_detect_target|def _finalize_target|def _validate_manual_target|class Target|Target\.current|mcpu' "$target_file" -A 80 -B 20Repository: tile-ai/tilelang-metax
Length of output: 4967
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== tilelang/testing/__init__.py ==="
cat -n tilelang/testing/__init__.py | sed -n '1,220p'
echo "=== MACA target registration / detectors ==="
rg -n --glob '*.py' 'register_target_detector|register_target_normalizer|determine_target\("auto"|Target\.current|auto_detect_target|maca' tilelang testing -A 30 -B 10
echo "=== pytest fixtures / conftest mentioning target ==="
rg -n --glob 'conftest.py' --glob '*.py' 'determine_target|Target\.current|auto_detect_target|target' testing tilelang -A 20 -B 10Repository: tile-ai/tilelang-metax
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== MACA target detector registration ==="
rg -n --glob 'tilelang/**/target.py' --glob 'tilelang/**/*.py' 'register_target_detector\(|def _detect_.*target|check_.*availability|Target\.current\(allow_none=True\)|auto_detect_target\(' tilelang -A 20 -B 10
echo "=== test harness / fixtures that may set target ==="
rg -n --glob 'testing/**/*.py' --glob 'testing/**/conftest.py' 'determine_target\("auto"|determine_target\(return_object=True\)|Target\.current|set_current|target =' testing -A 20 -B 10
echo "=== MACA test file imports / top-level setup ==="
cat -n testing/python/maca/test_maca_mma_accum_dtype.py | sed -n '1,120p'Repository: tile-ai/tilelang-metax
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== tilelang/testing/__init__.py (relevant lines) ==="
rg -n 'def requires_|def main|Target\.current|determine_target\("auto"|determine_target\(return_object=True\)|pytest|fixture' tilelang/testing/__init__.py -A 30 -B 10
echo "=== MACA test harness / conftest ==="
git ls-files | rg '^testing/python/maca/.*(conftest\.py|test_.*\.py)$' | while read -r f; do
echo "=== $f ==="
rg -n 'Target\.current|determine_target|requires_|maca|cuda|fixture' "$f" -A 20 -B 10
doneRepository: tile-ai/tilelang-metax
Length of output: 9339
Guard the int8 path behind a MACA target
_emitter(..., a_dtype=T.int8, b_dtype=T.int8, accum_dtype=T.int32) still calls determine_target(return_object=True) during emitter init, so this test will fail before the assertion when no MACA target is active or installed. Add a MACA skip/fixture or inject a fixed target here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@testing/python/maca/test_maca_mma_accum_dtype.py` around lines 50 - 56, Make
test_int32_accumulation_is_still_accepted independent of the host environment by
supplying a fixed MACA target or applying the existing MACA-target skip/fixture
before _emitter initialization. Ensure determine_target(return_object=True)
cannot fail when no MACA target is installed, while preserving the int32
accumulation assertion.
28c1d62 to
8fb9503
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
testing/python/maca/test_maca_mma_accum_dtype.py (1)
46-46: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a raw string for the regular expression pattern.
The pattern passed to
match=contains regex metacharacters (.*), but it is not a raw string. Using a raw string (r"...") ensures that escape sequences and metacharacters are treated correctly.♻️ Proposed refactor
- with pytest.raises(ValueError, match="does not support .* accumulation"): + with pytest.raises(ValueError, match=r"does not support .* accumulation"):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/maca/test_maca_mma_accum_dtype.py` at line 46, Update the pytest.raises match pattern in the accumulation dtype test to use a raw string literal, preserving the existing “does not support .* accumulation” regex content.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@testing/python/maca/test_maca_mma_accum_dtype.py`:
- Line 46: Update the pytest.raises match pattern in the accumulation dtype test
to use a raw string literal, preserving the existing “does not support .*
accumulation” regex content.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 296198de-199d-4419-81cd-775189f24794
📒 Files selected for processing (3)
testing/python/maca/test_maca_mma_accum_dtype.pytilelang/maca/intrinsics/macro/mma_macro_generator.pytilelang/maca/intrinsics/macro/mma_sp_macro_generator.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tilelang/maca/intrinsics/macro/mma_sp_macro_generator.py
- tilelang/maca/intrinsics/macro/mma_macro_generator.py
8fb9503 to
f546229
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
testing/python/maca/test_maca_mma_accum_dtype.py (1)
45-45: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the regular-expression intent explicit.
Ruff RUF043 flags this
pytest.raises(match=...)string because.*is an unescaped regular-expression metacharacter. Use a raw string to preserve the current matching behavior and remove the warning. (docs.astral.sh)Proposed fix
- with pytest.raises(ValueError, match="does not support .* accumulation"): + with pytest.raises(ValueError, match=r"does not support .* accumulation"):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/maca/test_maca_mma_accum_dtype.py` at line 45, Update the pytest.raises match argument in the relevant test to use a raw string literal, preserving the existing “does not support .* accumulation” regular-expression behavior while resolving Ruff RUF043.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@testing/python/maca/test_maca_mma_accum_dtype.py`:
- Line 45: Update the pytest.raises match argument in the relevant test to use a
raw string literal, preserving the existing “does not support .* accumulation”
regular-expression behavior while resolving Ruff RUF043.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 686aebe1-d5f4-46ff-8b31-73b94f8d2f7c
📒 Files selected for processing (3)
testing/python/maca/test_maca_mma_accum_dtype.pytilelang/maca/intrinsics/macro/mma_macro_generator.pytilelang/maca/intrinsics/macro/mma_sp_macro_generator.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tilelang/maca/intrinsics/macro/mma_sp_macro_generator.py
- tilelang/maca/intrinsics/macro/mma_macro_generator.py
f546229 to
dfb545c
Compare
…not hold The MMA lowering in src/maca/codegen/codegen_maca.cc synthesises the builtin as "__builtin_mxc_mma_" + prefix and casts the C operand to dtype_map[C_dtype]. The float matrix-core builtins take and return a float32x4, so any float accumulator narrower than 32 bits hands them the wrong vector type and mxcc rejects the call, several stages after the emitter accepted it. On upstream, with fp16 inputs, a C500 compiles accum_dtype=float32 and int32 and rejects float16, bfloat16, float8_e4m3 and float8_e5m2. The check is therefore on the width rather than on a single dtype, and it lives in one place that both the dense and the sparse emitter call. The default was float16, which no MACA build could compile, so it moves to float32. Signed-off-by: 林晨 (Leo Cheng) <chengkelfan@qq.com>
dfb545c to
eec41aa
Compare
Closes #69.
Problem
The MMA lowering in
src/maca/codegen/codegen_maca.ccsynthesises the builtin name as"__builtin_mxc_mma_" + prefixand casts the C operand todtype_map[C_dtype], wheredtype_map["float16x4"] == "float16x4". The MACA float matrix-core builtins take and return afloat32x4, so a narrow float accumulator hands them the wrong vector type:The emitter accepted it, so this only surfaced several stages later, inside mxcc.
accum_dtypedefaulted tofloat16in all three emitters, so the default configuration could never compile on a MACA target.Changes
The check is on the accumulator width rather than on one dtype. On upstream, with fp16 inputs, a C500 compiles
float32andint32and rejects every float type narrower than 32 bits:accum_dtypefloat32int32float16RuntimeErrorfrom mxccbfloat16RuntimeErrorfrom mxccfloat8_e4m3RuntimeErrorfrom mxccfloat8_e5m2RuntimeErrorfrom mxcc_reject_narrow_float_accumraises aValueErrornaming the offending dtype at emitter construction. It sits in one place and both the dense and the sparse emitter call it, so the two cannot drift apart. Integer shapes accumulate inint32and are untouched.The three
accum_dtypedefaults move fromfloat16tofloat32.Test Plan
testing/python/maca/test_maca_mma_accum_dtype.py, on a MetaX C500 (MACA 3.5.3.20):float16,bfloat16,float8_e4m3,float8_e5m2— is rejected, on all three emittersfloat32, on all three emitters. The helper omitsaccum_dtypeentirely for this case, so it is the emitter's own default that is exercisedint32accumulation withint8inputs still constructs, so the rejection does not spill over onto the integer shapesTest Result
Only the
int8case needs a live target:_initialize_k_dimcallsdetermine_target()on the 8-bit branch alone, to readmcpu. It carries@tilelang.testing.requires_cuda, which in this fork resolves to_check_is_maca(). The others run everywhere.Summary
ValueError.float32accumulation.int32accumulation support for integer MMA.float32defaults, and denseint32accumulation support.MACA tests passed for the added coverage: 8 tests passed. The test runner was later OOM-killed with exit code 137.