Customizing a PyTorch operation

At a glance

Project structure

Source bundle/
├── softshrink.py
├── run_sample.py
├── CustomSoftshrink.h
├── CustomSoftshrink.mm
└── compiler.py

Structure observations

• The entry/composition boundary and renderer or operation boundary are separate in the source; Objective-C++ custom op is the principal feature coordinator.
• GPU kernel source is embedded in a native header and compiled into a Metal library at runtime.
• The tree above is intentionally pruned to composition, resource, and shader files.

Overall architecture

flowchart LR
    N1["Python model/script"]
    N2["Objective-C++ custom op"]
    N3["Python-to-native custom operator bridge"]
    N4["runtime-compiled Metal kernel on PyTorch's MPS stream"]
    N1 --> N2
    N2 --> N3
    N3 --> N4

Reference code

softshrink.py:27 — feature handoff or setup anchor

    def forward(self, input):
        return compiled_lib.mps_softshrink(input, self.lambd)

Interpretation

The Python module calls a pybind11 extension. Objective-C++ validates the tensor, compiles the embedded Metal kernel, borrows PyTorch’s MPS command buffer/queue, and encodes the compute dispatch.

Ownership and state

classDiagram
    run_sample *-- CustomMPSSoftshrinkModel : custom_model

Ownership evidence

run_sample.py:24 — run_sample creates for a local scope custom_model

    custom_model = CustomMPSSoftshrinkModel().to(mps_device)

    # Measures time.
    for _ in range(100):
        start = time.time()
        default_model.forward(x)
        torch.mps.synchronize()
        default_softshrink += time.time() - start

Class and protocol design

No local protocol abstraction is present; the sample uses concrete framework, operation, or model types at this boundary.

Access control

This sample does not use Swift private, fileprivate, or public for its native boundary. Objective-C/Objective-C++ use header versus implementation placement, C++ uses access specifiers/linkage, Python uses module conventions, and Metal entry points cross a compiled-library boundary; these are not Swift access levels.

Logic ownership and placement

Shader boundary reference

CustomSoftshrink.h:19 — representative GPU entry/helper

kernel void softshrink_kernel(constant T*     input  [[buffer(0)]],
                              device   T*     output [[buffer(1)]],
                              constant float& lambda [[buffer(2)]],
                              uint index [[thread_position_in_grid]]) {
    output[index] = input[index] >  lambda ? input[index] - lambda :
                    input[index] < -lambda ? input[index] + lambda : 0;
}

Design patterns

Naming conventions

• Role suffixes make ownership visible: CustomMPSSoftshrinkModel, SoftshrinkModel.
• Method names describe setup or encoding actions: __init__, forward, extra_repr, test_speedup, test_correctness, test_softshrink.
• Shader entry points use stage/operation names: softshrink_kernel.
• Names favor concrete domain roles and target-local types; no broad public library namespace is introduced.

Architecture takeaways

• Keep Python model/script focused on composition; Objective-C++ custom op is the owner of Python-to-native custom operator bridge.
• Treat runtime-compiled Metal kernel on PyTorch’s MPS stream as a separate execution/compilation boundary with explicit resource and data-layout contracts.
• The verified ownership edge is run_sample → custom_model; broader exclusive ownership is not inferred.
• Access-control rationale follows concrete language boundaries rather than translating every header or shader symbol into Swift terms.

Source map

• Official Apple documentation