added input/output keys to umami and matrixelementapi

madgraph/iolibs/template_files/madmatrix/umami.cc

@@ -210,6 +210,37 @@ extern "C"
     return UMAMI_SUCCESS;
   }
 
+  UmamiStatus umami_supported_inputs( bool const** supported, int* count )
+  {
+    // MOMENTA, ALPHA_S, FLAVOR_INDEX, RANDOM_COLOR, RANDOM_HELICITY, RANDOM_DIAGRAM,
+    // HELICITY_INDEX=false, DIAGRAM_INDEX=true, CHANNEL_INDEX=false
+    static const bool data[UMAMI_INPUT_KEY_COUNT] = { true, true, true, true, true, true, false, true };
+    *supported = data;
+    *count = UMAMI_INPUT_KEY_COUNT;
+    return UMAMI_SUCCESS;
+  }
+
+  UmamiStatus umami_required_inputs( bool const** required, int* count )
+  {
+    static const bool data[UMAMI_INPUT_KEY_COUNT] = { true }; // MOMENTA only
+    *required = data;
+    *count = UMAMI_INPUT_KEY_COUNT;
+    return UMAMI_SUCCESS;
+  }
+
+  UmamiStatus umami_supported_outputs( bool const** supported, int* count )
+  {
+    // MATRIX_ELEMENT, DIAGRAM_AMP2, COLOR_INDEX, HELICITY_INDEX, DIAGRAM_INDEX, GPU_STREAM
+#ifdef MGONGPUCPP_GPUIMPL
+    static const bool data[UMAMI_OUTPUT_KEY_COUNT] = { true, true, true, true, true, true };
+#else
+    static const bool data[UMAMI_OUTPUT_KEY_COUNT] = { true, true, true, true, true };
+#endif
+    *supported = data;
+    *count = UMAMI_OUTPUT_KEY_COUNT;
+    return UMAMI_SUCCESS;
+  }
+
   UmamiStatus umami_initialize( UmamiHandle* handle, char const* param_card_path )
   {
     CPPProcess process;

madgraph/iolibs/template_files/mg7/api.cpp

@@ -26,6 +26,30 @@ UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result) {
     return UMAMI_SUCCESS;
 }
 
