[PyTorch] Centralize TE↔torch binary boundary in a single facade file

qa/L0_pytorch_lint/check_torch_boundary.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Enforce the single TE<->PyTorch binary boundary.
+
+Only ``transformer_engine/pytorch/csrc/torch_backend.{h,cpp}`` may include
+libtorch/ATen/c10 headers or name ``at::`` / ``c10::`` / ``c10d::`` /
+``torch::`` symbols. Every other translation unit must talk to PyTorch through
+the aliases and free functions that ``torch_backend.h`` exposes.
+
+Files that have not been migrated yet are grandfathered in via
+``torch_boundary_allowlist.txt`` (paths relative to the csrc dir). The guard
+fails if:
+  * a non-boundary, non-allowlisted file touches the torch ABI, or
+  * an allowlisted file no longer touches it (stale entry -> remove it), or
+  * an allowlisted path does not exist.
+
+Comments and string literals are stripped before scanning, so mentioning the
+tokens in a comment is fine.
+
+Usage: check_torch_boundary.py [TE_ROOT]   (default: $TE_PATH or cwd)
+"""
+
+import os
+import re
+import sys
+from pathlib import Path
+
+CSRC_REL = "transformer_engine/pytorch/csrc"
+BOUNDARY = {"torch_backend.h", "torch_backend.cpp"}
+SOURCE_SUFFIXES = {".cpp", ".h", ".hpp", ".cuh", ".cu"}
+
+# Include of a libtorch/ATen/c10 header, or a qualified at::/c10::/c10d::/torch:: name.
+INCLUDE_RE = re.compile(r'#\s*include\s*[<"](?:torch|ATen|c10)/')
+SYMBOL_RE = re.compile(r'\b(?:at|c10|c10d|torch)::')
+
+
+def strip_comments_and_strings(text: str) -> str:
+    """Blank out // and /* */ comments and "..."/'...' literals (newlines kept)."""
+    out = []
+    i, n = 0, len(text)
+    while i < n:
+        c = text[i]
+        two = text[i:i + 2]
+        if two == "//":
+            i += 2
+            while i < n and text[i] != "\n":
+                i += 1
+        elif two == "/*":
+            i += 2
+            while i < n and text[i:i + 2] != "*/":
+                out.append("\n" if text[i] == "\n" else " ")
+                i += 1
+            i += 2
+        elif c in "\"'":
+            quote = c
+            out.append(" ")
+            i += 1
+            while i < n and text[i] != quote:
+                if text[i] == "\\" and i + 1 < n:
+                    i += 1
+                out.append("\n" if text[i] == "\n" else " ")
+                i += 1
+            i += 1
+            out.append(" ")
+        else:
+            out.append(c)
+            i += 1
+    return "".join(out)
+
+
+def scan(path: Path):
+    """Return list of (lineno, text) lines that touch the torch ABI."""
+    raw = path.read_text(encoding="utf-8", errors="replace")
+    code = strip_comments_and_strings(raw)
+    hits = []
+    for lineno, line in enumerate(code.splitlines(), start=1):
+        if INCLUDE_RE.search(line) or SYMBOL_RE.search(line):
+            hits.append(lineno)
+    return hits
+
+
+def main() -> int:
+    root = Path(sys.argv[1] if len(sys.argv) > 1 else os.environ.get("TE_PATH", ".")).resolve()
+    csrc = root / CSRC_REL
+    if not csrc.is_dir():
+        print(f"error: csrc dir not found: {csrc}", file=sys.stderr)
+        return 2
+
+    allowlist_file = root / "qa/L0_pytorch_lint/torch_boundary_allowlist.txt"
+    allowlist = set()
+    if allowlist_file.is_file():
+        for line in allowlist_file.read_text().splitlines():
+            line = line.split("#", 1)[0].strip()
+            if line:
+                allowlist.add(line)
+
+    violations = []          # (rel, lineno) touching torch outside the boundary
+    still_allowlisted = set()  # allowlisted files that still touch torch (expected)
+
+    for path in sorted(csrc.rglob("*")):
+        if path.suffix not in SOURCE_SUFFIXES or not path.is_file():
+            continue
+        rel = path.relative_to(csrc).as_posix()
+        if rel in BOUNDARY:
+            continue
+        hits = scan(path)
+        if not hits:
+            continue
+        if rel in allowlist:
+            still_allowlisted.add(rel)
+        else:
+            violations.extend((rel, ln) for ln in hits)
+
+    stale = sorted(allowlist - still_allowlisted)
+    missing = sorted(p for p in allowlist if not (csrc / p).is_file())
+
+    ok = True
+    if violations:
+        ok = False
+        print("TE<->torch boundary violations (route these through torch_backend.h):")
+        for rel, ln in violations:
+            print(f"  {CSRC_REL}/{rel}:{ln}")
+    if stale:
+        ok = False
+        print("\nStale allowlist entries (migrated -- remove from "
+              "torch_boundary_allowlist.txt):")
+        for rel in stale:
+            print(f"  {rel}")
+    if missing:
+        ok = False
+        print("\nAllowlist entries pointing at non-existent files (remove them):")
+        for rel in missing:
+            print(f"  {rel}")
+
+    if ok:
+        print(f"torch boundary OK: only torch_backend.* touches the torch ABI "
+              f"({len(still_allowlisted)} file(s) still pending migration).")
+        return 0
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

