api: extend dropout attribute with offset host_scalars and s64

Author: mgouicem

doc/programming_model/attributes_dropout.md

@@ -16,32 +16,51 @@ on a deterministic PRNG (current implementation uses a variation of Philox
 algorithm) and transforms the values as follows:
 
 \f[
-    \mathrm{mask}[:] = (\mathrm{PRNG}(S, ...) > P) \\
+    \mathrm{mask}[:] = (\mathrm{PRNG}(S, \mathrm{off}, :) > P) \\
     \mathrm{dst}[:] = \mathrm{mask}[:] \cdot {{\mathrm{dst}[:]} \over {1 - P}}
 \f]
 
 where:
 
-* \f$\mathrm{mask}\f$ is the output buffer (always of the same dimensions and
-usually of the same layout as \f$\mathrm{dst}\f$, but potentially differing from
-it in type that can only be `u8`) whose values may be either 0 if the
-corresponding value in \f$\mathrm{dst}\f$ got zeroed (a.k.a. dropped out) or 1
-otherwise
-* \f$S\f$ is the integer seed for the PRNG algorithm
+* \f$\mathrm{mask}\f$ values may be either 0 if the corresponding value in
+  \f$\mathrm{dst}\f$ got zeroed (a.k.a. dropped out) or 1, otherwise.
+* \f$S, off\f$ are the seed and the offset for the PRNG algorithm.
 * \f$P\f$ is the probability for any given value to get dropped out,
-\f$0 \leq P \leq 1\f$
+  \f$0 \leq P \leq 1\f$
 
 ## API
 
 - C: @ref dnnl_primitive_attr_get_dropout, @ref dnnl_primitive_attr_set_dropout
 - C++: @ref dnnl::primitive_attr::get_dropout, @ref dnnl::primitive_attr::set_dropout
 
-If the dropout operation gets specified in the primitive's attributes, the user
-must provide three additional buffers to it on execution:
+The dropout primitive attribute has the following parameters:
 
-* `DNNL_ARG_ATTR_DROPOUT_MASK`: through this ID the user has to pass the
-\f$\mathrm{mask}\f$ output buffer
-* `DNNL_ARG_ATTR_DROPOUT_PROBABILITY`: this is a single-value `f32` input buffer
-that holds \f$P\f$
-* `DNNL_ARG_ATTR_DROPOUT_SEED`: this is a single-value `s32` input buffer that
-holds \f$S\f$
+* `mask_desc`: when set to a zero (or empty) memory descriptor, mask values are
+  not written to the memory. Otherwise, it should have the same dimensions and
+  the same layout as \f$\mathrm{dst}\f$, as well as `u8` data type.
+* `seed_dt`: data type of the seed argument \f$S\f$, `s64` is recommended, `s32`
+  is supported as a backward compatibility option.
+* `use_offset`: boolean to express if an offset argument will be provided
+  by the user at the execution time. When false, an offset of 0 is assumed.
+* `use_host_scalars`: boolean specifying if probability, seed, and offset memory
+  arguments will be passed as host_scalar memory objects when `true`, or
+  as device memory objects, otherwise.
+
+When the dropout primitive attribute is set, the user must provide two
+additional memory arguments to the primitive execution:
+
+* `DNNL_ARG_ATTR_DROPOUT_PROBABILITY`: this is a single-value `f32` input memory
+  argument that holds \f$P\f$.
+* `DNNL_ARG_ATTR_DROPOUT_SEED`: this is a single-value input memory argument
+  that holds \f$S\f$. Its data type is specified by the `seed_dt` primitive
+  attribute parameter and can be either `s32` or `s64`.
+
+Additionally, the following arguments conditionally need to be passed
+at the execution time as well:
+
+* `DNNL_ARG_ATTR_DROPOUT_MASK`: if the `mask_desc` primitive attribute parameter
+  is not a zero memory descriptor, the user must pass the \f$\mathrm{mask}\f$
+  through this output memory argument.
+* `DNNL_ARG_ATTR_DROPOUT_OFFSET`: if the `use_offset` primitive attribute
+  parameter is set, the user must pass the \f$\mathrm{off}\f$ through this
+  input memory argument. This is a single-value `s64` memory argument.

include/oneapi/dnnl/dnnl.h

