better handling of R values

DESCRIPTION

@@ -66,5 +66,4 @@ Collate:
     'tree.R'
     'utils.R'
     'zzz.R'
-Config/roxygen2/version: 8.0.0
-RoxygenNote: 7.3.3
+Config/roxygen2/version: 8.1.0

NAMESPACE

@@ -115,13 +115,17 @@ importFrom(Rcpp,sourceCpp)
 importFrom(bit64,integer64)
 importFrom(cli,cli_abort)
 importFrom(rlang,"%||%")
-importFrom(safetensors,safe_tensor_buffer)
-importFrom(safetensors,safe_tensor_meta)
-importFrom(tengen,as_array)
-importFrom(tengen,as_dtype)
-importFrom(tengen,as_raw)
-importFrom(tengen,device)
-importFrom(tengen,dtype)
-importFrom(tengen,shape)
+importFrom(safetensors,
+  safe_tensor_buffer,
+  safe_tensor_meta
+)
+importFrom(tengen,
+  as_array,
+  as_dtype,
+  as_raw,
+  device,
+  dtype,
+  shape
+)
 importFrom(utils,hashtab)
 useDynLib(pjrt, .registration = TRUE)

NEWS.md

@@ -1,5 +1,21 @@
 # pjrt (development version)
 
+## Breaking changes
+
+* `dispatcher()`'s cache key no longer carries an `ambiguous` bit, and an
+  `AnvlArray` leaf is now keyed apart from bare R data of the same dtype and
+  shape: the two compile to different programs. The compile callback's
+  `avals` are `list(dtype, shape)`, its `out_avals` likewise, and the wrapped
+  outputs no longer carry `$ambiguous`.
+
+## Features
+
+* A compile callback may return `input_dtypes`, naming the dtype each
+  execute-time input is supplied at. A bare R leaf is then uploaded at that
+  dtype instead of its default (`"f32"` for a double, `"i32"` for an integer,
+  `"pred"` for a logical), which is how a caller whose program consumes an R
+  double as `f64` gets the exact value rather than one rounded through `f32`.
+
 ## Other
 
 * pjrt no longer Suggests anvl and stablehlo for it's tests

R/dispatch.R

@@ -10,7 +10,8 @@
 #'
 #' @details
 #' Each [`dispatch()`] call flattens the inputs and builds a cache key: a
-#' dynamic leaf contributes its dtype, shape and `ambiguous` flag, a static leaf
+#' dynamic leaf contributes its kind (an array, or bare R data), its dtype and
+#' its shape, a static leaf
 #' its value (compared with [identical()]). On a hit the cached executable runs
 #' immediately; on a miss `compile` is called to produce a new cache entry.
 #'
@@ -29,8 +30,8 @@
 #'   inputs contribute their `$data` buffer, bare R literals and arrays are
 #'   uploaded with the same dtype defaults as
 #'   [`pjrt_scalar()`]/[`pjrt_buffer()`], and the outputs are wrapped back into
-#'   `"AnvlArray"`s -- lists of `$data`, `$dtype`, `$shape`, `$device`,
-#'   `$ambiguous` and `$backend` -- and re-nested via `out_tree`, all without
+#'   `"AnvlArray"`s -- lists of `$data`, `$dtype`, `$shape`, `$device` and
+#'   `$backend` -- and re-nested via `out_tree`, all without
 #'   leaving C++.
 #' * any other `backend` calls the compiled R closure `compile` returned, which
 #'   returns the call's finished value. Execution, output wrapping and input
@@ -52,10 +53,10 @@
 #'   * `in_tree`: its `RTree` (see [`build_tree`]),
 #'   * `leaves`: the flat leaf list (see [`flatten`]),
 #'   * `is_static`: a `logical()` mask over `leaves`,
-#'   * `avals`: per leaf, `NULL` if static, else the `list(dtype, shape,
-#'     ambiguous)` the cache key was built from. `dtype` is a canonical dtype
-#'     string (`"f32"`, `"i64"`, ...), `shape` an `integer()`, empty for a
-#'     scalar,
+#'   * `avals`: per leaf, `NULL` if static, else the `list(dtype, shape)` the
+#'     cache key was built from. `dtype` is a canonical dtype string (`"f32"`,
+#'     `"i64"`, ...) -- for a bare R leaf, the dtype it would be uploaded at by
+#'     default -- and `shape` an `integer()`, empty for a scalar,
 #'   * `default_device`: the device this call resolved because no array input
 #'     named one -- the device the cache key was built on, so `compile` must
 #'     compile for it rather than resolve a default of its own. `NULL` when an
@@ -67,13 +68,22 @@
 #'     compiled for,
 #'   * `out_tree`: the `RTree` of the outputs (see [`build_tree`]),
 #'   * `out_avals`: one aval per output leaf of `out_tree`, each a
