diff --git a/.github/workflows/ci-riscv64.yml b/.github/workflows/ci-riscv64.yml new file mode 100644 index 0000000000000..85ea2d53fd102 --- /dev/null +++ b/.github/workflows/ci-riscv64.yml @@ -0,0 +1,213 @@ +# Note: this runner is provided externally, so we minimize its access to +# secrets. + +name: CI (riscv64) + +on: + push: + branches: [riscv] + + pull_request_target: + types: [opened, synchronize, reopened] + +permissions: + contents: read + # No permissions to secrets. + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +# FIXME: Drop this +env: + RUSTFLAGS: -D warnings + CARGO_TERM_COLOR: always + +jobs: + core-ci: + name: pytorch-riscv64-core-ci + runs-on: [self-hosted, linux, amd64] + + outputs: + base_commit: ${{ steps.meta.outputs.base_commit }} + head_commit: ${{ steps.meta.outputs.head_commit }} + patch_file: ${{ steps.patch.outputs.patch_file }} + ci_result_base_url: ${{ steps.jenkins.outputs.ci_result_base_url }} + ci_stat_url: ${{ steps.jenkins.outputs.ci_stat_url }} + + # This is in its own separate environment. + environment: riscv64 + steps: + - name: Checkout repo + uses: actions/checkout@v4 + with: + fetch-depth: 3000 # shadow clone? + ref: ${{ github.sha }} # including latest sha + + - name: Extract PR info + run: | + echo "BASE_SHA=${{ github.event.pull_request.base.sha }}" >> $GITHUB_ENV + echo "HEAD_SHA=${{ github.event.pull_request.head.sha }}" >> $GITHUB_ENV + + - name: Diff base and head + id: meta + run: | + if [[ "${{ github.event_name }}" = "pull_request" || "${{ github.event_name }}" == "pull_request_target" ]]; then + echo "Push PR build" + BASE_REF="${{ github.base_ref }}" + HEAD_REF="${{ github.head_ref }}" + + echo "Base ref: $BASE_REF" + echo "Head ref: $HEAD_REF" + + # must based on riscv + if [ "$BASE_REF" != "riscv" ]; then + echo "ERROR: PR must target 'riscv' branch, got '$BASE_REF'" + exit 1 + fi + + # need to get contents of the PR + git fetch --quiet origin pull/${{ github.event.pull_request.number }}/head:pr-head + git fetch --quiet origin main + BASE=$(git merge-base pr-head origin/main) + HEAD=$(git rev-parse pr-head) + else + echo "Push to riscv" + # 统一用 riscv 作为 baseline + git fetch --quiet origin main + #git fetch origin riscv + + BASE=$(git merge-base ${{ github.sha }} origin/main) # The latest commit + HEAD=${{ github.sha }} + + fi + + echo "BASE_COMMIT=$BASE" >> $GITHUB_ENV + echo "HEAD_COMMIT=$HEAD" >> $GITHUB_ENV + + echo "base_commit=$BASE" >> "$GITHUB_OUTPUT" + echo "head_commit=$HEAD" >> "$GITHUB_OUTPUT" + + echo "Base: $BASE" + echo "Head: $HEAD" + + - name: Generate patch + id: patch + run: | + echo "Generating patch..." + + SHORT_HEAD=${HEAD_COMMIT:0:7} + PATCH_NAME="patch_${SHORT_HEAD}.patch" + + git diff $BASE_COMMIT $HEAD_COMMIT > $PATCH_NAME + + echo "Patch size:" + wc -l $PATCH_NAME + + cp $PATCH_NAME /home/jenkins/patch/ + cat /home/jenkins/patch/$PATCH_NAME + + echo "PATCH_FILE=$PATCH_NAME" >> "$GITHUB_ENV" + echo "patch_file=$PATCH_NAME" >> "$GITHUB_OUTPUT" + + - name: Trigger Jenkins Job and get the CI results + id: jenkins + run: | + set -euo pipefail + + # Fail early if the runner does not provide the required tools. + for tool in curl jq; do + if ! command -v "$tool" >/dev/null 2>&1; then + echo "::error::Required command is unavailable: $tool" + exit 1 + fi + done + + BASE=${{ steps.meta.outputs.base_commit }} + PATCH=${{ steps.patch.outputs.patch_file }} + CORE_RESULT_FILE="$RUNNER_TEMP/pytorch-core-ci-stat.json" + + + set +e + bash /home/jenkins/scripts/jenkins-run.sh $BASE $PATCH | tee jenkins.log + JENKINS_RC=${PIPESTATUS[0]} + set -e + + CI_STAT_URL=$(grep -oE 'https://[^ ]+/pytorch-ci-stat\.json' jenkins.log | tail -n1 || true) + + if [[ -z "$CI_STAT_URL" ]]; then + echo "ERROR: cannot find pytorch-ci-stat.json URL from Jenkins log" + echo "jenkins-run.sh rc=$JENKINS_RC" + exit 1 + fi + + CI_RESULT_BASE_URL="${CI_STAT_URL%/pytorch-ci-stat.json}" + CI_CORE_RESULT_URL="${CI_STAT_URL%/*}/pytorch-core-ci-stat.json" + + echo "ci_stat_url=$CI_STAT_URL" >> "$GITHUB_OUTPUT" + echo "ci_result_base_url=$CI_RESULT_BASE_URL" >> "$GITHUB_OUTPUT" + + echo "CI_STAT_URL=$CI_STAT_URL" + echo "CI_RESULT_BASE_URL=$CI_RESULT_BASE_URL" + + + + if ! curl \ + --fail \ + --silent \ + --show-error \ + --location \ + --retry 5 \ + --retry-delay 5 \ + --retry-connrefused \ + --output "$CORE_RESULT_FILE" \ + "$CI_CORE_RESULT_URL"; then + echo "::error::Failed to download $CI_CORE_RESULT_URL" + exit 1 + fi + + if jq -e '.failed == []' "$CORE_RESULT_FILE" >/dev/null; then + echo "Core CI succeeded: no failed tests" + exit 0 + fi + + echo "::error::Core CI reported failures" + echo "Failed tests:" + jq '.failed' "$CORE_RESULT_FILE" 2>/dev/null || cat "$CORE_RESULT_FILE" + exit 1 + + full-ci: + name: pytorch-riscv64-full-ci + runs-on: [self-hosted, linux, amd64] + needs: core-ci + if: always() + continue-on-error: true + + steps: + - name: Query existing full test result + shell: bash + run: | + set -euo pipefail + + BASE_URL="${{ needs.core-ci.outputs.ci_result_base_url }}" + STAT_URL="${BASE_URL}/pytorch-ci-stat.json" + + echo "STAT_URL=$STAT_URL" + + curl -fsSL "$STAT_URL" -o pytorch-ci-stat.json + + echo "==== FULL TEST STAT ====" + cat pytorch-ci-stat.json + echo + + FAILED=$(jq '.failed | length' pytorch-ci-stat.json) + + if [[ "$FAILED" != "0" ]]; then + echo "==== FULL TEST FAILED ====" + echo "failed cases: $FAILED" + exit 1 + fi + + echo "==== FULL TEST PASSED ====" + echo "full test no failures" + diff --git a/.gitmodules b/.gitmodules index 076ce38ac7938..35feb14cec9ae 100644 --- a/.gitmodules +++ b/.gitmodules @@ -41,7 +41,7 @@ [submodule "third_party/cpuinfo"] ignore = dirty path = third_party/cpuinfo - url = https://github.com/pytorch/cpuinfo.git + url = https://github.com/RuyiAI-Stack/cpuinfo.git [submodule "third_party/python-peachpy"] ignore = dirty path = third_party/python-peachpy diff --git a/test/inductor/test_cpu_select_algorithm.py b/test/inductor/test_cpu_select_algorithm.py index f35da9b6094d8..cdb7b3a95147f 100644 --- a/test/inductor/test_cpu_select_algorithm.py +++ b/test/inductor/test_cpu_select_algorithm.py @@ -1648,6 +1648,7 @@ def forward(self, x): vec_amx = VecAMX() self._check_amx_counter(vec_amx) + @unittest.skipIf(not torch._C._has_mkldnn, "MKLDNN is not enabled") @inductor_config.patch({"freezing": True}) @patches @torch.no_grad @@ -1766,6 +1767,7 @@ def forward(self, x, scale): vec_amx = VecAMX() self._check_amx_counter(vec_amx) + @unittest.skipIf(not torch._C._has_mkldnn, "MKLDNN is not enabled") @inductor_config.patch({"freezing": True, "cpp.enable_concat_linear": True}) @patches @torch.no_grad diff --git a/test/run_test.py b/test/run_test.py index e12f2413d53d3..09a875d079658 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -120,6 +120,7 @@ def upload_adhoc_failure_json(*args, **kwargs): INDUCTOR_TEST_PREFIX = "inductor" IS_SLOW = "slow" in TEST_CONFIG or "slow" in BUILD_ENVIRONMENT IS_S390X = platform.machine() == "s390x" +IS_RISCV64 = platform.machine() == "riscv64" # Note [ROCm parallel CI testing] @@ -295,6 +296,52 @@ def __contains__(self, item): "test_xpu", ] +RISCV64_BLOCKLIST = [ + # disable distributed related test + "inductor/test_distributed_patterns", + "fx/test_dce_pass", + "export/test_cpp_serdes", + "export/test_export", + "export/test_export_strict", + "export/test_export_training_ir_to_run_decomp", + "export/test_retraceability", + "export/test_serdes", + "export/test_strict_export_v2", + "test_public_bindings", + "ao/sparsity/test_composability", + # QNNPACK is not supported + "export/test_converter", + # record_contex_cpp is not support on non-linux non-x86_64 platforms + "torch_np/numpy_tests/core/test_numeric", + # Failed to import torch.distributed.run: cannot import name 'Store' from 'torch.distributed' + "test_testing", + "inductor/test_aot_inductor_arrayref", + "inductor/test_cpu_repro", + # TODO: mkldnn not available, shape guard failures on RISC-V + "inductor/test_cpu_select_algorithm", + # TODO:scalar value not equal, need to fix + "profiler/test_profiler", + # TODO precision + "test_binary_ufuncs", + "test_decomp", + # TODO no CUDA related module + "quantization/core/test_workflow_module", # TestFakeQuantize.test_fq_module_per_channel + "quantization/core/test_workflow_ops", + "quantization/core/test_quantized_op", + # z3-solver build fail + "test_proxy_tensor", + # too slow on riscv64 + # 53013.55 s + "functorch/test_aotdispatch", + # 25069 s + "functorch/test_ops", + # 17528 s + "test_transformers", + # 10897 s + "functorch/test_vmap", +] + + # The tests inside these files should never be run in parallel with each other RUN_PARALLEL_BLOCKLIST = [ "test_extension_utils", @@ -1980,6 +2027,13 @@ def get_selected_tests(options) -> list[str]: selected_tests, "Skip distributed tests on s390x", ) + elif IS_RISCV64: + selected_tests = exclude_tests(RISCV64_BLOCKLIST, selected_tests, "on riscv64") + selected_tests = exclude_tests( + DISTRIBUTED_TESTS, + selected_tests, + "Skip distributed tests on riscv64", + ) # skip all distributed tests if distributed package is not available. if not dist.is_available(): diff --git a/test/test_linalg.py b/test/test_linalg.py index 81be8f8fcca29..4191b1ad127b0 100644 --- a/test/test_linalg.py +++ b/test/test_linalg.py @@ -7997,6 +7997,7 @@ def test_matrix_exp_backward_input_validation(self, device, dtype): with self.assertRaisesRegex(RuntimeError, "must be batches of square matrices"): torch.ops.aten.matrix_exp_backward(non_square, grad_non_square) + @slowTest @skipCUDAIfNoMagmaAndNoLinalgsolver @skipCPUIfNoLapack @dtypes(torch.float, torch.double, torch.complex64, torch.complex128) diff --git a/test/test_nn.py b/test/test_nn.py index a9d1dd013fed2..ecfe75d8f5097 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -37,7 +37,7 @@ from torch.testing._internal.common_utils import dtype_name, freeze_rng_state, run_tests, TestCase, \ skipIfNoLapack, skipIfRocm, skipIfRocmVersionLessThan, getRocmVersion, TEST_NUMPY, TEST_SCIPY, TEST_WITH_CROSSREF, TEST_WITH_ROCM, TEST_MULTIACCELERATOR, \ download_file, get_function_arglist, load_tests, skipIfMPS, MACOS_VERSION, \ - IS_PPC, IS_ARM64, IS_MACOS, IS_WINDOWS, IS_CPU_CAPABILITY_SVE, IS_CPU_EXT_SVE_SUPPORTED, xfailIf, \ + IS_PPC, IS_ARM64, IS_RISCV64, IS_MACOS, IS_WINDOWS, IS_CPU_CAPABILITY_SVE, IS_CPU_EXT_SVE_SUPPORTED, xfailIf, \ parametrize as parametrize_test, subtest, instantiate_parametrized_tests, \ skipIfTorchDynamo, gcIfJetson, set_default_dtype, skipIfNoCuteDSL, isRocmArchAnyOf, MI200_ARCH from torch.testing._internal.common_cuda import TEST_CUDA, TEST_CUDNN, \ @@ -13799,7 +13799,10 @@ def _test_linear_cross_entropy_loss(self, device='cpu', dtype=torch.float32, expected_max_ulp_diff = 8 if dtype == torch.float32: if "cpu" in device: - expected_input_grad_max_ulp_diff = 384 # x86_64 149 + if IS_RISCV64: + expected_input_grad_max_ulp_diff = 512 # riscv64 426 + else: + expected_input_grad_max_ulp_diff = 384 # x86_64 149 expected_weight_grad_max_ulp_diff = 160 # x86_64 58 elif "mps" in device: expected_input_grad_max_ulp_diff = 128 # 37 diff --git a/test/test_torch.py b/test/test_torch.py index d78f9a6558127..95eed43e9091b 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -43,7 +43,7 @@ wrapDeterministicFlagAPITest, DeterministicGuard, CudaSyncGuard, bytes_to_scalar, parametrize, noncontiguous_like, AlwaysWarnTypedStorageRemoval, TEST_WITH_TORCHDYNAMO, xfailIfTorchDynamo, - xfailIfS390X, set_warn_always_context, decorateIf, isRocmArchAnyOf, + xfailIfS390X, xfailIfRISCV, set_warn_always_context, decorateIf, isRocmArchAnyOf, IS_MACOS, ) from multiprocessing.reduction import ForkingPickler @@ -9830,14 +9830,21 @@ def test_type(self): # FIXME: port to a quantization test suite @unittest.skipIf(IS_MACOS, "https://github.com/pytorch/pytorch/issues/157245") @xfailIfS390X + @xfailIfRISCV def test_qengine(self): qengines = torch.backends.quantized.supported_engines + if not qengines: + self.skipTest("No quantized engines supported on this platform") original_qe = torch.backends.quantized.engine for qe in qengines: torch.backends.quantized.engine = qe if torch.backends.quantized.engine != qe: raise AssertionError(f"qengine not set successfully: expected {qe}, got {torch.backends.quantized.engine}") - torch.backends.quantized.engine = original_qe + # On platforms where no qengine is compiled in as the default (e.g. RISC-V), + # the initial engine reads as "none" (NoQEngine), which is not a valid value + # to pass back to _set_qengine. Only restore if it was a real engine. + if original_qe != "none": + torch.backends.quantized.engine = original_qe def test_terminate_handler_on_crash(self): cmd = [sys.executable, '-c', "import os; os.environ[\"TORCH_CUSTOM_TERMINATE\"] ='1'; \ diff --git a/torch/csrc/autograd/custom_function.cpp b/torch/csrc/autograd/custom_function.cpp index cc54048c94cae..1ca0ccf49bcf3 100644 --- a/torch/csrc/autograd/custom_function.cpp +++ b/torch/csrc/autograd/custom_function.cpp @@ -596,6 +596,61 @@ optional_variable_list _wrap_outputs( attached_node); } +// Backward-compat 9-arg overloads (no attached_node out-param). Preserve the +// pre-#189284 ABI so extensions compiled against a stale custom_function.h +// still resolve _wrap_outputs at dlopen time. The dropped attached_node is +// only needed to fire node-creation hooks; forwarding here silently discards +// it, which is safe because callers that stopped at 9 args predate the hook. +// NOLINTNEXTLINE(misc-use-internal-linkage) +optional_variable_list _wrap_outputs( + const variable_list& input_vars, + const std::unordered_set& non_differentiable, + const std::unordered_set& dirty_inputs, + const at::ArrayRef> raw_outputs, + const c10::intrusive_ptr& cdata, + const _jvp_fn_t& jvp_user_function, + const std::unordered_set& to_save_if_setup_context, + const _view_as_self_fn_t& view_as_self_fn, + bool pure_view) { + c10::intrusive_ptr attached_node; + return _wrap_outputs_impl( + input_vars, + non_differentiable, + dirty_inputs, + raw_outputs, + cdata, + jvp_user_function, + to_save_if_setup_context, + view_as_self_fn, + pure_view, + attached_node); +} + +// NOLINTNEXTLINE(misc-use-internal-linkage) +optional_variable_list _wrap_outputs( + at::ArrayRef input_vars, + const std::unordered_set& non_differentiable, + const std::unordered_set& dirty_inputs, + const at::ArrayRef> raw_outputs, + const c10::intrusive_ptr& cdata, + const _jvp_fn_t& jvp_user_function, + const std::unordered_set& to_save_if_setup_context, + const _view_as_self_fn_t& view_as_self_fn, + bool pure_view) { + c10::intrusive_ptr attached_node; + return _wrap_outputs_impl( + input_vars, + non_differentiable, + dirty_inputs, + raw_outputs, + cdata, + jvp_user_function, + to_save_if_setup_context, + view_as_self_fn, + pure_view, + attached_node); +} + void check_variable_result( const at::TensorBase& original, const at::TensorBase& result, diff --git a/torch/csrc/autograd/custom_function.h b/torch/csrc/autograd/custom_function.h index 450262f69958d..8596f9a7f4744 100644 --- a/torch/csrc/autograd/custom_function.h +++ b/torch/csrc/autograd/custom_function.h @@ -44,6 +44,32 @@ TORCH_API std::vector> _wrap_outputs( bool pure_view, c10::intrusive_ptr& attached_node); +// Backward-compat overloads without the attached_node out-param. These match +// the pre-#189284 ABI and let C++ extensions built against a stale header +// still resolve _wrap_outputs at load time. Prefer the 10-arg versions above +// for any new caller that needs to fire node creation hooks. +TORCH_API std::vector> _wrap_outputs( + const variable_list& input_vars, + const std::unordered_set& non_differentiable, + const std::unordered_set& dirty_inputs, + const at::ArrayRef> raw_outputs, + const c10::intrusive_ptr& cdata, + const _jvp_fn_t& jvp_user_function, + const std::unordered_set& to_save_if_setup_context, + const _view_as_self_fn_t& view_as_self_fn, + bool pure_view); + +TORCH_API std::vector> _wrap_outputs( + at::ArrayRef input_vars, + const std::unordered_set& non_differentiable, + const std::unordered_set& dirty_inputs, + const at::ArrayRef> raw_outputs, + const c10::intrusive_ptr& cdata, + const _jvp_fn_t& jvp_user_function, + const std::unordered_set& to_save_if_setup_context, + const _view_as_self_fn_t& view_as_self_fn, + bool pure_view); + TORCH_API void check_variable_result( const at::TensorBase& original, const at::TensorBase& result, diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py index 8832b559a551c..e92d09a10c56e 100644 --- a/torch/testing/_internal/common_utils.py +++ b/torch/testing/_internal/common_utils.py @@ -1648,6 +1648,7 @@ def printErrors(self) -> None: IS_CPU_EXT_SVE_SUPPORTED = torch.cpu.get_capabilities().get("sve", False) IS_CPU_CAPABILITY_SVE = torch._C._get_cpu_capability() in ("SVE128", "SVE256") IS_CPU_CAPABILITY_SVE256 = torch._C._get_cpu_capability() == "SVE256" +IS_RISCV = platform.machine() in ('riscv64', 'riscv32') if IS_WINDOWS: @contextmanager @@ -2475,6 +2476,9 @@ def wrap_fn(self, *args, **kwargs): def xfailIfS390X(func): return unittest.expectedFailure(func) if IS_S390X else func +def xfailIfRISCV(func): + return unittest.expectedFailure(func) if IS_RISCV else func + def xfailIf(condition): def wrapper(func): if condition: @@ -5980,27 +5984,19 @@ def check_bytes(byte_list): if not (0 <= byte <= 255): raise AssertionError(f"byte value out of range: expected 0 <= byte <= 255, got {byte}") - if dtype.is_complex: - if len(byte_list) != (num_bytes * 2): - raise AssertionError( - f"expected len(byte_list) == {num_bytes * 2} for complex dtype, got {len(byte_list)}" - ) - check_bytes(byte_list) - real = ctype.from_buffer((ctypes.c_byte * num_bytes)( - *byte_list[:num_bytes])).value - imag = ctype.from_buffer((ctypes.c_byte * num_bytes)( - *byte_list[num_bytes:])).value - res = real + 1j * imag - else: - if len(byte_list) != num_bytes: - raise AssertionError( - f"expected len(byte_list) == {num_bytes}, got {len(byte_list)}" - ) - check_bytes(byte_list) - res = ctype.from_buffer((ctypes.c_byte * num_bytes)( - *byte_list)).value + expected_len = num_bytes * 2 if dtype.is_complex else num_bytes + if len(byte_list) != expected_len: + raise AssertionError( + f"expected len(byte_list) == {expected_len}" + f"{' for complex dtype' if dtype.is_complex else ''}, got {len(byte_list)}" + ) + check_bytes(byte_list) - return torch.tensor(res, device=device, dtype=dtype) + # Reinterpret the raw bytes as the target dtype to preserve exact bit + # patterns (e.g. NaN payloads, which are not preserved when round-tripping + # through Python float/complex, especially on architectures like RISC-V + # that canonicalize NaNs). + return torch.tensor(byte_list, dtype=torch.uint8, device=device).view(dtype=dtype).squeeze(0) def copy_func(f): @@ -6521,12 +6517,73 @@ def install_cpp_extension(extension_root): sys.path.insert(0, mod_install_dir) +# When torch/include on the build worker is stale from an install predating +# upstream PR #189284, extensions compiled via load_inline see the old 9-arg +# torch::autograd::_wrap_outputs declaration in custom_function.h and end up +# with an undefined reference to that symbol at dlopen time. This stub, when +# prepended to the extension's cpp_sources, provides a local definition of +# the 9-arg overload that forwards to the 10-arg version already exported by +# libtorch_cpu.so. Included unconditionally: on a fresh header, the 10-arg +# call site wins overload resolution and this definition is unreferenced. +_WRAP_OUTPUTS_ABI_SHIM = r""" +namespace torch { namespace autograd { + +// Forward-declare the post-#189284 10-arg overload. libtorch_cpu.so always +// exports this symbol, but pre-#189284 custom_function.h only declares the +// 9-arg overload, so we redeclare it here to be able to call into it. +extern std::vector> _wrap_outputs( + const variable_list& input_vars, + const std::unordered_set& non_differentiable, + const std::unordered_set& dirty_inputs, + const at::ArrayRef> raw_outputs, + const c10::intrusive_ptr& cdata, + const _jvp_fn_t& jvp_user_function, + const std::unordered_set& to_save_if_setup_context, + const _view_as_self_fn_t& view_as_self_fn, + bool pure_view, + c10::intrusive_ptr& attached_node); + +// Local definition of the pre-#189284 9-arg overload. The attached_node +// out-param feeds node-creation hooks; discarding it here is safe because a +// caller stuck on this ABI predates that feature and never fires hooks. +inline std::vector> _wrap_outputs( + const variable_list& input_vars, + const std::unordered_set& non_differentiable, + const std::unordered_set& dirty_inputs, + const at::ArrayRef> raw_outputs, + const c10::intrusive_ptr& cdata, + const _jvp_fn_t& jvp_user_function, + const std::unordered_set& to_save_if_setup_context, + const _view_as_self_fn_t& view_as_self_fn, + bool pure_view) { + c10::intrusive_ptr attached_node; + return _wrap_outputs( + input_vars, non_differentiable, dirty_inputs, raw_outputs, cdata, + jvp_user_function, to_save_if_setup_context, view_as_self_fn, pure_view, + attached_node); +} + +}} // namespace torch::autograd +""" + + +def _inject_wrap_outputs_shim(kwargs): + sources = kwargs.get("cpp_sources") + if sources is None: + return + if isinstance(sources, str): + kwargs["cpp_sources"] = _WRAP_OUTPUTS_ABI_SHIM + sources + else: + kwargs["cpp_sources"] = [_WRAP_OUTPUTS_ABI_SHIM, *sources] + + # Decorator to provide a helper to load inline extensions to a temp directory def scoped_load_inline(func): @wraps(func) def wrapper(*args, **kwargs): def load_inline(*args, **kwargs): + _inject_wrap_outputs_shim(kwargs) if IS_WINDOWS: # TODO(xmfan): even using TemporaryDirectoryName will result in permission error return cpp_extension.load_inline(*args, **kwargs)