[STF] Locality-domain tests: single-domain devices have nothing to split

cudax/include/cuda/experimental/__places/exec/locality_domain.cuh

@@ -224,7 +224,9 @@ inline bool locality_domain_memory_disabled()
  *
  * Each SM split method (see `locality_domain_sm_split`) gets its own set of
  * green contexts, so places built with different methods for the same domain
- * are distinct places with distinct SM partitions.
+ * are distinct places with distinct SM partitions. On a device with a single
+ * domain (or when the driver cannot report domains) the method is ignored and
+ * every method maps to one whole-device green context.
  *
  * Entries are created lazily per (device, split method) and kept alive for
  * the process, which also guarantees the (non-owning)
@@ -253,8 +255,10 @@ public:
     // Whole-device degrade ignores the split method (documented contract).
     // Canonicalize the cache key so every method resolves to the SAME green
     // context, stream pool, and execution-place identity instead of one
-    // whole-device context per requested method.
-    if (locality_domain_native_raw_count(dev_id) == 0)
+    // whole-device context per requested method. A device with a single
+    // domain degrades the same way: the domain is the device, and there is
+    // nothing to split.
+    if (locality_domain_native_raw_count(dev_id) <= 1)
     {
       split = locality_domain_sm_split::backfill;
     }
@@ -282,15 +286,19 @@ private:
     // The public count never reports 0: when the driver cannot answer the
     // locality-domain query, the device degrades to a single whole-device
     // domain, so we build one green context spanning the full SM resource
-    // instead of splitting by domain.
+    // instead of splitting by domain. A device with a single domain is
+    // handled the same way: asking the driver to split by domain there is
+    // pointless at best, and a driver without real locality-domain support
+    // (e.g. a forward-compat user-mode driver on an older kernel driver) is
+    // known to hand back an unsplit device without reporting an error.
     const int raw_domains = locality_domain_native_raw_count(dev_id);
     const int num_domains = (raw_domains > 0) ? raw_domains : 1;
 
     CUdevResource sm_resource;
     cuda_try(cuDeviceGetDevResource(device, &sm_resource, CU_DEV_RESOURCE_TYPE_SM));
 
     ::std::vector<CUdevResource> domain_sms(num_domains);
-    if (raw_domains == 0)
+    if (raw_domains <= 1)
     {
       // Whole-device degrade: the single domain covers all SMs, whatever the
       // requested split method.
@@ -895,10 +903,11 @@ inline data_place data_place::locality_domain(int dev_id, int domain_id)
  * The SM split method (`split`, native backend only) selects how the place's
  * SM partition is carved out of the device -- whole-device coverage with the
  * default `backfill`, or strictly per-domain partitions with `aligned` /
- * `fine`; see `locality_domain_sm_split` for the tradeoffs. Places built with
- * different methods for the same domain are distinct places (distinct green
- * contexts) sharing the same affine data place. The fallback backend and the
- * fake-topology override accept and ignore the method.
+ * `fine`; see `locality_domain_sm_split` for the tradeoffs. On a device with
+ * several domains, places built with different methods for the same domain
+ * are distinct places (distinct green contexts) sharing the same affine data
+ * place. Single-domain devices, the fallback backend and the fake-topology
+ * override accept and ignore the method.
  */
 inline exec_place exec_place::locality_domain(const locality_domain_view& view, locality_domain_sm_split split)
 {

cudax/test/stf/places/locality_domain_alloc.cu

@@ -70,12 +70,14 @@ int main()
     // With the native backend, the allocation must report the requested
     // domain ordinal (unless localization is disabled, or the fake topology
     // override is active, in which case it must report "not localized",
-    // i.e. -1).
+    // i.e. -1). On a single-domain device the driver documents that a
+    // localized allocation on domain 0 is a plain device allocation, which
+    // also reports -1.
     int ordinal = -2;
     cuda_safe_call(cuPointerGetAttribute(
       &ordinal, CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL, reinterpret_cast<CUdeviceptr>(ptr)));
     if (std::getenv("CUDASTF_DISABLE_LOCALIZED_MEMORY") != nullptr
-        || std::getenv("CUDASTF_FAKE_LOCALITY_DOMAINS") != nullptr)
+        || std::getenv("CUDASTF_FAKE_LOCALITY_DOMAINS") != nullptr || ndomains == 1)
     {
       EXPECT(ordinal == -1);
     }

cudax/test/stf/places/locality_domain_split_method.cu

@@ -26,6 +26,7 @@
 #include <cuda/experimental/stf.cuh>
 
 #include <cstdio>
+#include <cstdlib>
 
 using namespace cuda::experimental::stf;
 
@@ -88,75 +89,89 @@ int main()
          == exec_place::locality_domain(dev, 0, locality_domain_sm_split::backfill));
 
 #if _CCCL_CTK_AT_LEAST(13, 4) && !defined(CUDAX_PLACES_FORCE_LOCALITY_DOMAIN_FALLBACK)
