UCP/AM: Fix request datatype state during CM switch

When switching between transports, we can add new memry registration
handles to req->send.state.dt by calling ucp_send_request_add_reg_lane()
from ucp_do_am_zcopy_single(). Should not save 'state' before calling
add_reg_lane() - otherwise the new memory registration will be
overridden.

This fixes failure on r-vmb-ppc-jenkins in test_ucp_sockaddr_cm_switch
test with the symptom "Error: Input/output error". It is actually a
local protection error (PD violation) due to using wrong uct_mem_h.

src/ucp/proto/proto_am.inl

@@ -297,12 +297,15 @@ ucs_status_t ucp_do_am_zcopy_single(uct_pending_req_t *self, uint8_t am_id,
     ucp_ep_t *ep         = req->send.ep;
     size_t max_iov       = ucp_ep_config(ep)->am.max_iov;
     uct_iov_t *iov       = ucs_alloca(max_iov * sizeof(uct_iov_t));
-    ucp_dt_state_t state = req->send.state.dt;
+    ucp_dt_state_t state;
     ucs_status_t status;
 
     req->send.lane = ucp_ep_get_am_lane(ep);
     ucp_send_request_add_reg_lane(req, req->send.lane);
 
+    /* Cache datatype state only after registering the lane, since registering
+       the lane can change the state */
+    state  = req->send.state.dt;
     status = ucp_am_zcopy_common(req, hdr, hdr_size, user_hdr_desc, user_hdr_size,
                                  iov, max_iov, req->send.length + user_hdr_size,
                                  am_id, &state);