[NVGPU] Fix nvdsl examples #156830

castigli · 2025-09-04T08:52:37Z

This PR aims at fixing the nvdsl examples which got a bit out of sync not being tested in the CI.

The fixed bugs were related to the following PRs:

move to nanobind [mlir python] Port Python core code to nanobind. #118583
split gpu module initialization [mlir][gpu] Change GPU modules to globals #135478

github-actions · 2025-09-04T08:52:56Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

Wolfram70 · 2025-09-04T18:51:13Z

Thanks for bringing this to our attention!

I looked into this a bit, and it does look like this crash is occurring in the --gpu-module-to-binary pass.
For Ch4.py, dumping the .mlir file and extracting the LLVMIR generated during --gpu-module-to-binary, we get
extracted-llvmir.txt. Running llc -mtriple=nvptx64 -mcpu=sm_90a -mattr=+ptx80 on this reproduces this crash exactly so it seems to be an issue during codegen. I am not sure why this is occurring in this specific case.

@durga4github @abhilash1910 Do you have any idea why this might be happening?

abhilash1910 · 2025-09-05T03:55:59Z

Taking a look at the codegen. Thanks for highlighting. IR does not seem incorrect though at first glance.
Edit: Fix is in progress.

@durga4github

Context: Highlighted from #156830 , this is an Isel lowering issue in the NVPTX backend for prefetch.tensormap intrinsic. It is caused by unchecked pattern rewrite during infer-address-space pass. This intrinsic is valid only for const, param and generic address-spaces. Any other address space is invalid. Currently, this intrinsic gets falsely re-written to target AS(1), when the pointer-argument of the intrinsic comes as an argument of a kernel function. So, this patch adds a check on the correct address-spaces before re-writing them. cc @durga4github FYI: @Wolfram70 @rupprecht @castigli

llvmbot · 2025-10-02T09:36:39Z

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-nvgpu

@llvm/pr-subscribers-mlir-gpu

Author: Giacomo Castiglioni (castigli)

Changes

This PR aims at fixing the nvdsl examples which got a bit out of sync not being tested in the CI.

The fixed bugs were related to the following PRs:

move to nanobind #118583
split gpu module initialization #135478

There is one remaining bug that I think #153134 introduced. When running the Ch4 and Ch5 the nvvm.prefetch tensormap intrisic leads to the following error on sm_90a

LLVM ERROR: Cannot select: intrinsic %llvm.nvvm.prefetch.tensormap
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.	Program arguments: mlir-opt before.mlir --gpu-module-to-binary
1.	Running pass 'Function Pass Manager' on module 'LLVMDialectModule'.
2.	Running pass 'NVPTX DAG-&gt;DAG Pattern Instruction Selection' on function '@<!-- -->gemm_multistage_kernel'
...

Perahps @Wolfram70 or @grypp could help me out with the last bug?
Could be the solution to revert momentarily to inline ptx?
[edit] this was resolved in #159253

Full diff: https://github.com/llvm/llvm-project/pull/156830.diff

3 Files Affected:

(modified) mlir/test/Examples/NVGPU/Ch5.py (+1-1)
(modified) mlir/test/Examples/NVGPU/tools/nvdsl.py (+3-4)
(modified) mlir/test/Examples/NVGPU/tools/nvgpucompiler.py (+3-1)

diff --git a/mlir/test/Examples/NVGPU/Ch5.py b/mlir/test/Examples/NVGPU/Ch5.py
index f98cfd758a75f..91c346c837dda 100644
--- a/mlir/test/Examples/NVGPU/Ch5.py
+++ b/mlir/test/Examples/NVGPU/Ch5.py
@@ -156,7 +156,7 @@ def producer_loop(
 ):
     phase = const(True, ty=T.bool())
 
