feat: users can pass shared_memory_carveout to @cuda.jit

[kaeun97]

This PR requires rebase after merging this: #629.

This PR allows users to pass a new argument shared_memory_carveout to the @cuda.jit decorator as described here.

An alternative to consider is to use enum for cudaSharedCarveout but I think the current implementation is cleaner.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[kaeun97]

@gmarkall I didn't know this was still a draft 😅 would appreciate your review! Also, is there anything in particular that I can look into next? Excited to see what's next!

[greptile-apps]

Greptile Summary

This PR adds support for configuring shared memory carveout via the shared_memory_carveout parameter to @cuda.jit, implementing the feature requested in issue #617. The implementation validates input, parses string values to integers, and applies the configuration during kernel binding.

• Adds validation logic in decorators.py that accepts string values ("default", "MaxL1", "MaxShared") or integer values (-1 to 100)
• Implements parsing and application in dispatcher.py via _parse_carveout and bind() methods
• Includes serialization support to preserve carveout settings across kernel caching
• Adds test coverage for both valid and invalid carveout values

The implementation follows the existing patterns in the codebase and properly integrates with the kernel compilation and binding flow. The previous review comments have identified validation issues with boolean types being incorrectly accepted and error message inconsistencies that should be addressed.

Confidence Score: 3/5

• Safe to merge with minor validation improvements recommended
• The implementation is functionally sound with proper integration into the kernel compilation flow. However, the validation logic has edge cases (boolean acceptance) and usability issues (error message inconsistency) that are already identified in previous review threads and should be addressed for better code quality and user experience.
• Pay attention to numba_cuda/numba/cuda/decorators.py for validation improvements

Important Files Changed

Filename	Overview
numba_cuda/numba/cuda/decorators.py	Adds shared_memory_carveout parameter to @cuda.jit with validation logic. Has validation issues with boolean types and error message inconsistency.
numba_cuda/numba/cuda/dispatcher.py	Implements carveout parsing and application logic. Properly handles serialization and applies carveout during kernel binding.
numba_cuda/numba/cuda/tests/cudapy/test_dispatcher.py	Adds test coverage for valid and invalid carveout values. Missing tests for boolean and float types that are incorrectly accepted/not explicitly tested.

[greptile-apps]

Additional Comments (1)

1. numba_cuda/numba/cuda/decorators.py, line 33 (link)

   style: Missing documentation for the new shared_memory_carveout parameter in the docstring. Should add:

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

