Most non-transform operators working with JIT

[cliffburdick]

Converted most of the element-wise operators to have JIT support by adding all the capabilities and strings. ND operators still not working yet and need some investigation.

Other changes

• Split up large unit test file that was taking long to compile
• Added GLOBAL_KERNEL capability that says whether it can operate as separate CTAs or not
• Added JIT-enabled tests when MATX_EN_JIT is enabled along with the parameter sweeps. This can take a very long time and we'll have to figure out when to enable this.

[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.

[cliffburdick]

/build

[greptile-apps]

Greptile Summary

• Added JIT compilation support to most element-wise operators by implementing SUPPORTS_JIT capability queries and JIT string generation
• Introduced GLOBAL_KERNEL capability to distinguish between global-level and block-level kernel execution patterns
• Split large test file operator_func_test.cu into 8 smaller files by category to reduce compilation time

Confidence Score: 3/5

• This PR has a critical loop bound bug that will cause incorrect execution but is otherwise well-structured
• Score reflects one critical logical error in jit_kernel.h that must be fixed before merging (loop bound in matxOpT4StrideKernelBlock), but the overall JIT infrastructure design is sound
• Pay close attention to include/matx/executors/jit_kernel.h - contains a critical loop bound bug at line 274

Important Files Changed

Filename	Overview
include/matx/core/capabilities.h	Added GLOBAL_KERNEL capability enum and attributes, changed SUPPORTS_JIT default from true to false, added JIT support check for non-MatX operators
include/matx/executors/jit_kernel.h	Split kernels into Block and global variants, updated to use CurrentCapabilities template parameter, contains loop bound bug in matxOpT4StrideKernelBlock (line 274)
include/matx/executors/jit_cuda.h	Added logic to check GLOBAL_KERNEL capability and select appropriate kernel variants, integrated new JIT infrastructure with executor

Sequence Diagram

sequenceDiagram
    participant User
    participant CUDAJITExecutor
    participant CapabilitySystem
    participant KernelProvider
    participant NVRTC
    participant GPU

    User->>CUDAJITExecutor: Exec(op)
    CUDAJITExecutor->>CapabilitySystem: Check SUPPORTS_JIT capability
    CapabilitySystem-->>CUDAJITExecutor: JIT supported
    CUDAJITExecutor->>CapabilitySystem: Query GLOBAL_KERNEL capability
    CapabilitySystem-->>CUDAJITExecutor: global_kernel flag
    CUDAJITExecutor->>KernelProvider: create_kernel_provider(sizes, jit=true, global_kernel)
    CUDAJITExecutor->>KernelProvider: find_best_launch_params(op, kernel_provider)
    KernelProvider-->>CUDAJITExecutor: ept, shm_size, block_size, groups_per_block
    CUDAJITExecutor->>CUDAJITExecutor: get_grid_dims(blocks, threads, sizes, ept)
    CUDAJITExecutor->>NVRTC: nvrtc_compile_and_run(op, sizes, blocks, threads, ept, stride, shm_size)
    NVRTC->>GPU: Launch JIT-compiled kernel
    GPU-->>User: Results

Loading

[greptile-apps]

_{91 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile
React with 👍 or 👎 to share your feedback on this new summary format}

[greptile-apps]

_{92 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile
React with 👍 or 👎 to share your feedback on this new summary format}

[cliffburdick]

91 files reviewed, 1 comment

Edit Code Review Agent Settings | Greptile React with 👍 or 👎 to share your feedback on this new summary format

Critical loop bug fixed. Please look again

[greptile-apps]

Skipped: This PR changes more files than the configured file change limit: (105 files found, 100 file limit)

[cliffburdick]

/build

[cliffburdick]

/build

[cliffburdick]

/build