Skip to content

Fix GPU lambda constants for casted jit_variable captures #535

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ZwFink merged 2 commits into from
Oct 9, 2026
Merged

Fix GPU lambda constants for casted jit_variable captures #535

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2112975
Add test for casted jit_variable lambda captures
ZwFink Sep 28, 2026
eedef17
Use closure slot type in lambda schema metadata
ZwFink Sep 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions src/pass/LambdaAnalysis.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -552,8 +552,11 @@ class LambdaPassImpl {
for (const auto &JitVarInfo : It->second) {
auto *SlotC = dyn_cast<ConstantInt>(JitVarInfo.Slot);
auto *OffsetC = dyn_cast<ConstantInt>(JitVarInfo.Offset);
auto RCType = getRCTypeForLLVMType(JitVarInfo.Type);
if (!SlotC || !OffsetC || !RCType)
if (!SlotC || !OffsetC)
reportFatalError("Failed to emit lambda schema metadata");
auto RCType = getRCTypeForLLVMType(
LambdaStorageType->getElementType(SlotC->getZExtValue()));
if (!RCType)
reportFatalError("Failed to emit lambda schema metadata");

LLVMContext &Ctx = M.getContext();
Expand Down
1 change: 1 addition & 0 deletions tests/cpu/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ CREATE_CPU_TEST(types_api types_api.cpp)
CREATE_CPU_TEST(lambda lambda.cpp)
CREATE_CPU_TEST(lambda_ptrtoint_capture lambda_ptrtoint_capture.cpp)
CREATE_CPU_TEST(lambda_def lambda_def.cpp)
CREATE_CPU_TEST(lambda_cast_capture lambda_cast_capture.cpp)
CREATE_CPU_TEST(lambda_nested lambda_nested.cpp)
CREATE_CPU_TEST(annot_nested annot_nested.cpp)
CREATE_CPU_TEST(lambda_def_register_once lambda_def_register_once.cpp)
Expand Down
52 changes: 52 additions & 0 deletions tests/cpu/lambda_cast_capture.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
// clang-format off
// RUN: rm -rf "%t.$$.proteus"
// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" %build/lambda_cast_capture | %FILECHECK %s --check-prefixes=CHECK
// RUN: rm -rf "%t.$$.proteus"
// clang-format on

#include <cstdio>

#include <proteus/JitInterface.h>

template <typename F> void run(F &&Func) { proteus::register_lambda(Func)(); }

Check warning on line 11 in tests/cpu/lambda_cast_capture.cpp

View workflow job for this annotation

GitHub Actions / cpp-linter

tests/cpu/lambda_cast_capture.cpp:11:28 [misc-use-internal-linkage] 

function 'run' can be made static or moved into an anonymous namespace to enforce internal linkage

int main() {
int I = -5;
unsigned U = 0x80000000u;
float F = 1.5f;
double D = 2.7;

auto IntToLong =
[=, X = (long)proteus::jit_variable(I)]()
__attribute__((annotate("jit"))) { printf("IntToLong %ld\n", X); };
run(IntToLong);

auto UnsignedToULong = [=, X = (unsigned long)proteus::jit_variable(U)]()
__attribute__((annotate("jit"))) {
printf("UnsignedToULong %lu\n", X);
};
run(UnsignedToULong);

auto FloatToDouble =
[=, X = (double)proteus::jit_variable(F)]()
__attribute__((annotate("jit"))) { printf("FloatToDouble %f\n", X); };
run(FloatToDouble);

auto IntToDouble =
[=, X = (double)proteus::jit_variable(I)]()
__attribute__((annotate("jit"))) { printf("IntToDouble %f\n", X); };
run(IntToDouble);

auto DoubleToInt =
[=, X = (int)proteus::jit_variable(D)]()
__attribute__((annotate("jit"))) { printf("DoubleToInt %d\n", X); };
run(DoubleToInt);

return 0;
}

// CHECK: IntToLong -5
// CHECK: UnsignedToULong 2147483648
// CHECK: FloatToDouble 1.500000
// CHECK: IntToDouble -5.000000
// CHECK: DoubleToInt 2
1 change: 1 addition & 0 deletions tests/gpu/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,7 @@ CREATE_GPU_TEST(shared_array shared_array.cpp)
CREATE_GPU_TEST(enable_disable enable_disable.cpp)
CREATE_GPU_TEST(lambda_multiple lambda_multiple.cpp)
CREATE_GPU_TEST(lambda_def lambda_def.cpp)
CREATE_GPU_TEST(lambda_cast_capture lambda_cast_capture.cpp)
CREATE_GPU_TEST(lambda_factory lambda_factory.cpp)
CREATE_GPU_TEST(lambda_factory_forwarded lambda_factory_forwarded.cpp)
CREATE_GPU_TEST(lambda_thread_privatizer lambda_thread_privatizer.cpp)
Expand Down
63 changes: 63 additions & 0 deletions tests/gpu/lambda_cast_capture.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
// clang-format off
// RUN: rm -rf "%t.$$.proteus"
// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" %build/lambda_cast_capture.%ext | %FILECHECK %s --check-prefixes=CHECK
// RUN: rm -rf "%t.$$.proteus"
// clang-format on

#include <cstdio>

#include "gpu_common.h"
#include <proteus/JitInterface.h>

template <typename T>
__global__ __attribute__((annotate("jit"))) void kernel(T LB) {

Check warning on line 13 in tests/gpu/lambda_cast_capture.cpp

View workflow job for this annotation

GitHub Actions / cpp-linter

tests/gpu/lambda_cast_capture.cpp:13:50 [misc-use-internal-linkage] 

function 'kernel' can be made static or moved into an anonymous namespace to enforce internal linkage
LB();
}

template <typename T> void run(T &&LB) {

Check warning on line 17 in tests/gpu/lambda_cast_capture.cpp

View workflow job for this annotation

GitHub Actions / cpp-linter

tests/gpu/lambda_cast_capture.cpp:17:28 [misc-use-internal-linkage] 

function 'run' can be made static or moved into an anonymous namespace to enforce internal linkage
kernel<<<1, 1>>>(proteus::register_lambda(LB));
gpuErrCheck(gpuDeviceSynchronize());
fflush(stdout);
}

int main() {
int I = -5;
unsigned U = 0x80000000u;
float F = 1.5f;
double D = 2.7;

auto IntToLong =
[=, X = (long)proteus::jit_variable(I)] __device__()
__attribute__((annotate("jit"))) { printf("IntToLong %ld\n", X); };
run(IntToLong);

auto UnsignedToULong =
[=, X = (unsigned long)proteus::jit_variable(U)] __device__()
__attribute__((annotate("jit"))) {
printf("UnsignedToULong %lu\n", X);
};
run(UnsignedToULong);

auto FloatToDouble =
[=, X = (double)proteus::jit_variable(F)] __device__()
__attribute__((annotate("jit"))) { printf("FloatToDouble %f\n", X); };
run(FloatToDouble);

auto IntToDouble =
[=, X = (double)proteus::jit_variable(I)] __device__()
__attribute__((annotate("jit"))) { printf("IntToDouble %f\n", X); };
run(IntToDouble);

auto DoubleToInt =
[=, X = (int)proteus::jit_variable(D)] __device__()
__attribute__((annotate("jit"))) { printf("DoubleToInt %d\n", X); };
run(DoubleToInt);

return 0;
}

// CHECK: IntToLong -5
// CHECK: UnsignedToULong 2147483648
// CHECK: FloatToDouble 1.500000
// CHECK: IntToDouble -5.000000
// CHECK: DoubleToInt 2
Loading