-    for iv, phase in scf.for_(0, (K // TILE_K), 1, [phase]):
+    for iv, phase, _ in scf.for_(0, (K // TILE_K), 1, [phase]):
         stage = iv % num_stages
         # Wait MMA to be done
         mbar_mma[stage].try_wait(phase)
diff --git a/mlir/test/Examples/NVGPU/tools/nvdsl.py b/mlir/test/Examples/NVGPU/tools/nvdsl.py
index 90dbb2355e1c8..d4c50fc9bc28d 100644
--- a/mlir/test/Examples/NVGPU/tools/nvdsl.py
+++ b/mlir/test/Examples/NVGPU/tools/nvdsl.py
@@ -84,8 +84,7 @@ def arrive(self, txcount: int = 0, predicate=None):
                 self.mbar_group_op, txcount_op, self.id_op, predicate=predicate
             )
         else:
-            nvgpu.mbarrier_arrive(
-                ir.Type.parse("!nvgpu.mbarrier.token"), self.mbar_group_op, self.id_op
+            nvgpu.mbarrier_arrive(self.mbar_group_op, self.id_op
             )
 
     def try_wait(self, phase: bool = False, ticks: int = 10000000):
@@ -144,7 +143,7 @@ def create_descriptor(self, device_ptr):
             device_ptr,
         )
         self.tma_descriptor = nvgpu.TmaCreateDescriptorOp(
-            tma_descriptor_ty, device_unranked_memref, map(const, self.tma_box_shape)
+            tma_descriptor_ty, device_unranked_memref, list(map(const, self.tma_box_shape))
         )
         return self.tma_descriptor.result
 
@@ -156,7 +155,7 @@ def load(self, dest, mbarrier: Mbarriers, coords=[0], predicate=None):
             dest,
             mbarrier.mbar_group_op,
             self.tma_descriptor,
-            coordinates=map(const, coords),
+            coordinates=list(map(const, coords)),
             mbarId=mbarrier.id_op,
             predicate=predicate,
         )
diff --git a/mlir/test/Examples/NVGPU/tools/nvgpucompiler.py b/mlir/test/Examples/NVGPU/tools/nvgpucompiler.py
index 1c9cc74fcd169..4b661f8df6a9f 100644
--- a/mlir/test/Examples/NVGPU/tools/nvgpucompiler.py
+++ b/mlir/test/Examples/NVGPU/tools/nvgpucompiler.py
@@ -35,9 +35,11 @@ def compile(self, module: ir.Module):
 
     def jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:
         """Wraps the module in a JIT execution engine."""
-        return execution_engine.ExecutionEngine(
+        ee = execution_engine.ExecutionEngine(
             module, opt_level=self.opt_level, shared_libs=self.shared_libs
         )
+        ee.initialize()
+        return ee
 
     def compile_and_jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:
         """Compiles and jits the module."""

durga4github · 2025-10-08T09:43:41Z

@castigli , I updated the commit-msg since the prefetch.tensormap issue is resolved.

Could you please rebase and push once? I can initiate the workflows to run, to get CI results.

castigli · 2025-10-08T12:23:43Z

@castigli , I updated the commit-msg since the prefetch.tensormap issue is resolved.

Could you please rebase and push once? I can initiate the workflows to run, to get CI results.

Done!

durga4github · 2025-10-08T13:01:19Z

@castigli , I updated the commit-msg since the prefetch.tensormap issue is resolved.
Could you please rebase and push once? I can initiate the workflows to run, to get CI results.

Done!

Thanks, initiated the CI.

github-actions · 2025-10-08T13:01:29Z

⚠️ Python code formatter, darker found issues in your code. ⚠️

You can test this locally with the following command:

darker --check --diff -r origin/main...HEAD mlir/test/Examples/NVGPU/Ch5.py mlir/test/Examples/NVGPU/tools/nvdsl.py mlir/test/Examples/NVGPU/tools/nvgpucompiler.py

⚠️
The reproduction instructions above might return results for more than one PR
in a stack if you are using a stacked PR workflow. You can limit the results by
changing origin/main to the base branch/commit you want to compare against.
⚠️

View the diff from darker here.

--- tools/nvdsl.py	2025-10-08 11:51:06.000000 +0000
+++ tools/nvdsl.py	2025-10-08 13:01:05.815271 +0000
@@ -82,12 +82,11 @@
             txcount_op = const(txcount)
             nvgpu.mbarrier_arrive_expect_tx(
                 self.mbar_group_op, txcount_op, self.id_op, predicate=predicate
             )
         else:
-            nvgpu.mbarrier_arrive(self.mbar_group_op, self.id_op
-            )
+            nvgpu.mbarrier_arrive(self.mbar_group_op, self.id_op)
 
     def try_wait(self, phase: bool = False, ticks: int = 10000000):
         ticks_op = const(ticks)
         phase_op = const(phase, T.bool())
         nvgpu.MBarrierTryWaitParityOp(
@@ -141,11 +140,13 @@
                 self.memref_ty.element_type, self.memref_ty.memory_space
             ),
             device_ptr,
         )
         self.tma_descriptor = nvgpu.TmaCreateDescriptorOp(
-            tma_descriptor_ty, device_unranked_memref, list(map(const, self.tma_box_shape))
+            tma_descriptor_ty,
+            device_unranked_memref,
+            list(map(const, self.tma_box_shape)),
         )
         return self.tma_descriptor.result
 
     def prefetch(self, predicate=None):
         nvgpu.tma_prefetch_descriptor(self.tma_descriptor, predicate=predicate)

durga4github · 2025-10-08T13:02:21Z

@grypp , Kindly take a look when you get a chance.

grypp

Thanks for fixing these

grypp · 2025-10-08T13:10:34Z

I feel like we should enable LIT testing without running for these test so at least they can get compiled.

castigli · 2025-10-08T16:45:50Z

I feel like we should enable LIT testing without running for these test so at least they can get compiled.

In principle I agree, but I don't have good ideas to check both compile-only and compile-and-run without mucking the code too much.
What about something like

...
# RUN: env MLIR_RUN_CUDA_SM90_TESTS=%mlir_run_cuda_sm90_tests
...
if run_if_cuda_sm90_enabled(lambda: saxpy(x, y, alpha)) is None:
    #  4. Verify MLIR with reference computation
    ref = np.ones((M, N), np.float32)
    ref += x * alpha
    np.testing.assert_allclose(y, ref, rtol=5e-03, atol=1e-01)
print("PASS")
# CHECK: PASS

with the util defined as

def run_if_cuda_sm90_enabled(func, *args, **kwargs):
    """Execute a function if CUDA SM90 tests are enabled, otherwise print a warning."""
    mlir_run_cuda_sm90_tests = os.getenv("MLIR_RUN_CUDA_SM90_TESTS")
    if mlir_run_cuda_sm90_tests == "1" or mlir_run_cuda_sm90_tests is None:
        return func(*args, **kwargs)
    else:
        print("warning: skipping test execution")
        return -1

castigli changed the title ~~Fix nvdsl examples~~ [NVGPU] Fix nvdsl examples Sep 4, 2025

rupprecht mentioned this pull request Sep 17, 2025

[MLIR][NVVM][NVGPU] Combine prefetch and prefetch.tensormap #153134

Merged

abhilash1910 mentioned this pull request Sep 17, 2025

[NVPTX] prefetch.tensormap pattern rewriter fix #159253

Merged

castigli force-pushed the fix-nvdsl branch from e904099 to d0ae04d Compare October 2, 2025 09:33

castigli marked this pull request as ready for review October 2, 2025 09:36

castigli requested a review from grypp as a code owner October 2, 2025 09:36

llvmbot added mlir:gpu mlir mlir:nvgpu labels Oct 2, 2025

fix nvdsl

a915cce

castigli force-pushed the fix-nvdsl branch from d0ae04d to a915cce Compare October 8, 2025 12:22

durga4github approved these changes Oct 8, 2025

View reviewed changes

grypp approved these changes Oct 8, 2025

View reviewed changes

format

49b7e6a

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[NVGPU] Fix nvdsl examples #156830

[NVGPU] Fix nvdsl examples #156830

castigli commented Sep 4, 2025 •

edited by durga4github

Loading

Uh oh!

github-actions bot commented Sep 4, 2025

Uh oh!

Wolfram70 commented Sep 4, 2025

Uh oh!

abhilash1910 commented Sep 5, 2025 •

edited

Loading

Uh oh!

llvmbot commented Oct 2, 2025 •

edited

Loading

Uh oh!

durga4github commented Oct 8, 2025

Uh oh!

castigli commented Oct 8, 2025

Uh oh!

durga4github commented Oct 8, 2025

Uh oh!

github-actions bot commented Oct 8, 2025

Uh oh!

durga4github commented Oct 8, 2025

Uh oh!

grypp left a comment

Uh oh!

grypp commented Oct 8, 2025

Uh oh!

castigli commented Oct 8, 2025

Uh oh!

Uh oh!

[NVGPU] Fix nvdsl examples #156830

Are you sure you want to change the base?

[NVGPU] Fix nvdsl examples #156830

Conversation

castigli commented Sep 4, 2025 • edited by durga4github Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Sep 4, 2025

Uh oh!

Wolfram70 commented Sep 4, 2025

Uh oh!

abhilash1910 commented Sep 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Oct 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

durga4github commented Oct 8, 2025

Uh oh!

castigli commented Oct 8, 2025

Uh oh!

durga4github commented Oct 8, 2025

Uh oh!

github-actions bot commented Oct 8, 2025

Uh oh!

durga4github commented Oct 8, 2025

Uh oh!

grypp left a comment

Choose a reason for hiding this comment

Uh oh!

grypp commented Oct 8, 2025

Uh oh!

castigli commented Oct 8, 2025

Uh oh!

Uh oh!

castigli commented Sep 4, 2025 •

edited by durga4github

Loading

abhilash1910 commented Sep 5, 2025 •

edited

Loading

llvmbot commented Oct 2, 2025 •

edited

Loading