@@ -271,24 +271,61 @@ dnnl_status_t DNNL_API dnnl_primitive_attr_clone(
 ///     otherwise.
 dnnl_status_t DNNL_API dnnl_primitive_attr_destroy(dnnl_primitive_attr_t attr);
 
-/// Returns probability for output dropout primitive attribute.
+/// Gets dropout primitive attribute.
 ///
 /// @param attr Primitive attributes.
-/// @param dropout_desc Output dropout memory descriptor
+/// @param mask_desc Output memory descriptor for dropout masks. If a default
+///     memory descriptor is returned, the mask values will not be written to
+///     the output memory buffer during the primitive execution.
 /// @returns #dnnl_success on success and a status describing the error
 ///     otherwise.
 dnnl_status_t DNNL_API dnnl_primitive_attr_get_dropout(
-        const_dnnl_primitive_attr_t attr,
-        const_dnnl_memory_desc_t *dropout_desc);
+        const_dnnl_primitive_attr_t attr, const_dnnl_memory_desc_t *mask_desc);
 
-/// Sets probability for output dropout primitive attribute.
+/// Sets dropout primitive attribute.
 ///
 /// @param attr Primitive attributes.
-/// @param dropout_desc Output dropout memory descriptor
+/// @param mask_desc Memory descriptor for dropout masks. If a default memory
+///     descriptor is passed, the mask values will not be written to the output
+///     memory buffer during the primitive execution.
 /// @returns #dnnl_success on success and a status describing the error
 ///     otherwise.
 dnnl_status_t DNNL_API dnnl_primitive_attr_set_dropout(
-        dnnl_primitive_attr_t attr, const_dnnl_memory_desc_t dropout_desc);
+        dnnl_primitive_attr_t attr, const_dnnl_memory_desc_t mask_desc);
+
+/// Gets dropout primitive attribute parameters.
+///
+/// @param attr Primitive attributes.
+/// @param mask_desc Output memory descriptor for dropout masks. If a default
+///     memory descriptor is returned, the mask values will not be written to
+///     the output memory buffer during the primitive execution.
+/// @param seed_dt Output datatype for seed argument.
+/// @param use_offset Output boolean. If true, an offset argument must be passed
+///     at the execution and will be used in random number generation.
+/// @param use_host_scalars Output boolean. If true, probability, seed and
+///     offset arguments are passed as host_scalar memory objects.
+/// @returns #dnnl_success on success and a status describing the error
+///     otherwise.
+dnnl_status_t DNNL_API dnnl_primitive_attr_get_dropout_v2(
+        const_dnnl_primitive_attr_t attr, const_dnnl_memory_desc_t *mask_desc,
+        dnnl_data_type_t *seed_dt, int *use_offset, int *use_host_scalars);
+
+/// Sets dropout primitive attribute parameters.
+///
+/// @param attr Primitive attributes.
+/// @param mask_desc Memory descriptor for dropout masks. If a default memory
+///     descriptor is passed, the mask values will not be written to the output
+///     memory buffer during the primitive execution.
+/// @param seed_dt Datatype for seed argument.
+/// @param use_offset If true, an offset argument must be passed at the
+///     execution and will be used in random number generation.
+/// @param use_host_scalars If true, probability, seed and offset arguments are
+///     passed as host_scalar memory objects.
+/// @returns #dnnl_success on success and a status describing the error
+///     otherwise.
+dnnl_status_t DNNL_API dnnl_primitive_attr_set_dropout_v2(
+        dnnl_primitive_attr_t attr, const_dnnl_memory_desc_t mask_desc,
+        dnnl_data_type_t seed_dt, int use_offset, int use_host_scalars);
 
 /// Returns the floating-point math mode primitive attribute.
 ///

include/oneapi/dnnl/dnnl.hpp