_{4 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[gmarkall]

@kaeun97 Apologies for the delay in getting to this, it was not because it was a draft that it didn't receive a review, but just that things have been a a bit heavy (and I had a little PTO) over the last few days.

[gmarkall]

/ok to test

[copy-pr-bot]

/ok to test

@gmarkall, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[gmarkall]

/ok to test bbdda61

[gmarkall]

Many thanks again for your efforts here!

I've had a quick skim - apologies for the delay. Some initial thoughts:

• The API design looks good to me.
• We should probably test with and without signatures to the @cuda.jit decorator, because the paths taken through the jit decorator are different when func_or_sig is a function and when it is a signature.
• As this is a user-facing feature, we should add it to the docstring for jit. I think the documentation is then derived from this docstring, so there shouldn't need to be a separate documentation change.

[greptile-apps]

Additional Comments (1)

1. numba_cuda/numba/cuda/dispatcher.py, line 479-481 (link)

   logic: carveout may not be applied when kernel runs on multiple devices

   the carveout is only set in bind(), but bind() is called once during compilation. when get_cufunc() creates a function for a different device (line 333-335 in codegen.py caches per device), the carveout won't be applied

_{6 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{5 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

[kaeun97]

Thank you for the quick response and feedback after your PTO (+ hope you got some chance to rest and have fun)! Much appreciated.

• add docstring for shared_memory_carveout
• add test with signature

Note that as I was testing with signatures, I realized that the shared_memory_carveout value was checked at different times (previous implementation: (without signature) raise error at compilation time, (with signature) raise at decoration time) , so made it consistent to check it at decoration time.

Also, would there be some other issue I can work on after this task?

[kaeun97]

Hi @gmarkall, would be much appreciated if we can get this merged :)

[gmarkall]

/ok to test 373b7c2

[greptile-apps]

_{3 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

The error message displays lowercase string options ["default", "maxl1", "maxshared"], but the documentation (lines 100-101) and tests (test_dispatcher.py line 851) use the capitalized forms "MaxL1", "MaxShared", "default". This inconsistency will confuse users who follow the documentation.

Consider updating the error message to show the capitalized forms that users should actually use:

Suggested change

	valid_strings = ["default", "maxl1", "maxshared"]
	if carveout.lower() not in valid_strings:
	raise ValueError(
	f"Invalid carveout value: {carveout}. "
	f"Must be -1 to 100 or one of {valid_strings}"
	)
	valid_strings = ["default", "MaxL1", "MaxShared"]
	if carveout.lower() not in [s.lower() for s in valid_strings]:
	raise ValueError(
	f"Invalid carveout value: {carveout}. "
	f"Must be -1 to 100 or one of {valid_strings}"
	)

This way the error message displays the same capitalization as the documentation, while still accepting case-insensitive input.

[greptile-apps]

The isinstance(carveout, int) check will also accept boolean values since bool is a subclass of int in Python. This means shared_memory_carveout=True would be treated as 1 and shared_memory_carveout=False as 0, which is likely unintended behavior.

Consider explicitly rejecting boolean types:

Suggested change

	elif isinstance(carveout, int):
	if not (-1 <= carveout <= 100):
	raise ValueError("Carveout must be between -1 and 100")
	elif isinstance(carveout, bool):
	raise TypeError(
	f"shared_memory_carveout must be str or int, got {type(carveout).__name__}"
	)
	elif isinstance(carveout, int):
	if not (-1 <= carveout <= 100):
	raise ValueError("Carveout must be between -1 and 100")

Note: The bool check must come before the int check since bool is a subclass of int.

[greptile-apps]

The shared_memory_carveout parameter is validated here but doesn't check if device=True is also set. Since shared memory carveout only applies to kernel launches (not device functions), using this parameter with device=True should raise an error rather than being silently ignored.

Consider adding validation:

Suggested change

	if shared_memory_carveout is not None:
	_validate_shared_memory_carveout(shared_memory_carveout)
	if shared_memory_carveout is not None:
	if device:
	raise ValueError(
	"shared_memory_carveout cannot be used with device=True. "
	"This parameter only applies to kernels, not device functions."
	)
	_validate_shared_memory_carveout(shared_memory_carveout)

This provides clearer user feedback when the parameter is misused.

[gmarkall]

/ok to test e5ba748

[greptile-apps]

_{3 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

The validation accepts boolean values because bool is a subclass of int in Python. This means shared_memory_carveout=True and shared_memory_carveout=False are incorrectly accepted.

True evaluates to 1 and False evaluates to 0, both within the valid range [-1, 100], but this is likely unintended behavior that could confuse users.

To fix this, check for boolean types explicitly before checking for int:

Suggested change

	elif isinstance(carveout, int):
	if not (-1 <= carveout <= 100):
	raise ValueError("Carveout must be between -1 and 100")
	elif isinstance(carveout, bool):
	raise TypeError(
	f"shared_memory_carveout must be str or int, got {type(carveout).__name__}"
	)
	elif isinstance(carveout, int):
	if not (-1 <= carveout <= 100):
	raise ValueError("Carveout must be between -1 and 100")

[greptile-apps]

The error message displays lowercase strings ["default", "maxl1", "maxshared"], but the documentation (lines 100-101) shows capitalized forms "MaxL1", "MaxShared", "default". This inconsistency will confuse users when they see an error.

The validation correctly accepts any case (due to carveout.lower() on line 320), but the error message should match what's documented.

Suggested change

	valid_strings = ["default", "maxl1", "maxshared"]
	if carveout.lower() not in valid_strings:
	raise ValueError(
	f"Invalid carveout value: {carveout}. "
	f"Must be -1 to 100 or one of {valid_strings}"
	valid_strings = ["default", "maxl1", "maxshared"]
	if carveout.lower() not in valid_strings:
	raise ValueError(
	f"Invalid carveout value: {carveout}. "
	f"Must be -1 to 100 or one of ['default', 'MaxL1', 'MaxShared']"
	)

[greptile-apps]

The test coverage for invalid values should include boolean types, which are currently incorrectly accepted due to the validation bug in decorators.py (line 325).

Consider adding test cases for:

• Boolean values: True, False (currently accepted but shouldn't be)
• Float values: 50.5, 0.0 (should be rejected with TypeError)
• Other invalid types: [], {}, () (should be rejected with TypeError)

Example addition to test_cases:

test_cases = [
    (150, ValueError, "must be between -1 and 100"),
    (-2, ValueError, "must be between -1 and 100"),
    (101, ValueError, "must be between -1 and 100"),
    ("InvalidOption", ValueError, "Invalid carveout value"),
    (True, TypeError, "must be str or int"),  # Currently fails - booleans are accepted
    (False, TypeError, "must be str or int"),  # Currently fails - booleans are accepted
    (50.5, TypeError, "must be str or int"),
]

[gmarkall]

/ok to test 1be8209

[greptile-apps]

Greptile's behavior is changing!

From now on, if a review finishes with no comments, we will not post an additional "statistics" comment to confirm that our review found nothing to comment on. However, you can confirm that we reviewed your changes in the status check section.

_{This feature can be toggled off in your Code Review Settings by deselecting "Create a status check for each PR".}

[gmarkall]

@kaeun97 Thanks for this, and apologies for the delay here. Things got a bit backlogged over the last month and the holiday period, and I'm aiming to get a bit more back on track again now.