Merge pull request #11549 from devreal/btl_check_device

janjust · web-flow · commit 619e5c6841ab · 2023-04-04T10:41:04.000-05:00
Have btls check for whether convertor involves device memory
diff --git a/ompi/mca/mtl/base/mtl_base_datatype.h b/ompi/mca/mtl/base/mtl_base_datatype.h
@@ -39,19 +39,19 @@ ompi_mtl_datatype_pack(struct opal_convertor_t *convertor,
 {
     struct iovec iov;
     uint32_t iov_count = 1;
+    bool is_accelerator = opal_convertor_on_device(convertor);
 #if !(OPAL_ENABLE_HETEROGENEOUS_SUPPORT)
     if (convertor->pDesc &&
 	!(convertor->flags & CONVERTOR_COMPLETED) &&
 	opal_datatype_is_contiguous_memory_layout(convertor->pDesc,
 						  convertor->count) &&
-        !(convertor->flags & CONVERTOR_ACCELERATOR)) {
+        !is_accelerator) {
 	    *free_after = false;
 	    *buffer = convertor->pBaseBuf + convertor->bConverted + convertor->pDesc->true_lb;
 	    *buffer_len = convertor->local_size;
 	    return OPAL_SUCCESS;
     }
 #endif
-    bool is_accelerator = convertor->flags & CONVERTOR_ACCELERATOR;
 
     opal_convertor_get_packed_size(convertor, buffer_len);
     *free_after = false;
@@ -94,7 +94,7 @@ ompi_mtl_datatype_recv_buf(struct opal_convertor_t *convertor,
         *buffer_len = 0;
         return OMPI_SUCCESS;
     }
-    bool is_accelerator = convertor->flags & CONVERTOR_ACCELERATOR;
+    bool is_accelerator = opal_convertor_on_device(convertor);;
 
     /* If we need buffers or we don't have accelerator support and it is a device buffer, we will need to copy */
     if (opal_convertor_need_buffers(convertor) || (is_accelerator && false == ompi_mtl_base_selected_component->accelerator_support)) {
@@ -122,7 +122,7 @@ ompi_mtl_datatype_unpack(struct opal_convertor_t *convertor,
 {
     struct iovec iov;
     uint32_t iov_count = 1;
-    bool is_accelerator = convertor->flags & CONVERTOR_ACCELERATOR;
+    bool is_accelerator = opal_convertor_on_device(convertor);
 
     /* If the buffer length is greater than 0 and we allocated buffers previously, we need to unpack them */
     if (buffer_len > 0 && (opal_convertor_need_buffers(convertor) || (is_accelerator && false == ompi_mtl_base_selected_component->accelerator_support))) {
diff --git a/opal/datatype/opal_convertor.h b/opal/datatype/opal_convertor.h
@@ -187,6 +187,23 @@ static inline int32_t opal_convertor_need_buffers(const opal_convertor_t *pConve
     return 1;
 }
 
+static inline int32_t opal_convertor_on_device(const opal_convertor_t *pConvertor)
+{
+    return !!(pConvertor->flags & CONVERTOR_ACCELERATOR);
+}
+
+static inline int32_t opal_convertor_on_discrete_device(const opal_convertor_t *pConvertor)
+{
+    return (CONVERTOR_ACCELERATOR == ((pConvertor->flags & CONVERTOR_ACCELERATOR) |
+                                      (pConvertor->flags & CONVERTOR_ACCELERATOR_UNIFIED)));
+}
+
+static inline int32_t opal_convertor_on_unified_device(const opal_convertor_t *pConvertor)
+{
+    return (!!(pConvertor->flags & CONVERTOR_ACCELERATOR) &&
+            !!(pConvertor->flags & CONVERTOR_ACCELERATOR_UNIFIED));
+}
+
 /**
  * Update the size of the remote datatype representation. The size will
  * depend on the configuration of the master convertor. In homogeneous
diff --git a/opal/mca/btl/portals4/btl_portals4.c b/opal/mca/btl/portals4/btl_portals4.c
@@ -533,7 +533,9 @@ mca_btl_base_descriptor_t *mca_btl_portals4_prepare_src(struct mca_btl_base_modu
                          ((struct mca_btl_portals4_module_t *) btl_base)->interface_num,
                          reserve, *size, max_data));
 
-    if (0 != reserve || 0 != opal_convertor_need_buffers(convertor)) {
+    if (0 != reserve ||
+        0 != opal_convertor_need_buffers(convertor) ||
+        0 != opal_convertor_on_device(convertor)) {
         OPAL_OUTPUT_VERBOSE((90, opal_btl_base_framework.framework_output,
                              "mca_btl_portals4_prepare_src NEED BUFFERS or RESERVE\n"));
         frag = (mca_btl_portals4_frag_t *) mca_btl_portals4_alloc(btl_base, peer, MCA_BTL_NO_ORDER,
diff --git a/opal/mca/btl/self/btl_self.c b/opal/mca/btl/self/btl_self.c
@@ -151,7 +151,7 @@ static struct mca_btl_base_descriptor_t *mca_btl_self_prepare_src(
     struct mca_btl_base_module_t *btl, struct mca_btl_base_endpoint_t *endpoint,
     struct opal_convertor_t *convertor, uint8_t order, size_t reserve, size_t *size, uint32_t flags)
 {
-    bool inline_send = !opal_convertor_need_buffers(convertor);
+    bool inline_send = !(opal_convertor_need_buffers(convertor) || opal_convertor_on_device(convertor));
     size_t buffer_len = reserve + (inline_send ? 0 : *size);
     mca_btl_self_frag_t *frag;
 
@@ -229,7 +229,9 @@ static int mca_btl_self_sendi(struct mca_btl_base_module_t *btl,
 {
     mca_btl_base_descriptor_t *frag;
 
-    if (!payload_size || !opal_convertor_need_buffers(convertor)) {
+    if (!payload_size ||
+        !(opal_convertor_need_buffers(convertor) ||
+          opal_convertor_on_device(convertor))) {
         void *data_ptr = NULL;
         if (payload_size) {
             opal_convertor_get_current_pointer(convertor, &data_ptr);
diff --git a/opal/mca/btl/sm/btl_sm_module.c b/opal/mca/btl/sm/btl_sm_module.c
@@ -429,7 +429,10 @@ static struct mca_btl_base_descriptor_t *sm_prepare_src(struct mca_btl_base_modu
     assert(NULL != data_ptr);
 
     /* in place send fragment */
-    if (OPAL_UNLIKELY(opal_convertor_need_buffers(convertor))) {
+    if (OPAL_UNLIKELY(opal_convertor_need_buffers(convertor) ||
+                      opal_convertor_on_discrete_device(convertor) ||
+                      (opal_convertor_on_unified_device(convertor) &&
+                       total_size > (size_t) mca_btl_sm_component.max_inline_send))) {
         uint32_t iov_count = 1;
         struct iovec iov;
 
diff --git a/opal/mca/btl/sm/btl_sm_sendi.c b/opal/mca/btl/sm/btl_sm_sendi.c
@@ -59,7 +59,7 @@ int mca_btl_sm_sendi(struct mca_btl_base_module_t *btl, struct mca_btl_base_endp
         opal_convertor_get_current_pointer(convertor, &data_ptr);
     }
 
-    if (!(payload_size && opal_convertor_need_buffers(convertor))
+    if (!(payload_size && (opal_convertor_need_buffers(convertor) || opal_convertor_on_device(convertor)))
         && mca_btl_sm_fbox_sendi(endpoint, tag, header, header_size, data_ptr, payload_size)) {
         return OPAL_SUCCESS;
     }
diff --git a/opal/mca/btl/tcp/btl_tcp.c b/opal/mca/btl/tcp/btl_tcp.c
@@ -250,7 +250,7 @@ mca_btl_base_descriptor_t *mca_btl_tcp_prepare_src(struct mca_btl_base_module_t
     frag->segments[0].seg_len = reserve;
 
     frag->base.des_segment_count = 1;
-    if (opal_convertor_need_buffers(convertor)) {
+    if (opal_convertor_need_buffers(convertor) || opal_convertor_on_device(convertor)) {
 
         if (max_data + reserve > frag->size) {
             max_data = frag->size - reserve;
diff --git a/opal/mca/btl/ugni/btl_ugni_prepare.h b/opal/mca/btl/ugni/btl_ugni_prepare.h
@@ -179,6 +179,7 @@ mca_btl_ugni_prepare_src_send(struct mca_btl_base_module_t *btl, mca_btl_base_en
 
     send_in_place = (btl->btl_flags & MCA_BTL_FLAGS_SEND_INPLACE)
                     && !(opal_convertor_need_buffers(convertor)
+                         || opal_convertor_on_device(convertor)
                          || (use_eager_get && ((uintptr_t) data_ptr & 3)));
 
     if (send_in_place) {
diff --git a/opal/mca/btl/usnic/btl_usnic_compat.c b/opal/mca/btl/usnic/btl_usnic_compat.c
@@ -91,7 +91,7 @@ prepare_src_small(struct opal_btl_usnic_module_t *module, struct mca_btl_base_en
      * we might still use INLINE for the send, and in that case we do not want
      * to copy the data at all.
      */
-    if (OPAL_UNLIKELY(opal_convertor_need_buffers(convertor))) {
+    if (OPAL_UNLIKELY(opal_convertor_need_buffers(convertor) || opal_convertor_on_device(convertor))) {
         /* put user data just after end of 1st seg (upper layer header) */
         assert(payload_len <= module->max_frag_payload);
         usnic_convertor_pack_simple(convertor,
@@ -227,7 +227,7 @@ static opal_btl_usnic_send_frag_t *prepare_src_large(struct opal_btl_usnic_modul
     /* make sure upper header small enough */
     assert(reserve <= sizeof(lfrag->lsf_ompi_header));
 
-    if (OPAL_UNLIKELY(opal_convertor_need_buffers(convertor))) {
+    if (OPAL_UNLIKELY(opal_convertor_need_buffers(convertor) || opal_convertor_on_device(convertor))) {
         /* threshold == -1 means always pack eagerly */
         if (mca_btl_usnic_component.pack_lazy_threshold >= 0
             && *size >= (size_t) mca_btl_usnic_component.pack_lazy_threshold) {

Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,7 @@ int mca_btl_sm_sendi(struct mca_btl_base_module_t *btl, struct mca_btl_base_endp`
`59`	`59`	`opal_convertor_get_current_pointer(convertor, &data_ptr);`
`60`	`60`	`}`
`61`	`61`
`62`		`- if (!(payload_size && opal_convertor_need_buffers(convertor))`
	`62`	`+ if (!(payload_size && (opal_convertor_need_buffers(convertor) \|\| opal_convertor_on_device(convertor)))`
`63`	`63`	`&& mca_btl_sm_fbox_sendi(endpoint, tag, header, header_size, data_ptr, payload_size)) {`
`64`	`64`	`return OPAL_SUCCESS;`
`65`	`65`	`}`