@@ -4061,7 +4061,9 @@ struct primitive_attr : public handle<dnnl_primitive_attr_t> {
 
     /// Returns the parameters of a dropout attribute.
     ///
-    /// @param mask_desc Output memory descriptor of a dropout mask.
+    /// @param mask_desc Output memory descriptor for dropout masks. If a
+    ///     default memory descriptor is returned, the mask values will not be
+    ///     written to the output memory buffer during the primitive execution.
     void get_dropout(memory::desc &mask_desc) const {
         const_dnnl_memory_desc_t cdesc;
         error::wrap_c_api(dnnl_primitive_attr_get_dropout(get(), &cdesc),
@@ -4072,15 +4074,66 @@ struct primitive_attr : public handle<dnnl_primitive_attr_t> {
         mask_desc = memory::desc(cloned_md);
     }
 
+    /// Returns the parameters of a dropout attribute.
+    ///
+    /// @param mask_desc Output memory descriptor for dropout masks. If a
+    ///     default memory descriptor is returned, the mask values will not be
+    ///     written to the output memory buffer during the primitive execution.
+    /// @param seed_dt Output datatype for seed argument.
+    /// @param use_offset Output boolean. If true, an offset argument must be
+    ///     passed at the execution and will be used in random number
+    ///     generation.
+    /// @param use_host_scalars Output boolean. If true, probability, seed and
+    ///     offset arguments are passed as host_scalar memory objects.
+    void get_dropout(memory::desc &mask_desc, memory::data_type &seed_dt,
+            bool &use_offset, bool &use_host_scalars) const {
+        const_dnnl_memory_desc_t cdesc;
+        dnnl_data_type_t c_seed_dt;
+        int c_use_offset;
+        int c_use_host_scalars;
+        error::wrap_c_api(
+                dnnl_primitive_attr_get_dropout_v2(get(), &cdesc, &c_seed_dt,
+                        &c_use_offset, &c_use_host_scalars),
+                "could not get parameters of a dropout attribute");
+        dnnl_memory_desc_t cloned_md = nullptr;
+        error::wrap_c_api(dnnl_memory_desc_clone(&cloned_md, cdesc),
+                "could not clone a memory descriptor");
+        mask_desc = memory::desc(cloned_md);
+        seed_dt = static_cast<memory::data_type>(c_seed_dt);
+        use_offset = c_use_offset;
+        use_host_scalars = c_use_host_scalars;
+    }
+
     /// Sets dropout probability.
     ///
-    /// @param mask_desc Output memory descriptor of a dropout mask.
+    /// @param mask_desc Memory descriptor for dropout masks. If a default
+    ///     memory descriptor is passed, the mask values will not be written to
+    ///     the output memory buffer during the primitive execution.
     void set_dropout(const memory::desc &mask_desc) {
         error::wrap_c_api(
                 dnnl_primitive_attr_set_dropout(get(), mask_desc.get()),
                 "could not set dropout primitive attribute");
     }
 
+    /// Sets dropout probability.
+    ///
+    /// @param mask_desc Memory descriptor for dropout masks. If a default
+    ///     memory descriptor is passed, the mask values will not be written to
+    ///     the output memory buffer during the primitive execution.
+    /// @param seed_dt Datatype for seed argument.
+    /// @param use_offset If true, an offset argument must be passed at the
+    ///     execution and will be used in random number generation.
+    /// @param use_host_scalars If true, probability, seed and offset arguments
+    ///     are passed as host_scalar memory objects.
+    void set_dropout(const memory::desc &mask_desc, memory::data_type seed_dt,
+            bool use_offset, bool use_host_scalars) {
+        error::wrap_c_api(
+                dnnl_primitive_attr_set_dropout_v2(get(), mask_desc.get(),
+                        memory::convert_to_c(seed_dt), use_offset,
+                        use_host_scalars),
+                "could not set dropout primitive attribute");
+    }
+
     /// Returns the fpmath mode
     fpmath_mode get_fpmath_mode() const {
         dnnl_fpmath_mode_t result;

include/oneapi/dnnl/dnnl_types.h

@@ -2694,6 +2694,9 @@ typedef const struct dnnl_primitive *const_dnnl_primitive_t;
 /// A special mnemonic for shift argument of normalization primitives.
 #define DNNL_ARG_DIFF_SHIFT 256
 
+/// Dropout offset value passed via a buffer
+#define DNNL_ARG_ATTR_DROPOUT_OFFSET 507
+
 /// Rounding mode seed for stochastic rounding
 /// Single seed needed independently of how many arguments need stochastic rounding
 #define DNNL_ARG_ATTR_ROUNDING_SEED 508

scripts/verbose_converter/src/benchdnn_generator.py

@@ -250,6 +250,13 @@ def dropout(self):
         result = "0.5:12345"
         if dropout.tag:
             result += f":{dropout.tag}"
+        # Seed dt is always s64 in benchdnn and is not passed to driver
+        if dropout.use_offset == "1":
+            result += ":987654321"
+        else:
+            result += ":0"
+        if dropout.use_host_scalars == "1":
+            result += f":{dropout.use_host_scalars}"
         return f"--attr-dropout={result}"
 
     deterministic = attribute_flag("deterministic")

scripts/verbose_converter/src/ir.py

@@ -158,9 +158,14 @@ def __hash_str__(self):
 @dataclass(eq=False)
 class Dropout(Mapping):
     tag: Optional[str] = None
+    seed_dt: Optional[str] = None
+    use_offset: Optional[bool] = None
+    use_host_scalars: Optional[bool] = None
 
     def __str__(self):
-        return self.tag or ""
+        return ":".join(
+            [self.tag, self.seed_dt, self.use_offset, self.use_host_scalars]
+        )
 
 
 class FormattedMapping(Mapping, ABC):

scripts/verbose_converter/src/parse.py

@@ -390,7 +390,9 @@ def parse_binary_post_op(spec, alg) -> ir.BinaryPostOp:
 
     @staticmethod
     def parse_dropout(args: str) -> ir.Dropout:
-        return ir.Dropout(tag=args if args else None)
+        spec = ParseSpec(args)
+        fields = args.split(":")
+        return ir.Dropout(fields[0], fields[1], fields[2], fields[3])
 
     @staticmethod
     def parse_per_argument(attr, name, parse):

src/common/math_utils.hpp

@@ -455,15 +455,20 @@ inline bool is_eltwise_ok(
     return eltwise_use_src || eltwise_use_dst;
 }
 
-inline uint32_t philox4x32(uint32_t idx, uint32_t seed) {
+inline uint32_t philox4x32(uint64_t idx, uint64_t seed, uint64_t offset) {
+    // This impl is aligned with PyTorch at
+    // https://github.com/pytorch/pytorch/blob/09c950c/aten/src/ATen/core/PhiloxRNGEngine.h
+    // - both offset and idx are used to fill ctr
+    // - seed/offset/idx are uint64_t
+
     // Note 1: This impl computes 4 different int32_t rand
     //   values. Even though this is redundundant for sequential ref,
     //   keeping vector version to guide optimized implementations.
-    // Note 2: this can be used for 8x16 as well by changing indexing.
 
-    uint32_t x = (idx & ~3L);
-    uint32_t ctr[4] = {x + 0, x + 1, x + 2, x + 3};
-    uint32_t key[2] = {uint32_t(seed), uint32_t(seed)};
+    uint64_t x = (idx & ~3L);
+    uint32_t ctr[4] = {uint32_t(offset), uint32_t(offset >> 32), uint32_t(x),
+            uint32_t(x >> 32)};
+    uint32_t key[2] = {uint32_t(seed), uint32_t(seed >> 32)};
 
     auto mulhilo32 = [&](uint32_t a, uint32_t b, uint32_t &hi, uint32_t &lo) {
         const uint64_t product = static_cast<uint64_t>(a) * b;
@@ -490,30 +495,25 @@ inline uint32_t philox4x32(uint32_t idx, uint32_t seed) {
         key[0] += PHILOX_W4x32_0;
         key[1] += PHILOX_W4x32_1;
     };
-
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
-    philox4x32round();
-    philox4x32bumpkey();
+    constexpr int nrounds = 10;
+    for (int i = 0; i < (nrounds - 1); ++i) {
+        philox4x32round();
+        philox4x32bumpkey();
+    }
     philox4x32round();
 
     return ctr[idx & 3L];
 }
 
+inline uint32_t philox4x32(uint32_t idx, uint32_t seed) {
+    // Note: this is for compatibility with impls that don't support s64 rand
+    uint64_t x = idx & ~3L;
+    uint64_t idx_64 = ((x + 3) << 32) + (x + 2);
+    uint64_t offset_64 = ((x + 1) << 32) + x;
+    uint64_t seed_64 = (uint64_t(seed) << 32) + seed;
+    return philox4x32(idx_64, seed_64, offset_64);
+}
+
 inline uint16_t philox8x16(uint32_t idx, uint32_t seed) {
     // we split the index in two parts:
     // - 31 msb are used to generate 32 random bits