-  namespace places = ::cuda::experimental::places;
-
-  // (Native backend only: the whole-device fallback and the fake-topology
-  // override deliberately ignore the split method, so their places compare
-  // equal across methods.)
-  if (places::locality_domain_fake_count() == 0)
+  // The fake-topology override deliberately ignores the split method, so its
+  // places compare equal across methods and carry no per-domain partition.
+  if (std::getenv("CUDASTF_FAKE_LOCALITY_DOMAINS") != nullptr)
+  {
+    fprintf(stderr, "Fake locality-domain topology active: SM partition checks waived.\n");
+  }
+  else
   {
-    // ==== Distinct methods select distinct SM partitions ====
-
-    exec_place p_bf = exec_place::locality_domain(dev, 0, locality_domain_sm_split::backfill);
-    exec_place p_al = exec_place::locality_domain(dev, 0, locality_domain_sm_split::aligned);
-    exec_place p_fi = exec_place::locality_domain(dev, 0, locality_domain_sm_split::fine);
-    EXPECT(p_bf != p_al);
-    EXPECT(p_bf != p_fi);
-    EXPECT(p_al != p_fi);
-
-    // ==== Structural properties of the partitions (adaptive) ====
-    //
-    // Only meaningful when the driver actually splits by domain (not under
-    // the whole-device degrade, where every method spans all SMs).
-    // Under the whole-device degrade the split method is documented to be
-    // ignored: every method must resolve to the SAME cached whole-device
-    // green context (same execution-place identity), not one context per
-    // requested method.
-    if (places::locality_domain_native_raw_count(dev) == 0)
+    // The SM partition of a place, read through the public API: activate the
+    // place and ask the current (green) context for its SM resource.
+    auto partition_of = [](const exec_place& p) {
+      CUdevResource r;
+      auto scope = p.activate();
+      cuda_safe_call(cuCtxGetDevResource(cuda_try<cuCtxGetCurrent>(), &r, CU_DEV_RESOURCE_TYPE_SM));
+      return r.sm;
+    };
+
+    CUdevice device;
+    cuda_safe_call(cuDeviceGet(&device, dev));
+    CUdevResource total;
+    cuda_safe_call(cuDeviceGetDevResource(device, &total, CU_DEV_RESOURCE_TYPE_SM));
+    const unsigned int alignment = total.sm.smCoscheduledAlignment;
+    EXPECT(alignment >= 1);
+
+    if (ndomains == 1)
     {
-      const auto& e_bf = places::locality_domain_ctx_cache::instance().get(dev, 0, locality_domain_sm_split::backfill);
-      const auto& e_al = places::locality_domain_ctx_cache::instance().get(dev, 0, locality_domain_sm_split::aligned);
-      const auto& e_fi = places::locality_domain_ctx_cache::instance().get(dev, 0, locality_domain_sm_split::fine);
-      EXPECT(e_bf.g_ctx == e_al.g_ctx);
-      EXPECT(e_bf.g_ctx == e_fi.g_ctx);
+      // ==== Single domain: the domain is the device ====
+      //
+      // Whether the driver reports one domain or cannot answer the query at
+      // all, every method resolves to the same whole-device place: the
+      // method is documented to be ignored, so the places compare equal and
+      // the partition spans every SM. None of the multi-domain structural
+      // properties below apply.
+      exec_place p_bf = exec_place::locality_domain(dev, 0, locality_domain_sm_split::backfill);
+      EXPECT(p_bf == exec_place::locality_domain(dev, 0, locality_domain_sm_split::aligned));
+      EXPECT(p_bf == exec_place::locality_domain(dev, 0, locality_domain_sm_split::fine));
+      EXPECT(partition_of(p_bf).smCount == total.sm.smCount);
     }
-
-    if (places::locality_domain_native_raw_count(dev) > 0)
+    else
     {
-      CUdevice device;
-      cuda_safe_call(cuDeviceGet(&device, dev));
-      CUdevResource total;
-      cuda_safe_call(cuDeviceGetDevResource(device, &total, CU_DEV_RESOURCE_TYPE_SM));
-      const unsigned int alignment = total.sm.smCoscheduledAlignment;
-      EXPECT(alignment >= 1);
+      // ==== Distinct methods select distinct SM partitions ====
+
+      exec_place p_bf = exec_place::locality_domain(dev, 0, locality_domain_sm_split::backfill);
+      exec_place p_al = exec_place::locality_domain(dev, 0, locality_domain_sm_split::aligned);
+      exec_place p_fi = exec_place::locality_domain(dev, 0, locality_domain_sm_split::fine);
+      EXPECT(p_bf != p_al);
+      EXPECT(p_bf != p_fi);
+      EXPECT(p_al != p_fi);
+
+      // ==== Structural properties of the partitions (adaptive) ====
 
       // backfill sizes every group to an even share of the device (rounded
       // down to the required multiple of 2).
       const unsigned int share             = total.sm.smCount / ndomains;
       const unsigned int expected_backfill = (share < 2) ? 2 : (share - share % 2);
 
-      auto domain_sm_count = [&](unsigned int i, locality_domain_sm_split m) {
-        const auto& entry = places::locality_domain_ctx_cache::instance().get(dev, static_cast<int>(i), m);
-        CUdevResource r;
-        cuda_safe_call(cuGreenCtxGetDevResource(entry.g_ctx, &r, CU_DEV_RESOURCE_TYPE_SM));
-        return r.sm.smCount;
-      };
-
       unsigned int sum_aligned = 0, sum_fine = 0, sum_backfill = 0;
       for (unsigned int i = 0; i < ndomains; i++)
       {
-        const unsigned int c_bf = domain_sm_count(i, locality_domain_sm_split::backfill);
-        const unsigned int c_al = domain_sm_count(i, locality_domain_sm_split::aligned);
-        const unsigned int c_fi = domain_sm_count(i, locality_domain_sm_split::fine);
-
-        EXPECT(c_bf == expected_backfill); // even share of the device
-        EXPECT(c_al % alignment == 0); // complete co-scheduled groups only
-        EXPECT(c_fi % 2 == 0); // finest granularity: groups of 2
-        EXPECT(c_fi >= c_al); // fine recovers at least the aligned SMs
-
-        sum_aligned += c_al;
-        sum_fine += c_fi;
-        sum_backfill += c_bf;
+        const auto bf =
+          partition_of(exec_place::locality_domain(dev, static_cast<int>(i), locality_domain_sm_split::backfill));
+        const auto al =
+          partition_of(exec_place::locality_domain(dev, static_cast<int>(i), locality_domain_sm_split::aligned));
+        const auto fi =
+          partition_of(exec_place::locality_domain(dev, static_cast<int>(i), locality_domain_sm_split::fine));
+
+        // The driver must have honored the request, not silently handed
+        // back an unsplit device: every group echoes the requested domain.
+        for (const auto* r : {&bf, &al, &fi})
+        {
+          EXPECT((r->flags & CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID) != 0);
+          EXPECT(r->localityDomainId == i);
+        }
+        EXPECT((bf.flags & CU_DEV_SM_RESOURCE_GROUP_BACKFILL) != 0);
+
+        EXPECT(bf.smCount == expected_backfill); // even share of the device
+        EXPECT(al.smCount % alignment == 0); // complete co-scheduled groups only
+        EXPECT(fi.smCount % 2 == 0); // finest granularity: groups of 2
+        EXPECT(fi.smCount >= al.smCount); // fine recovers at least the aligned SMs
+
+        sum_aligned += al.smCount;
+        sum_fine += fi.smCount;
+        sum_backfill += bf.smCount;
       }
 
       // The strict methods never exceed the device, and together the
@@ -166,14 +181,6 @@ int main()
       EXPECT(sum_backfill <= total.sm.smCount);
       EXPECT(sum_backfill == expected_backfill * ndomains);
     }
-    else
-    {
-      fprintf(stderr, "Driver reports no locality domains: structural partition checks waived.\n");
-    }
-  }
-  else
-  {
-    fprintf(stderr, "Fake locality-domain topology active: method distinction checks waived.\n");
   }
 #endif // _CCCL_CTK_AT_LEAST(13, 4) && !defined(CUDAX_PLACES_FORCE_LOCALITY_DOMAIN_FALLBACK)
 

docs/cudax/places.rst

@@ -195,11 +195,12 @@ offset the locality benefit the partitioning was meant to capture. Prefer
 whole-device coverage.
 
 The split method only affects the execution side: data places take no
-method, and places built with different methods for the same
-``(device, domain)`` are distinct places sharing the same (equal) affine
-data place. Backends without native locality-domain support (pre-13.4
-toolkits, the whole-device degrade, the fake-topology override) accept and
-ignore the method.
+method, and on a device with several domains, places built with different
+methods for the same ``(device, domain)`` are distinct places sharing the
+same (equal) affine data place. Single-domain devices and backends without
+native locality-domain support (pre-13.4 toolkits, the whole-device degrade,
+the fake-topology override) accept and ignore the method: the domain is the
+device.
 
 All three methods rest on the CUDA 13.4 driver surface that locality-domain
 places already require (splitting by locality domain with