qa/L0_pytorch_lint/test.sh

@@ -15,6 +15,8 @@ then
   echo "Checking C++ files"
   python3 -m cpplint --recursive --exclude=transformer_engine/common/include --exclude=transformer_engine/build_tools/build transformer_engine/common
   python3 -m cpplint --recursive transformer_engine/pytorch
+  echo "Checking TE<->PyTorch binary boundary"
+  python3 qa/L0_pytorch_lint/check_torch_boundary.py "$TE_PATH"
 fi
 if [ -z "${CPP_ONLY}" ]
 then

qa/L0_pytorch_lint/torch_boundary_allowlist.txt

@@ -0,0 +1,11 @@
+# Files in transformer_engine/pytorch/csrc that still talk to the PyTorch ABI
+# (at::/c10::/torch::) directly instead of going through torch_backend.h.
+#
+# This list only shrinks: when a file is migrated to the torch_backend.h
+# facade, remove it here. The guard (check_torch_boundary.py) fails if a listed
+# file no longer touches the torch ABI (stale) or if an unlisted file starts to.
+#
+# Paths are relative to transformer_engine/pytorch/csrc/.
+#
+# STATUS: empty -- the entire pytorch/csrc extension now routes every torch/ATen/
+# c10 access through torch_backend.h. Nothing is grandfathered. Keep it that way.

transformer_engine/pytorch/csrc/common.cpp

@@ -40,15 +40,15 @@ std::array<size_t, 2> get_2d_dims(NVTEShape shape, bool transpose) {
   }
 }
 
