spin_unlock_irqrestore(&tid_lock, flags);
 }
 
-static int send_mad(struct ib_sa_query *query, int timeout_ms)
+static int send_mad(struct ib_sa_query *query, int timeout_ms, gfp_t gfp_mask)
 {
        unsigned long flags;
        int ret, id;
 
 retry:
-       if (!idr_pre_get(&query_idr, GFP_ATOMIC))
+       if (!idr_pre_get(&query_idr, gfp_mask))
                return -ENOMEM;
        spin_lock_irqsave(&idr_lock, flags);
        ret = idr_get_new(&query_idr, query, &id);
 
        *sa_query = &query->sa_query;
 
-       ret = send_mad(&query->sa_query, timeout_ms);
+       ret = send_mad(&query->sa_query, timeout_ms, gfp_mask);
        if (ret < 0)
                goto err2;
 
 
        *sa_query = &query->sa_query;
 
-       ret = send_mad(&query->sa_query, timeout_ms);
+       ret = send_mad(&query->sa_query, timeout_ms, gfp_mask);
        if (ret < 0)
                goto err2;
 
 
        *sa_query = &query->sa_query;
 
-       ret = send_mad(&query->sa_query, timeout_ms);
+       ret = send_mad(&query->sa_query, timeout_ms, gfp_mask);
        if (ret < 0)
                goto err2;