-#'     `list(dtype = <string>, shape = <integer>, ambiguous = <logical(1)>)`
-#'     (`ambiguous` is optional and defaults to `FALSE`). The outputs are
-#'     wrapped from these.
+#'     `list(dtype = <string>, shape = <integer>)`. The outputs are wrapped
+#'     from these.
 #'   * `const_arrays` (optional): buffers prepended to the inputs,
 #'   * `phantom_specs` (optional): a list of `list(dtype = <string>, shape =
 #'     <integer>)` donation-output buffers to allocate fresh per call.
 #'
+#'   Either kind of result may additionally carry:
+#'   * `input_dtypes` (optional): a `character()` with one entry per dynamic
+#'     leaf, in order, naming the dtype that input is supplied at. A bare R
+#'     leaf is uploaded at that dtype instead of its default (a double at
+#'     `"f32"`, an integer at `"i32"`, a logical at `"pred"`), which is how a
+#'     caller whose program consumes an R double as `f64` gets the exact value
+#'     rather than one rounded through `f32` first. `NA` leaves an input alone;
+#'     an array input is passed through whatever this says, so `NA` is the
+#'     meaningful entry for one.
+#'
 #'   For any other `backend` it must return a named list with:
 #'   * `r_fun`: a function called with the list of the call's dynamic leaves, in
 #'     order and with an array leaf contributing its `$data`, returning the
@@ -122,8 +132,8 @@
 #'   object a backend hands out stays alive for the dispatcher's lifetime.
 #' @param extractor (`function` | `NULL`)\cr
 #'   Reads a non-`"pjrt"` array's metadata via the backend's accessors, called as
-#'   `extractor(leaf)` and returning `list(aval = list(dtype, shape, ambiguous),
-#'   device, backend)` -- `dtype` a tengen `DataType`, `shape` an `integer()`.
+#'   `extractor(leaf)` and returning `list(aval = list(dtype, shape), device,
+#'   backend)` -- `dtype` a tengen `DataType`, `shape` an `integer()`.
 #'   Required for any backend other than `"pjrt"`; ignored for `"pjrt"` (see
 #'   *Backends*).
 #' @return [`dispatcher()`] returns a `Dispatcher`.

src/dispatch.cpp

@@ -231,7 +231,6 @@ SEXP impl_dispatch_run(SEXP dispatcher, Rcpp::List args) {
     kl.kind = KeyLeaf::kRData;
     kl.aval.dtype = rd->dtype;
     kl.aval.shape = std::move(rd->shape);
-    kl.aval.ambiguous = true;  // bare R data is dtype-ambiguous (to_avals)
     key.leaves.push_back(std::move(kl));
     exec_inputs.push_back({leaf, &key.leaves.back().aval, true});
   }
@@ -278,8 +277,7 @@ SEXP impl_dispatch_run(SEXP dispatcher, Rcpp::List args) {
       // the callback always sees a string, whichever backend the leaf is from.
       avals[i] = Rcpp::List::create(
           Rcpp::Named("dtype") = anvl_dtype_name(kl.aval.dtype),
-          Rcpp::Named("shape") = shp,
-          Rcpp::Named("ambiguous") = kl.aval.ambiguous);
+          Rcpp::Named("shape") = shp);
     }
     Rcpp::List info = Rcpp::List::create(
         Rcpp::Named("args") = args,
@@ -297,6 +295,9 @@ SEXP impl_dispatch_run(SEXP dispatcher, Rcpp::List args) {
     // The engine validates the result and builds its entry material.
     CacheEntry e;
     engine.build_entry(res, e);
+    // Engine-agnostic: the dtype each input is supplied at, which the callback
+    // declares because only the compiled program knows what it takes.
+    read_input_dtypes(res, exec_inputs.size(), e);
 
     // Root every SEXP the inserted key holds, so it outlives this call; the
     // entry drops them when it is evicted. The key's device token needs no
@@ -308,6 +309,22 @@ SEXP impl_dispatch_run(SEXP dispatcher, Rcpp::List args) {
     entry = d.cache().get(key);
   }
 
-  // 5. The engine runs the call and returns the finished value.
+  // 5. Stamp each input with the dtype the entry was compiled to take it at,
+  // then let the engine run the call and return the finished value.
+  if (!entry->input_dtypes.empty()) {
+    // Sizes agree by construction: the key fixes which leaves are static, so
+    // every call filed under it supplies the same number of inputs as the one
+    // the entry was validated against.
+    if (entry->input_dtypes.size() != exec_inputs.size()) {
+      Rcpp::stop(
+          "internal error: cache entry declares %d input dtypes but "
+          "this call supplies %d inputs",
+          static_cast<int>(entry->input_dtypes.size()),
+          static_cast<int>(exec_inputs.size()));
+    }
+    for (std::size_t k = 0; k < exec_inputs.size(); ++k) {
+      exec_inputs[k].dtype = entry->input_dtypes[k];
+    }
+  }
   return engine.run(*entry, exec_inputs);
 }