-std::vector<size_t> getTensorShape(const at::Tensor& t) {
+std::vector<size_t> getTensorShape(const Tensor& t) {
   std::vector<size_t> shape;
   for (auto s : t.sizes()) {
     shape.push_back(s);
   }
   return shape;
 }
 
-NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape) {
+NVTEShape convertTorchShape(const IntArrayRef torch_shape) {
   NVTEShape ret;
   ret.ndim = torch_shape.size();
   constexpr int max_dimensions = sizeof(ret.data) / sizeof(size_t);
@@ -132,7 +132,7 @@ TensorWrapper makeTransformerEngineTensor(py::handle tensor, py::handle quantize
              "Unexpected quantization params type.");
 
   // Regular pyTorch tensor
-  at::Tensor torch_tensor = tensor.cast<at::Tensor>();
+  Tensor torch_tensor = tensor.cast<Tensor>();
 
   // #TODO (pgadzinski) - needed in attention for non-contiguous tensors.
   //if (!torch_tensor.is_contiguous()) {
@@ -156,7 +156,7 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(
   return transformer_engine::TensorWrapper(data_ptr, shape, type);
 }
 
-transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor) {
+transformer_engine::TensorWrapper makeTransformerEngineTensor(Tensor tensor) {
   transformer_engine::DType dtype = GetTransformerEngineDType(tensor.scalar_type());
   std::vector<size_t> shape;
   for (auto s : tensor.sizes()) {
@@ -167,7 +167,7 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor)
 
 std::tuple<std::vector<transformer_engine::TensorWrapper>, std::vector<std::vector<NVTETensor>>,
            std::vector<NVTETensor*>, size_t, size_t>
-makeTransformerEngineTensorList(std::vector<std::vector<at::Tensor>> at_tensor_lists) {
+makeTransformerEngineTensorList(std::vector<std::vector<Tensor>> at_tensor_lists) {
   size_t num_lists = at_tensor_lists.size();
 
   NVTE_CHECK(num_lists > 0, "List of tensors is empty.");
@@ -238,18 +238,18 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(
   return ret;
 }
 
-transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor, at::Tensor amax,
-                                                              const at::Tensor scale,
-                                                              at::Tensor scale_inv,
+transformer_engine::TensorWrapper makeTransformerEngineTensor(Tensor tensor, Tensor amax,
+                                                              const Tensor scale,
+                                                              Tensor scale_inv,
                                                               NVTEScalingMode scaling_mode) {
   transformer_engine::DType dtype = GetTransformerEngineDType(tensor.scalar_type());
 
   auto tensor_shape = getTensorShape(tensor);
   auto scale_inv_shape = getTensorShape(scale_inv);
 
-  NVTE_CHECK(amax.scalar_type() == at::kFloat);
-  NVTE_CHECK(scale.scalar_type() == at::kFloat);
-  NVTE_CHECK(scale_inv.scalar_type() == at::kFloat);
+  NVTE_CHECK(GetTransformerEngineDType(amax.scalar_type()) == DType::kFloat32);
+  NVTE_CHECK(GetTransformerEngineDType(scale.scalar_type()) == DType::kFloat32);
+  NVTE_CHECK(GetTransformerEngineDType(scale_inv.scalar_type()) == DType::kFloat32);
 
   return makeTransformerEngineTensor(tensor.data_ptr(), tensor_shape, dtype, amax.data_ptr(),
                                      scale.data_ptr(), scale_inv.data_ptr(), scale_inv_shape,
@@ -286,45 +286,37 @@ std::vector<size_t> nvte_shape_to_vector(const NVTEShape& nvte_shape) {
   return shape;
 }
 
-at::Tensor allocateSpace(const std::vector<size_t>& shape, const transformer_engine::DType type,
+Tensor allocateSpace(const std::vector<size_t>& shape, const transformer_engine::DType type,
                          bool init_to_zeros) {
   std::vector<int64_t> shape_int64(shape.begin(), shape.end());
-  c10::IntArrayRef ar_shape(shape_int64);
-  if (init_to_zeros) {
-    return at::zeros(ar_shape, at::CUDA(GetATenDType(type)));
-  } else {
-    return at::empty(ar_shape, at::CUDA(GetATenDType(type)));
-  }
+  return new_cuda_tensor(shape_int64, GetATenDType(type), init_to_zeros);
 }
 
-at::Tensor allocateSpace(const NVTEShape& shape, const transformer_engine::DType type,
+Tensor allocateSpace(const NVTEShape& shape, const transformer_engine::DType type,
                          bool init_to_zeros) {
   auto size = shape.ndim;
-  if (size == 2 && init_to_zeros) {
-    return at::zeros({static_cast<int64_t>(shape.data[0]), static_cast<int64_t>(shape.data[1])},
-                     at::CUDA(GetATenDType(type)));
-  } else if (size == 2) {
-    return at::empty({static_cast<int64_t>(shape.data[0]), static_cast<int64_t>(shape.data[1])},
-                     at::CUDA(GetATenDType(type)));
-  } else if (size == 1 && init_to_zeros) {
-    return at::zeros({static_cast<int64_t>(shape.data[0])}, at::CUDA(GetATenDType(type)));
+  if (size == 2) {
+    return new_cuda_tensor(
+        {static_cast<int64_t>(shape.data[0]), static_cast<int64_t>(shape.data[1])},
+        GetATenDType(type), init_to_zeros);
   } else if (size == 1) {
-    return at::empty({static_cast<int64_t>(shape.data[0])}, at::CUDA(GetATenDType(type)));
+    return new_cuda_tensor({static_cast<int64_t>(shape.data[0])}, GetATenDType(type),
+                           init_to_zeros);
   }
   NVTE_ERROR("Unsupported tensor allocation: ndim=", size, ", init_to_zeros=", init_to_zeros,
              ". Only 1D and 2D tensors are supported.");
 }
 
-at::Tensor allocateTorchTensor(int M, int N, transformer_engine::DType dtype) {
-  return at::empty({static_cast<int64_t>(M), static_cast<int64_t>(N)},
-                   at::CUDA(GetATenDType(dtype)));
+Tensor allocateTorchTensor(int M, int N, transformer_engine::DType dtype) {
+  return new_cuda_tensor({static_cast<int64_t>(M), static_cast<int64_t>(N)}, GetATenDType(dtype),
+                         /*zero_init=*/false);
 }
 
-at::Tensor allocateTorchTensor(int M, transformer_engine::DType dtype) {
-  return at::empty({static_cast<int64_t>(M)}, at::CUDA(GetATenDType(dtype)));
+Tensor allocateTorchTensor(int M, transformer_engine::DType dtype) {
+  return new_cuda_tensor({static_cast<int64_t>(M)}, GetATenDType(dtype), /*zero_init=*/false);
 }
 
-void* getDataPtr(at::Tensor tensor, int offset) {
+void* getDataPtr(Tensor tensor, int offset) {
   void* dptr = nullptr;
   if (tensor.numel() > 0) {
     dptr = tensor.data_ptr();
@@ -348,17 +340,17 @@ size_t roundup(size_t value, size_t multiple) {
 
 size_t ceildiv(size_t numer, size_t denom) { return (numer + denom - 1) / denom; }
 
-void philox_unpack(at::PhiloxCudaState arg, int64_t* rng_state_ptr) {
+void philox_unpack(PhiloxCudaState arg, int64_t* rng_state_ptr) {
   NVTE_SCOPED_GIL_RELEASE({
     nvte_extract_seed_and_offset(rng_state_ptr, arg.captured_, arg.seed_.ptr, arg.seed_.val,
                                  arg.offset_.ptr, arg.offset_.val, arg.offset_intragraph_,
-                                 at::cuda::getCurrentCUDAStream());
+                                 getCurrentCUDAStream());
   });
 }
 
 // extract PhiloxCudaState from CUDA random number generator
-at::PhiloxCudaState init_philox_state(at::CUDAGeneratorImpl* gen, size_t elts_per_thread) {
-  at::PhiloxCudaState philox_args;
+PhiloxCudaState init_philox_state(CUDAGeneratorImpl* gen, size_t elts_per_thread) {
+  PhiloxCudaState philox_args;
   std::lock_guard<std::mutex> lock(gen->mutex_);
   philox_args = gen->philox_cuda_state(elts_per_thread);
   return philox_args;