+UmamiStatus umami_supported_inputs(bool const** supported, int* count) {
+    // MOMENTA, ALPHA_S, FLAVOR_INDEX, RANDOM_COLOR, RANDOM_HELICITY, RANDOM_DIAGRAM
+    static const bool data[UMAMI_INPUT_KEY_COUNT] = { true, true, true, true, true, true };
+    *supported = data;
+    *count = UMAMI_INPUT_KEY_COUNT;
+    return UMAMI_SUCCESS;
+}
+
+UmamiStatus umami_required_inputs(bool const** required, int* count) {
+    // MOMENTA and FLAVOR_INDEX (both dereferenced unconditionally in umami_matrix_element)
+    static const bool data[UMAMI_INPUT_KEY_COUNT] = { true, false, true };
+    *required = data;
+    *count = UMAMI_INPUT_KEY_COUNT;
+    return UMAMI_SUCCESS;
+}
+
+UmamiStatus umami_supported_outputs(bool const** supported, int* count) {
+    // MATRIX_ELEMENT, DIAGRAM_AMP2, COLOR_INDEX, HELICITY_INDEX, DIAGRAM_INDEX
+    static const bool data[UMAMI_OUTPUT_KEY_COUNT] = { true, true, true, true, true };
+    *supported = data;
+    *count = UMAMI_OUTPUT_KEY_COUNT;
+    return UMAMI_SUCCESS;
+}
+
 UmamiStatus umami_initialize(UmamiHandle* handle, char const* param_card_path) {
     CPPProcess* process = new CPPProcess(param_card_path);
     std::vector<double*>& momenta = process->getMomenta();

madgraph/iolibs/template_files/mg7/api.h

@@ -14,6 +14,7 @@
 #ifndef UMAMI_HEADER
 #define UMAMI_HEADER 1
 
+#include <stdbool.h>
 #include <stddef.h>
 
 #ifdef __cplusplus
@@ -66,6 +67,9 @@ typedef enum {
     UMAMI_IN_DIAGRAM_INDEX,
 } UmamiInputKey;
 
+/** Number of values in `UmamiInputKey` */
+#define UMAMI_INPUT_KEY_COUNT (UMAMI_IN_DIAGRAM_INDEX + 1)
+
 typedef enum {
     UMAMI_OUT_MATRIX_ELEMENT,
     UMAMI_OUT_DIAGRAM_AMP2,
@@ -77,6 +81,9 @@ typedef enum {
     // color: LC-ME, FC-ME
 } UmamiOutputKey;
 
+/** Number of values in `UmamiOutputKey` */
+#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_GPU_STREAM + 1)
+
 typedef void* UmamiHandle;
 
 
@@ -94,6 +101,33 @@ typedef void* UmamiHandle;
  */
 UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result);
 
+/**
+ * Reports which input keys can be passed to this matrix element implementation.
+ * Optional — callers that fail to resolve this symbol should treat it as unavailable.
+ *
+ * @param supported pointer set to an implementation-owned boolean array
+ * @param count pointer set to the length of `*supported`
+ */
+UmamiStatus umami_supported_inputs(bool const** supported, int* count);
+
+/**
+ * Reports which input keys are mandatory. Subset of `umami_supported_inputs`.
+ * Optional — callers that fail to resolve this symbol should treat it as unavailable.
+ *
+ * @param required pointer set to an implementation-owned boolean array
+ * @param count pointer set to the length of `*required`
+ */
+UmamiStatus umami_required_inputs(bool const** required, int* count);
+
+/**
+ * Reports which output keys can be requested from this matrix element implementation.
+ * Optional — callers that fail to resolve this symbol should treat it as unavailable.
+ *
+ * @param supported pointer set to an implementation-owned boolean array
+ * @param count pointer set to the length of `*supported`
+ */
+UmamiStatus umami_supported_outputs(bool const** supported, int* count);
+
 /**
  * Creates an instance of the matrix element. Each instance is independent, so thread
  * safety can be achieved by creating a separate one for every thread.

madspace/include/madspace/driver/context.hpp

@@ -47,6 +47,38 @@ class MatrixElementApi {
     }
     std::size_t index() const { return _index; }
     const std::string& file_name() const { return _file_name; }
+    std::vector<bool> supported_inputs() const {
+        bool const* data; int count;
+        check_umami_status(_supported_inputs(&data, &count));
+        std::vector<bool> result(UMAMI_INPUT_KEY_COUNT, false);
+        for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i)
+            result[i] = data[i];
+        return result;
+    }
+    std::vector<bool> required_inputs() const {
+        bool const* data; int count;
+        check_umami_status(_required_inputs(&data, &count));
+        std::vector<bool> supported = supported_inputs();
+        for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) {
+            if (data[i] && !supported[i]) {
+                throw_error(std::format(
+                    "input key {} is reported as required but not as supported", i
+                ));
+            }
+        }
+        std::vector<bool> result(UMAMI_INPUT_KEY_COUNT, false);
+        for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i)
+            result[i] = data[i];
+        return result;
+    }
+    std::vector<bool> supported_outputs() const {
+        bool const* data; int count;
+        check_umami_status(_supported_outputs(&data, &count));
+        std::vector<bool> result(UMAMI_OUTPUT_KEY_COUNT, false);
+        for (int i = 0; i < count && i < UMAMI_OUTPUT_KEY_COUNT; ++i)
+            result[i] = data[i];
+        return result;
+    }
 
     void call(
         UmamiHandle handle,
@@ -89,6 +121,9 @@ class MatrixElementApi {
     [[noreturn]] void throw_error(const std::string& message) const;
     std::unique_ptr<void, std::function<void(void*)>> _shared_lib;
     decltype(&umami_get_meta) _get_meta;
+    decltype(&umami_supported_inputs) _supported_inputs;
+    decltype(&umami_required_inputs) _required_inputs;
+    decltype(&umami_supported_outputs) _supported_outputs;
     decltype(&umami_initialize) _initialize;
     decltype(&umami_matrix_element) _matrix_element;
     decltype(&umami_free) _free;

madspace/include/madspace/umami.h

@@ -14,6 +14,7 @@
 #ifndef UMAMI_HEADER
 #define UMAMI_HEADER 1
 
+#include <stdbool.h>
 #include <stddef.h>
 
 #ifdef __cplusplus
@@ -94,6 +95,9 @@ typedef enum {
     UMAMI_IN_CHANNEL_INDEX,
 } UmamiInputKey;
 
+/** Number of values in `UmamiInputKey` */
+#define UMAMI_INPUT_KEY_COUNT (UMAMI_IN_CHANNEL_INDEX + 1)
+
 typedef enum {
     /** value of the matrix element, type: `double`, shape: `()` */
     UMAMI_OUT_MATRIX_ELEMENT,
@@ -110,6 +114,9 @@ typedef enum {
     UMAMI_OUT_GPU_STREAM,
 } UmamiOutputKey;
 
+/** Number of values in `UmamiOutputKey` */
+#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_GPU_STREAM + 1)
+
 /** Implementation-defined pointer to a matrix element instance */
 typedef void* UmamiHandle;
 
@@ -125,6 +132,64 @@ typedef void* UmamiHandle;
  */
 UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result);
 
+/**
+ * Reports which input keys can be passed to this matrix element implementation, to
+ * be written by whoever implements the UMAMI interface. Implementing this function
+ * is optional; a caller that fails to resolve it should treat it as unavailable.
+ *
+ * @param supported
+ *     pointer set on return to an implementation-owned array of booleans, where
+ *     element `i` indicates whether the input key with numeric value `i`
+ *     (see `UmamiInputKey`) is accepted by this implementation.
+ * @param count
+ *     pointer set on return to the number of elements in `*supported`, reflecting
+ *     the number of `UmamiInputKey` values known to this implementation. A caller
+ *     built against a newer header should treat any key index beyond `*count` as
+ *     not supported.
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_supported_inputs(bool const** supported, int* count);
+
+/**
+ * Reports which input keys must be passed to this matrix element implementation for
+ * `umami_matrix_element` to succeed, to be written by whoever implements the UMAMI
+ * interface. This is a subset of the keys reported by `umami_supported_inputs`.
+ * Implementing this function is optional; a caller that fails to resolve it should
+ * treat it as unavailable.
+ *
+ * @param required
+ *     pointer set on return to an implementation-owned array of booleans, where
+ *     element `i` indicates whether the input key with numeric value `i`
+ *     (see `UmamiInputKey`) is mandatory for this implementation.
+ * @param count
+ *     pointer set on return to the number of elements in `*required`. A caller
+ *     built against a newer header should treat any key index beyond `*count` as
+ *     not required.
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_required_inputs(bool const** required, int* count);
+
+/**
+ * Reports which output keys can be requested from this matrix element
+ * implementation, to be written by whoever implements the UMAMI interface.
+ * Implementing this function is optional; a caller that fails to resolve it should
+ * treat it as unavailable.
+ *
+ * @param supported
+ *     pointer set on return to an implementation-owned array of booleans, where
+ *     element `i` indicates whether the output key with numeric value `i`
+ *     (see `UmamiOutputKey`) can be produced by this implementation.
+ * @param count
+ *     pointer set on return to the number of elements in `*supported`. A caller
+ *     built against a newer header should treat any key index beyond `*count` as
+ *     not supported.
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_supported_outputs(bool const** supported, int* count);
+
 /**
  * Creates an instance of the matrix element. Each instance is independent, so thread
  * safety can be achieved by creating a separate one for every thread.