diff --git a/DESCRIPTION b/DESCRIPTION index 85ef5796..0cf9f6d7 100644 --- a/DESCRIPTION +++ b/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 diff --git a/NAMESPACE b/NAMESPACE index 452b9e98..e1113d97 100644 --- a/NAMESPACE +++ b/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) diff --git a/NEWS.md b/NEWS.md index 739c6e75..a0d9bfda 100644 --- a/NEWS.md +++ b/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 diff --git a/R/dispatch.R b/R/dispatch.R index e6c7222e..f6465d06 100644 --- a/R/dispatch.R +++ b/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 = , shape = , ambiguous = )` -#' (`ambiguous` is optional and defaults to `FALSE`). The outputs are -#' wrapped from these. +#' `list(dtype = , shape = )`. The outputs are wrapped +#' from these. #' * `const_arrays` (optional): buffers prepended to the inputs, #' * `phantom_specs` (optional): a list of `list(dtype = , shape = #' )` 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`. diff --git a/man/dispatcher.Rd b/man/dispatcher.Rd index bfb9de60..a0301a04 100644 --- a/man/dispatcher.Rd +++ b/man/dispatcher.Rd @@ -27,9 +27,10 @@ classify the inputs a second time: \item \code{in_tree}: its \code{RTree} (see \code{\link{build_tree}}), \item \code{leaves}: the flat leaf list (see \code{\link{flatten}}), \item \code{is_static}: a \code{logical()} mask over \code{leaves}, -\item \code{avals}: per leaf, \code{NULL} if static, else the \code{list(dtype, shape, ambiguous)} the cache key was built from. \code{dtype} is a canonical dtype -string (\code{"f32"}, \code{"i64"}, ...), \code{shape} an \code{integer()}, empty for a -scalar, +\item \code{avals}: per leaf, \code{NULL} if static, else the \code{list(dtype, shape)} the +cache key was built from. \code{dtype} is a canonical dtype string (\code{"f32"}, +\code{"i64"}, ...) -- for a bare R leaf, the dtype it would be uploaded at by +default -- and \code{shape} an \code{integer()}, empty for a scalar, \item \code{default_device}: the device this call resolved because no array input named one -- the device the cache key was built on, so \code{compile} must compile for it rather than resolve a default of its own. \code{NULL} when an @@ -43,13 +44,24 @@ For \code{backend = "pjrt"} it must return a named list with: compiled for, \item \code{out_tree}: the \code{RTree} of the outputs (see \code{\link{build_tree}}), \item \code{out_avals}: one aval per output leaf of \code{out_tree}, each a -\verb{list(dtype = , shape = , ambiguous = )} -(\code{ambiguous} is optional and defaults to \code{FALSE}). The outputs are -wrapped from these. +\verb{list(dtype = , shape = )}. The outputs are wrapped +from these. \item \code{const_arrays} (optional): buffers prepended to the inputs, \item \code{phantom_specs} (optional): a list of \verb{list(dtype = , shape = )} donation-output buffers to allocate fresh per call. } +Either kind of result may additionally carry: +\itemize{ +\item \code{input_dtypes} (optional): a \code{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 +\code{"f32"}, an integer at \code{"i32"}, a logical at \code{"pred"}), which is how a +caller whose program consumes an R double as \code{f64} gets the exact value +rather than one rounded through \code{f32} first. \code{NA} leaves an input alone; +an array input is passed through whatever this says, so \code{NA} is the +meaningful entry for one. +} + For any other \code{backend} it must return a named list with: \itemize{ \item \code{r_fun}: a function called with the list of the call's dynamic leaves, in @@ -105,7 +117,7 @@ object a backend hands out stays alive for the dispatcher's lifetime.} \item{extractor}{(\code{function} | \code{NULL})\cr Reads a non-\code{"pjrt"} array's metadata via the backend's accessors, called as -\code{extractor(leaf)} and returning \code{list(aval = list(dtype, shape, ambiguous), device, backend)} -- \code{dtype} a tengen \code{DataType}, \code{shape} an \code{integer()}. +\code{extractor(leaf)} and returning \code{list(aval = list(dtype, shape), device, backend)} -- \code{dtype} a tengen \code{DataType}, \code{shape} an \code{integer()}. Required for any backend other than \code{"pjrt"}; ignored for \code{"pjrt"} (see \emph{Backends}).} } @@ -123,7 +135,8 @@ compile only on a cache miss. } \details{ Each \code{\link[=dispatch]{dispatch()}} call flattens the inputs and builds a cache key: a -dynamic leaf contributes its dtype, shape and \code{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 \code{\link[=identical]{identical()}}). On a hit the cached executable runs immediately; on a miss \code{compile} is called to produce a new cache entry. @@ -143,8 +156,8 @@ sits behind it: inputs contribute their \verb{$data} buffer, bare R literals and arrays are uploaded with the same dtype defaults as \code{\link[=pjrt_scalar]{pjrt_scalar()}}/\code{\link[=pjrt_buffer]{pjrt_buffer()}}, and the outputs are wrapped back into -\code{"AnvlArray"}s -- lists of \verb{$data}, \verb{$dtype}, \verb{$shape}, \verb{$device}, -\verb{$ambiguous} and \verb{$backend} -- and re-nested via \code{out_tree}, all without +\code{"AnvlArray"}s -- lists of \verb{$data}, \verb{$dtype}, \verb{$shape}, \verb{$device} and +\verb{$backend} -- and re-nested via \code{out_tree}, all without leaving C++. \item any other \code{backend} calls the compiled R closure \code{compile} returned, which returns the call's finished value. Execution, output wrapping and input diff --git a/src/dispatch.cpp b/src/dispatch.cpp index 22e12f62..7977a67b 100644 --- a/src/dispatch.cpp +++ b/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(entry->input_dtypes.size()), + static_cast(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); } diff --git a/src/dispatch_engine.cpp b/src/dispatch_engine.cpp index a725dafa..28eba114 100644 --- a/src/dispatch_engine.cpp +++ b/src/dispatch_engine.cpp @@ -139,6 +139,38 @@ std::string leaf_subject(const RTree& in_tree, std::size_t leaf_index) { return "input `" + tree_path(in_tree, static_cast(leaf_index) + 1) + "`"; } +// ---- The compile callback's declared input dtypes -------------------------- + +void read_input_dtypes(const Rcpp::List& res, std::size_t n_inputs, + CacheEntry& e) { + if (!res.containsElementNamed("input_dtypes")) return; + SEXP v = res["input_dtypes"]; + if (v == R_NilValue) return; + if (TYPEOF(v) != STRSXP) { + Rcpp::stop("compile result: `input_dtypes` must be a character vector"); + } + const R_xlen_t n = XLENGTH(v); + if (static_cast(n) != n_inputs) { + Rcpp::stop( + "compile result: `input_dtypes` has %d entries but the call supplies " + "%d inputs", + static_cast(n), static_cast(n_inputs)); + } + e.input_dtypes.assign(static_cast(n), AnvlDtype::kInvalid); + for (R_xlen_t k = 0; k < n; ++k) { + SEXP el = STRING_ELT(v, k); + if (el == NA_STRING) continue; // this input is passed through as it is + const AnvlDtype d = anvl_dtype_from_name(CHAR(el)); + if (d == AnvlDtype::kInvalid) { + Rcpp::stop( + "compile result: `input_dtypes[[%d]]` is not a dtype anvl can " + "represent: \"%s\"", + static_cast(k + 1), CHAR(el)); + } + e.input_dtypes[static_cast(k)] = d; + } +} + // ---- Engine: device canonicalization and the generic Aval read ------------- SEXP Engine::canonical_device(SEXP device) { @@ -168,10 +200,10 @@ static void check_dtype_representable(const Aval& a, const RTree& in_tree, } } -// Build an Aval from a backend extractor's outputs: a tengen DataType object, -// an integer shape, and the ambiguous bit. -static Aval aval_from_tengen(SEXP dtype, SEXP shape, bool ambiguous, - const RTree& in_tree, std::size_t leaf_index) { +// Build an Aval from a backend extractor's outputs: a tengen DataType object +// and an integer shape. +static Aval aval_from_tengen(SEXP dtype, SEXP shape, const RTree& in_tree, + std::size_t leaf_index) { if (dtype == R_NilValue || TYPEOF(shape) != INTSXP) { Rcpp::stop( "invalid %s: the backend extractor must return an aval with a dtype " @@ -180,7 +212,6 @@ static Aval aval_from_tengen(SEXP dtype, SEXP shape, bool ambiguous, } Aval a; a.dtype = anvl_dtype_from_tengen(dtype); - a.ambiguous = ambiguous; const R_xlen_t nd = XLENGTH(shape); a.shape.reserve(nd); for (R_xlen_t j = 0; j < nd; ++j) a.shape.push_back(INTEGER(shape)[j]); @@ -220,7 +251,7 @@ class ClosureEngine : public Engine { extractor_(require_extractor(extractor)) {} // Read metadata through the backend's accessors: extractor(leaf) returns - // list(aval = list(dtype, shape, ambiguous), device, backend). `$data` is the + // list(aval = list(dtype, shape), device, backend). `$data` is the // one field read directly (contract-guaranteed). The Aval is built only for a // leaf of this dispatcher's backend; a foreign leaf carries its true tag back // for the core to reject, and its metadata is neither required nor validated. @@ -242,9 +273,7 @@ class ClosureEngine : public Engine { Rcpp::List av = meta["aval"]; SEXP dtype = av.containsElementNamed("dtype") ? av["dtype"] : R_NilValue; SEXP shape = av.containsElementNamed("shape") ? av["shape"] : R_NilValue; - const bool amb = - av.containsElementNamed("ambiguous") && field_true(av["ambiguous"]); - al.aval = aval_from_tengen(dtype, shape, amb, in_tree, leaf_index); + al.aval = aval_from_tengen(dtype, shape, in_tree, leaf_index); } return al; } @@ -283,6 +312,15 @@ class ClosureEngine : public Engine { // ---- PjrtEngine ------------------------------------------------------------- +// The dtype name to upload a bare R leaf at: the one the entry declared, or the +// leaf's own default when it declared none. "bool" is spelled "pred" at the +// buffer-facing layer (see anvl_dtype_name). +static const char* upload_dtype_name(AnvlDtype d, const char* fallback) { + if (d == AnvlDtype::kInvalid) return fallback; + if (d == AnvlDtype::kBool) return "pred"; + return anvl_dtype_name(d); +} + // One output-donation phantom buffer to allocate per call (CPU memory mgmt). struct PhantomSpec { PJRT_Buffer_Type dtype = PJRT_Buffer_Type_INVALID; @@ -326,11 +364,10 @@ class PjrtEngine : public Engine { // Reads the native way: dtype/shape/device all come off the PJRTBuffer in // `$data`, which caches them natively and so cannot drift from the array's - // own fields (this engine never consults them). `$ambiguous` is the only - // R-list field the Aval needs -- the buffer carries no such anvl type-system - // bit -- and `$backend` the only other, for the reject path. The buffer is - // interpreted only once the leaf's tag matches, so a foreign leaf carries its - // tag back for the core to reject rather than failing the buffer check here. + // own fields (this engine never consults them). `$backend` is the only + // R-list field it needs, for the reject path. The buffer is interpreted only + // once the leaf's tag matches, so a foreign leaf carries its tag back for + // the core to reject rather than failing the buffer check here. std::optional read_array(SEXP leaf, const RTree& in_tree, std::size_t leaf_index) override { if (!is_anvl_array(leaf)) return std::nullopt; @@ -346,7 +383,6 @@ class PjrtEngine : public Engine { Rcpp::XPtr buf(al.data); al.aval.dtype = anvl_dtype_from_pjrt(buf->element_type()); al.aval.shape = buf->dimensions(); - al.aval.ambiguous = field_true(anvl_field(leaf, "ambiguous")); check_dtype_representable(al.aval, in_tree, leaf_index); // Device from the buffer, not $device: interned by PJRT_Device* (see // canonical_device) so the token still matches a literal-only call's @@ -492,18 +528,25 @@ class PjrtEngine : public Engine { inputs[pos++] = in.value; continue; } + // The dtype the program was compiled to take this input at. The + // callback declares it because only the trace knows what the value is + // used for -- an R double that meets an f64 array has to arrive as f64, + // not rounded through the f32 default first. switch (TYPEOF(in.value)) { case REALSXP: inputs[pos++] = impl_client_buffer_from_double( - pe->client, pe->device, in.value, in.aval->shape, "f32"); + pe->client, pe->device, in.value, in.aval->shape, + upload_dtype_name(in.dtype, "f32")); break; case INTSXP: inputs[pos++] = impl_client_buffer_from_integer( - pe->client, pe->device, in.value, in.aval->shape, "i32"); + pe->client, pe->device, in.value, in.aval->shape, + upload_dtype_name(in.dtype, "i32")); break; default: inputs[pos++] = impl_client_buffer_from_logical( - pe->client, pe->device, in.value, in.aval->shape, "pred"); + pe->client, pe->device, in.value, in.aval->shape, + upload_dtype_name(in.dtype, "pred")); break; } } @@ -545,13 +588,13 @@ class PjrtEngine : public Engine { private: // One template AnvlArray per output, built on the compile (cold) path from // the avals the callback declared: a named list (data = NULL, dtype, shape, - // device, ambiguous, backend) of class "AnvlArray" -- the wrapper layout an - // pjrt leaf carries, which PjrtEngine::read_array reads back as an input. The - // hot path only shallow-copies a template and drops the output buffer into - // its `$data` slot. + // device, backend) of class "AnvlArray" -- the wrapper layout an pjrt leaf + // carries, which PjrtEngine::read_array reads back as an input. The hot path + // only shallow-copies a template and drops the output buffer into its + // `$data` slot. // - // `out_avals[[i]]` is list(dtype = , shape = , ambiguous = - // , the last optional). See ?dispatcher. + // `out_avals[[i]]` is list(dtype = , shape = ). See + // ?dispatcher. Rcpp::List build_templates(SEXP out_avals, SEXP device) const { const R_xlen_t n_out = XLENGTH(out_avals); Rcpp::Environment tengen = Rcpp::Environment::namespace_env("tengen"); @@ -577,13 +620,11 @@ class PjrtEngine : public Engine { // C-API spelling "pred" and the MLIR spelling "i1" are accepted aliases). if (dt == "pred" || dt == "i1") dt = "bool"; Rcpp::IntegerVector shape = Rcpp::as(aval["shape"]); - const bool amb = aval.containsElementNamed("ambiguous") && - Rf_asLogical(aval["ambiguous"]) == TRUE; Rcpp::List tmpl = Rcpp::List::create( Rcpp::Named("data") = R_NilValue, Rcpp::Named("dtype") = as_dtype(Rcpp::CharacterVector::create(dt)), Rcpp::Named("shape") = shape, Rcpp::Named("device") = device, - Rcpp::Named("ambiguous") = amb, Rcpp::Named("backend") = backend); + Rcpp::Named("backend") = backend); tmpl.attr("class") = cls; templates[i] = tmpl; } diff --git a/src/dispatch_engine.h b/src/dispatch_engine.h index 77530e40..77df47f1 100644 --- a/src/dispatch_engine.h +++ b/src/dispatch_engine.h @@ -81,6 +81,11 @@ struct ExecInput { SEXP value = R_NilValue; // an array leaf's `$data`, or the bare R leaf const Aval* aval = nullptr; // an upload needs its shape bool upload = false; // bare R data: upload it. Else: ready to use + // The dtype to upload bare R data at. kInvalid means the leaf's own default + // (double -> f32, integer -> i32, logical -> pred); the entry's + // `input_dtypes` overrides it, because only the compiled program knows what + // dtype it was compiled to take. Unused when `upload` is false. + AnvlDtype dtype = AnvlDtype::kInvalid; }; // An engine's per-entry material -- what the compile callback produced, in the @@ -105,6 +110,12 @@ struct CacheEntry { // and are rooted by it, not here. std::vector keep; std::unique_ptr data; + // The dtype each execute-time input is supplied at, in program order, as the + // compile callback declared it (`input_dtypes`). kInvalid at a position + // leaves that input alone -- an array input is passed through, and bare R + // data falls back to its own default dtype. Empty when the callback declared + // nothing, which is the same as all-kInvalid. + std::vector input_dtypes; // Root `x` for this entry's lifetime. R_NilValue is a no-op. void keep_alive(SEXP x) { @@ -166,6 +177,14 @@ class Engine { std::vector canonical_devices_; }; +// Read the compile callback's `input_dtypes` into the entry, validating it: a +// character vector, one element per execute-time input, each a canonical dtype +// name or NA. Absent (or NULL) leaves `e.input_dtypes` empty. `n_inputs` is +// what the call actually has to supply, so a callback that declares a different +// number is a malformed result rather than a silent mismatch at execute time. +void read_input_dtypes(const Rcpp::List& res, std::size_t n_inputs, + CacheEntry& e); + // `engine_name` is the R-facing selector: "pjrt" or "closure"; throws on any // other value. `backend` is the tag the dispatcher's arrays carry (the engine // needs it to stamp the arrays it wraps and to recognize its own leaves). diff --git a/src/dispatch_key.h b/src/dispatch_key.h index 5dd0e53f..16ad006e 100644 --- a/src/dispatch_key.h +++ b/src/dispatch_key.h @@ -110,19 +110,10 @@ inline const char *anvl_dtype_name(AnvlDtype d) { return "invalid"; } -// Translate a tengen DataType object to an AnvlDtype. It is a length-1 -// character vector classed "DataType" whose string is the canonical dtype -// name. tengen names more dtypes than the dispatcher supports (f16, bf16, -// f8*, complex, sub-byte ints); those yield kInvalid and the caller rejects -// them rather than keying approximately. -inline AnvlDtype anvl_dtype_from_tengen(SEXP dtype) { - if (TYPEOF(dtype) != STRSXP || XLENGTH(dtype) != 1) { - return AnvlDtype::kInvalid; - } - if (!Rf_inherits(dtype, "DataType")) { - return AnvlDtype::kInvalid; - } - const char *name = CHAR(STRING_ELT(dtype, 0)); +// Translate a canonical dtype name to an AnvlDtype. tengen names more dtypes +// than the dispatcher supports (f16, bf16, f8*, complex, sub-byte ints); those +// yield kInvalid and the caller rejects them rather than keying approximately. +inline AnvlDtype anvl_dtype_from_name(const char *name) { if (!std::strcmp(name, "bool")) return AnvlDtype::kBool; if (!std::strcmp(name, "i8")) return AnvlDtype::kI8; if (!std::strcmp(name, "i16")) return AnvlDtype::kI16; @@ -137,19 +128,28 @@ inline AnvlDtype anvl_dtype_from_tengen(SEXP dtype) { return AnvlDtype::kInvalid; } -// Per-leaf abstract value -- mirrors anvl's nv_aval(dtype, shape, ambiguous). -// dtype/shape are read off the leaf; `ambiguous` is an anvl type-system bit -// supplied per leaf (pjrt folds it into the key but never interprets it). The -// device is not part of it: it is a single per-call value on the CacheKey. +// The same, for a tengen DataType object: a length-1 character vector classed +// "DataType" whose string is the canonical dtype name. +inline AnvlDtype anvl_dtype_from_tengen(SEXP dtype) { + if (TYPEOF(dtype) != STRSXP || XLENGTH(dtype) != 1) { + return AnvlDtype::kInvalid; + } + if (!Rf_inherits(dtype, "DataType")) { + return AnvlDtype::kInvalid; + } + return anvl_dtype_from_name(CHAR(STRING_ELT(dtype, 0))); +} + +// Per-leaf abstract value -- mirrors anvl's nv_aval(dtype, shape). Both are +// read off the leaf. The device is not part of it: it is a single per-call +// value on the CacheKey. struct Aval { AnvlDtype dtype = AnvlDtype::kInvalid; std::vector shape; - bool ambiguous = false; }; inline std::uint64_t aval_hash(const Aval &a) { std::uint64_t h = static_cast(a.dtype); - h = hash_combine(h, a.ambiguous ? 1u : 0u); for (int64_t d : a.shape) { h = hash_combine(h, static_cast(d)); } @@ -157,7 +157,7 @@ inline std::uint64_t aval_hash(const Aval &a) { } inline bool aval_eq(const Aval &a, const Aval &b) { - return a.dtype == b.dtype && a.ambiguous == b.ambiguous && a.shape == b.shape; + return a.dtype == b.dtype && a.shape == b.shape; } // identical(), tightened for use as a cache key. @@ -205,15 +205,15 @@ struct KeyLeaf { SEXP value = R_NilValue; // kStatic: the leaf }; -// How a leaf contributes to the key: by its value, or by its Aval. +// How a leaf contributes to the key: by its value, or by its kind and Aval. // -// kArray and kRData are deliberately not distinguished. They differ only in -// where execution finds the input -- the leaf's own $data, or a fresh upload of -// the leaf -- and that is settled per call, from the call's own leaves, never -// from the cache entry. The program compiled for an Aval is the same either -// way, so keying them apart would compile it twice: `f(x, y)` and `f(x, 1)` -// with matching avals should share one executable. `kind` survives only to -// steer input assembly. +// kArray and kRData are distinguished, because the caller compiles them into +// different programs: an array leaf is an input of a dtype it already has, +// while bare R data has no dtype of its own until the program says what it is +// used as, and can therefore be uploaded as something other than its default +// (anvl's RData values, which let `x_f64 / sqrt(2)` see the exact double +// rather than one rounded through f32). `f(x, y)` and `f(x, 1)` are two +// programs, and two entries. // // CacheKeyHash and CacheKeyEq must agree on this, or two keys the map calls // equal would hash into different buckets. @@ -288,11 +288,10 @@ struct CacheKeyHash { for (const KeyLeaf &leaf : k.leaves) { // Folded before the per-leaf material, so a value-keyed leaf's hash // stream can never coincide with an Aval-keyed one's: the domain - // separator. Note it is `keyed_by_value`, not `kind` -- a kArray and a - // kRData leaf of the same Aval must hash alike, because CacheKeyEq calls - // them equal. + // separator. The kind goes in whole, so a kArray and a kRData leaf of + // the same Aval land in different buckets, as CacheKeyEq requires. const bool by_value = keyed_by_value(leaf.kind); - h = hash_combine(h, by_value ? 1u : 0u); + h = hash_combine(h, static_cast(leaf.kind)); if (!by_value) { h = hash_combine(h, aval_hash(leaf.aval)); continue; @@ -327,12 +326,12 @@ struct CacheKeyEq { for (std::size_t k = 0; k < a.leaves.size(); ++k) { const KeyLeaf &x = a.leaves[k]; const KeyLeaf &y = b.leaves[k]; - // A kArray and a kRData leaf of the same Aval are the same key: the two - // compile to one program (see keyed_by_value). Only value-keyed against - // Aval-keyed is a difference -- and tree_eq has already ruled that out, - // since static-ness follows the argument names it compares. + // A kArray and a kRData leaf of the same Aval are different keys: they + // compile to different programs (see keyed_by_value). Static-ness is + // already ruled out by tree_eq, which compares the argument names it + // follows, but the array/rdata split is not. + if (x.kind != y.kind) return false; const bool by_value = keyed_by_value(x.kind); - if (by_value != keyed_by_value(y.kind)) return false; if (by_value) { if (!r_identical(x.value, y.value)) return false; } else if (!aval_eq(x.aval, y.aval)) { diff --git a/src/test-dispatch.cpp b/src/test-dispatch.cpp index 9bcffa93..0d51a17c 100644 --- a/src/test-dispatch.cpp +++ b/src/test-dispatch.cpp @@ -45,11 +45,10 @@ RTree flat_tree(std::size_t n) { return t; } -Aval mk_aval(AnvlDtype dtype, std::vector shape, bool ambiguous) { +Aval mk_aval(AnvlDtype dtype, std::vector shape) { Aval a; a.dtype = dtype; a.shape = std::move(shape); - a.ambiguous = ambiguous; return a; } @@ -157,7 +156,7 @@ context("AnvlDtype") { } context("CacheKey: aval-keyed leaves") { - const Aval f32_2x3 = mk_aval(AnvlDtype::kF32, {2, 3}, false); + const Aval f32_2x3 = mk_aval(AnvlDtype::kF32, {2, 3}); test_that("equal signatures compare equal and hash alike") { CacheKey a = key_of({array_leaf(f32_2x3), array_leaf(f32_2x3)}); @@ -166,42 +165,34 @@ context("CacheKey: aval-keyed leaves") { expect_true(hash_of(a) == hash_of(b)); } - test_that("dtype, shape, ambiguity and arity each split the key") { + test_that("dtype, shape and arity each split the key") { CacheKey base = key_of({array_leaf(f32_2x3)}); - CacheKey dtype = - key_of({array_leaf(mk_aval(AnvlDtype::kI32, {2, 3}, false))}); + CacheKey dtype = key_of({array_leaf(mk_aval(AnvlDtype::kI32, {2, 3}))}); expect_false(eq(base, dtype)); expect_false(hash_of(base) == hash_of(dtype)); - CacheKey shape = - key_of({array_leaf(mk_aval(AnvlDtype::kF32, {3, 2}, false))}); + CacheKey shape = key_of({array_leaf(mk_aval(AnvlDtype::kF32, {3, 2}))}); expect_false(eq(base, shape)); expect_false(hash_of(base) == hash_of(shape)); - CacheKey rank = key_of({array_leaf(mk_aval(AnvlDtype::kF32, {}, false))}); + CacheKey rank = key_of({array_leaf(mk_aval(AnvlDtype::kF32, {}))}); expect_false(eq(base, rank)); expect_false(hash_of(base) == hash_of(rank)); - CacheKey ambig = - key_of({array_leaf(mk_aval(AnvlDtype::kF32, {2, 3}, true))}); - expect_false(eq(base, ambig)); - expect_false(hash_of(base) == hash_of(ambig)); - CacheKey arity = key_of({array_leaf(f32_2x3), array_leaf(f32_2x3)}); expect_false(eq(base, arity)); expect_false(hash_of(base) == hash_of(arity)); } - test_that("a kArray and a kRData leaf of one aval are one key") { - // They compile to the same program; only where execution finds the input - // differs, and that is decided per call. Keying them apart would compile - // `f(x, y)` and `f(x, 1)` twice. + test_that("a kArray and a kRData leaf of one aval are different keys") { + // They compile to different programs: bare R data has no dtype of its own + // until the program says what it is used as, so `f(x, 1)` may consume it + // at a dtype `f(x, y)` never asks for. CacheKey arr = key_of({array_leaf(f32_2x3)}); CacheKey lit = key_of({rdata_leaf(f32_2x3)}); - expect_true(eq(arr, lit)); - expect_true(hash_of(arr) == - hash_of(lit)); // or the map never compares them + expect_false(eq(arr, lit)); + expect_false(hash_of(arr) == hash_of(lit)); } test_that("an aval-keyed leaf never equals a value-keyed one") { @@ -213,7 +204,7 @@ context("CacheKey: aval-keyed leaves") { } context("CacheKey: device and tree") { - const Aval f32 = mk_aval(AnvlDtype::kF32, {2}, false); + const Aval f32 = mk_aval(AnvlDtype::kF32, {2}); test_that("the device token splits the key and is folded into the hash") { CacheKey a = key_of({array_leaf(f32)}); diff --git a/tests/testthat/test-dispatch.R b/tests/testthat/test-dispatch.R index 94490a31..91efc2ea 100644 --- a/tests/testthat/test-dispatch.R +++ b/tests/testthat/test-dispatch.R @@ -8,11 +8,11 @@ test_pjrt_device <- function() pjrt_device("cpu:0") test_device <- function(id = "cpu") structure(list(device = id), class = "QuickrDevice") test_quickr_device <- function() test_device("cpu") -# One output aval, as the compile callback declares it: the dtype/shape/ -# ambiguous pjrt stamps on that output's wrapper. Same shape as the input avals -# the callback receives in `info$avals`. -oav <- function(dtype = "f32", shape = 2L, ambiguous = FALSE) { - list(dtype = dtype, shape = as.integer(shape), ambiguous = ambiguous) +# One output aval, as the compile callback declares it: the dtype and shape +# pjrt stamps on that output's wrapper. Same shape as the input avals the +# callback receives in `info$avals`. +oav <- function(dtype = "f32", shape = 2L) { + list(dtype = dtype, shape = as.integer(shape)) } # The pjrt engine's full compile-callback contract for a single-output @@ -40,7 +40,7 @@ pjrt_entry <- function( # A pjrt array leaf, as anvl builds them: an "AnvlArray" whose $data is a buffer. parr <- function(buf) { structure( - list(data = buf, ambiguous = FALSE, device = tengen::device(buf), backend = "pjrt"), + list(data = buf, device = tengen::device(buf), backend = "pjrt"), class = "AnvlArray" ) } @@ -52,7 +52,6 @@ qarr <- function(v, dtype = "f64", device = test_quickr_device(), backend = "qui data = v, dtype = if (is.character(dtype)) tengen::as_dtype(dtype) else dtype, shape = as.integer(length(v)), - ambiguous = FALSE, device = device, backend = backend ), @@ -70,7 +69,7 @@ out <- function(res) as.numeric(tengen::as_array(await(res$data))) # backend whose accessors happen to be `$` reads. test_extractor <- function(leaf) { list( - aval = list(dtype = leaf$dtype, shape = leaf$shape, ambiguous = leaf$ambiguous), + aval = list(dtype = leaf$dtype, shape = leaf$shape), device = leaf$device, backend = leaf$backend ) @@ -293,21 +292,6 @@ test_that("the pjrt engine wraps its outputs and caches one entry per signature" } }) -test_that("the wrapped output's ambiguity is the callback's claim", { - skip_if_not(plugins_downloaded()) - wrapped <- function(amb) { - d <- dispatcher( - 10L, - function(info) pjrt_entry(binop_exec(), out_avals = list(oav(ambiguous = amb))), - default_device = test_pjrt_device - ) - x <- parr(pjrt_buffer(c(1, 2), dtype = "f32")) - dispatch(d, list(x = x, y = x)) - } - expect_false(wrapped(FALSE)$ambiguous) - expect_true(wrapped(TRUE)$ambiguous) -}) - test_that("a static argument compiles one entry per distinct value", { skip_if_not(plugins_downloaded()) d <- dispatcher( @@ -355,9 +339,15 @@ test_that("bare R data is uploaded at its default dtype, column-major", { expect_equal(out(dispatch(d, list(x = 50))), 50) expect_equal(dispatcher_size(d), 1L) + # An array leaf is not the same key material as bare R data, even at the same + # aval: the R value has no dtype of its own until the program says what it is + # used as, so the two compile to different programs and take separate entries. + expect_equal(out(dispatch(d, list(x = parr(pjrt_scalar(5, dtype = "f32"))))), 5) + expect_equal(dispatcher_size(d), 2L) + # An integer literal defaults to i32, a different aval and so a new entry. invisible(dispatch(d, list(x = 3L))) - expect_equal(dispatcher_size(d), 2L) + expect_equal(dispatcher_size(d), 3L) # An R array uploads column-major, like pjrt_buffer(). m <- matrix(c(1, 2, 3, 4), nrow = 2) @@ -859,14 +849,14 @@ test_that("an input pjrt cannot classify is rejected, naming the offending argum test_that("a closure backend can compute metadata via accessors, storing no fields", { # anvl's AnvlBackend contract guarantees only $data on a leaf; dtype/shape/ - # device/ambiguous/backend may be computed by the backend's accessors rather + # device/backend may be computed by the backend's accessors rather # than stored as fields. The dispatcher must read them through the extractor, # never by reaching for fields -- this array carries $data and nothing else. n_miss <- 0L dev <- test_device("cpu") extractor <- function(leaf) { list( - aval = list(dtype = tengen::as_dtype("f64"), shape = length(leaf$data), ambiguous = FALSE), + aval = list(dtype = tengen::as_dtype("f64"), shape = length(leaf$data)), device = dev, backend = "quickr" ) @@ -921,18 +911,17 @@ test_that("out_avals and out_tree are the callback's claim, and are honoured", { y <- parr(pjrt_buffer(c(3, 4), dtype = "f32")) # Every wrapped field comes from the declared aval, not from the buffer: the - # first output is stamped ambiguous although nothing about the buffer is. + # first output is stamped f64 although the buffer it holds is f32. d <- mk( build_tree(list(sum = 0, rest = list(prod = 0))), - list(oav(ambiguous = TRUE), oav()) + list(oav("f64"), oav()) ) res <- impl_dispatch_run(d, list(x, y)) expect_equal(out(res$sum), c(4, 6)) expect_equal(out(res$rest$prod), c(3, 8)) - expect_identical(res$sum$ambiguous, TRUE) - expect_identical(res$rest$prod$ambiguous, FALSE) + expect_identical(res$sum$dtype, tengen::as_dtype("f64")) + expect_identical(res$rest$prod$dtype, tengen::as_dtype("f32")) expect_identical(res$sum$shape, 2L) - expect_identical(res$sum$dtype, tengen::as_dtype("f32")) # An out_tree whose leaf count disagrees with the executable's actual output # count is the one half of the callback's claim pjrt can still settle, and it @@ -1026,3 +1015,67 @@ test_that("an AnvlArray that is not a list is rejected, naming the argument", { bad <- structure(c(data = 1), class = "AnvlArray") expect_error(impl_dispatch_run(d, list(x = bad)), "invalid input `x`") }) + +# --------------------------------------------------------------------------- +# `input_dtypes`: the dtype an execute-time input is supplied at. +# --------------------------------------------------------------------------- + +test_that("`input_dtypes` decides the dtype a bare R leaf is uploaded at", { + skip_if_not(plugins_downloaded()) + # An f64 program: without `input_dtypes` the bare R double would arrive as + # f32 and the executable would refuse it, and -- the point of the mechanism + # -- the value would already have been rounded to f32 on the way in. + src <- 'func.func @main(%x: tensor) -> tensor { + "func.return"(%x): (tensor) -> () + }' + exec <- pjrt_compile(pjrt_program(src = src)) + entry <- function(dtypes) { + function(info) { + pjrt_entry( + exec, + out_tree = build_tree(0), + out_avals = list(oav("f64", integer())), + input_dtypes = dtypes + ) + } + } + d <- dispatcher(10L, entry("f64"), default_device = test_pjrt_device) + res <- dispatch(d, list(x = sqrt(2))) + # Exact to the last bit: the R double was uploaded as f64, not widened from f32. + expect_identical(as.numeric(tengen::as_array(await(res$data))), sqrt(2)) + + # The same call keys the same entry whatever the value, so a second value + # is served by the entry compiled for the first one. + expect_identical(as.numeric(tengen::as_array(await(dispatch(d, list(x = pi))$data))), pi) + expect_equal(dispatcher_size(d), 1L) + + # An entry that declares nothing keeps the leaf's own default (f32 for a + # double), which this f64 executable rejects. + d2 <- dispatcher(10L, entry(NULL), default_device = test_pjrt_device) + expect_error(dispatch(d2, list(x = sqrt(2)))) +}) + +test_that("a malformed `input_dtypes` is rejected, not silently ignored", { + skip_if_not(plugins_downloaded()) + src <- 'func.func @main(%x: tensor) -> tensor { + "func.return"(%x): (tensor) -> () + }' + exec <- pjrt_compile(pjrt_program(src = src)) + cb <- function(dtypes) { + function(info) { + pjrt_entry(exec, out_avals = list(oav("f64", integer())), input_dtypes = dtypes) + } + } + expect_error( + dispatch(dispatcher(10L, cb(c("f64", "f64")), default_device = test_pjrt_device), list(x = 1)), + "2 entries but the call supplies 1" + ) + expect_error( + dispatch(dispatcher(10L, cb("f16"), default_device = test_pjrt_device), list(x = 1)), + "not a dtype anvl can represent" + ) + expect_error( + dispatch(dispatcher(10L, cb(64), default_device = test_pjrt_device), list(x = 1)), + "must be a character vector" + ) +})