<commit_before>/*************************************************************************
 *
 *  OpenOffice.org - a multi-platform office productivity suite
 *
 *  $RCSfile: fmtwrapinfluenceonobjpos.cxx,v $
 *
 *  $Revision: 1.4 $
 *
 *  last change: $Author: rt $ $Date: 2005-09-09 02:57:41 $
 *
 *  The Contents of this file are made available subject to
 *  the terms of GNU Lesser General Public License Version 2.1.
 *
 *
 *    GNU Lesser General Public License Version 2.1
 *    =============================================
 *    Copyright 2005 by Sun Microsystems, Inc.
 *    901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *    This library is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License version 2.1, as published by the Free Software Foundation.
 *
 *    This library is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *    Lesser General Public License for more details.
 *
 *    You should have received a copy of the GNU Lesser General Public
 *    License along with this library; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *    MA  02111-1307  USA
 *
 ************************************************************************/

#ifndef _FMTWRAPINFLUENCEONOBJPOS_HXX
#include <fmtwrapinfluenceonobjpos.hxx>
#endif

#ifndef _UNOMID_H
#include <unomid.h>
#endif

using namespace ::com::sun::star::uno;

TYPEINIT1(SwFmtWrapInfluenceOnObjPos, SfxPoolItem);

SwFmtWrapInfluenceOnObjPos::SwFmtWrapInfluenceOnObjPos(
                                            sal_Int16 _nWrapInfluenceOnPosition )
    : SfxPoolItem( RES_WRAP_INFLUENCE_ON_OBJPOS ),
    mnWrapInfluenceOnPosition( _nWrapInfluenceOnPosition )
{
}

SwFmtWrapInfluenceOnObjPos::SwFmtWrapInfluenceOnObjPos(
                                        const SwFmtWrapInfluenceOnObjPos& _rCpy )
    : SfxPoolItem( RES_WRAP_INFLUENCE_ON_OBJPOS ),
    mnWrapInfluenceOnPosition( _rCpy.GetWrapInfluenceOnObjPos() )
{
}

SwFmtWrapInfluenceOnObjPos::~SwFmtWrapInfluenceOnObjPos()
{
}

SwFmtWrapInfluenceOnObjPos& SwFmtWrapInfluenceOnObjPos::operator=(
                                    const SwFmtWrapInfluenceOnObjPos& _rSource )
{
    mnWrapInfluenceOnPosition = _rSource.GetWrapInfluenceOnObjPos();

    return *this;
}

int SwFmtWrapInfluenceOnObjPos::operator==( const SfxPoolItem& _rAttr ) const
{
    ASSERT( SfxPoolItem::operator==( _rAttr ), "keine gleichen Attribute" );
    return ( mnWrapInfluenceOnPosition ==
                    static_cast<const SwFmtWrapInfluenceOnObjPos&>(_rAttr).
                                                GetWrapInfluenceOnObjPos() );
}

SfxPoolItem* SwFmtWrapInfluenceOnObjPos::Clone( SfxItemPool *pPool ) const
{
    return new SwFmtWrapInfluenceOnObjPos(*this);
}

BOOL SwFmtWrapInfluenceOnObjPos::QueryValue( Any& rVal, BYTE nMemberId ) const
{
    nMemberId &= ~CONVERT_TWIPS;
    sal_Bool bRet = sal_True;
    switch ( nMemberId )
    {
        case MID_WRAP_INFLUENCE:
        {
            rVal <<= GetWrapInfluenceOnObjPos();
        }
        break;
        default:
            ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::QueryValue()> - unknown MemberId" );
            bRet = sal_False;
    }

    return bRet;
}

BOOL SwFmtWrapInfluenceOnObjPos::PutValue( const Any& rVal, BYTE nMemberId )
{
    nMemberId &= ~CONVERT_TWIPS;
    sal_Bool bRet = sal_True;

    switch ( nMemberId )
    {
        case MID_WRAP_INFLUENCE:
        {
            sal_Int16 nNewWrapInfluence;
            rVal >>= nNewWrapInfluence;
            // --> OD 2004-10-18 #i35017# - constant names have changed and
            // <ITERATIVE> has been added
            if ( nNewWrapInfluence == text::WrapInfluenceOnPosition::ONCE_SUCCESSIVE ||
                 nNewWrapInfluence == text::WrapInfluenceOnPosition::ONCE_CONCURRENT ||
                 nNewWrapInfluence == text::WrapInfluenceOnPosition::ITERATIVE )
            // <--
            {
                SetWrapInfluenceOnObjPos( nNewWrapInfluence );
            }
            else
            {
                ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::PutValue(..)> - invalid attribute value" );
                bRet = sal_False;
            }
        }
        break;
        default:
            ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::QueryValue()> - unknown MemberId" );
            bRet = sal_False;
    }

    return bRet;
}

void SwFmtWrapInfluenceOnObjPos::SetWrapInfluenceOnObjPos( sal_Int16 _nWrapInfluenceOnPosition )
{
    // --> OD 2004-10-18 #i35017# - constant names have changed and consider
    // new value <ITERATIVE>
    if ( _nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ONCE_SUCCESSIVE ||
         _nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ONCE_CONCURRENT ||
         _nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ITERATIVE )
    // <--
    {
        mnWrapInfluenceOnPosition = _nWrapInfluenceOnPosition;
    }
    else
    {
        ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::SetWrapInfluenceOnObjPos(..)> - invalid attribute value" );
    }
}

// --> OD 2004-10-18 #i35017# - add parameter <_bIterativeAsOnceConcurrent>
// to control, if value <ITERATIVE> has to be treated as <ONCE_CONCURRENT>
sal_Int16 SwFmtWrapInfluenceOnObjPos::GetWrapInfluenceOnObjPos(
                                const bool _bIterativeAsOnceConcurrent ) const
{
    sal_Int16 nWrapInfluenceOnPosition( mnWrapInfluenceOnPosition );

    if ( _bIterativeAsOnceConcurrent &&
         nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ITERATIVE )
    {
        nWrapInfluenceOnPosition = text::WrapInfluenceOnPosition::ONCE_CONCURRENT;
    }

    return nWrapInfluenceOnPosition;
}
// <--
<commit_msg>INTEGRATION: CWS pchfix02 (1.4.474); FILE MERGED 2006/09/01 17:51:15 kaib 1.4.474.1: #i68856# Added header markers and pch files<commit_after>/*************************************************************************
 *
 *  OpenOffice.org - a multi-platform office productivity suite
 *
 *  $RCSfile: fmtwrapinfluenceonobjpos.cxx,v $
 *
 *  $Revision: 1.5 $
 *
 *  last change: $Author: obo $ $Date: 2006-09-16 20:40:02 $
 *
 *  The Contents of this file are made available subject to
 *  the terms of GNU Lesser General Public License Version 2.1.
 *
 *
 *    GNU Lesser General Public License Version 2.1
 *    =============================================
 *    Copyright 2005 by Sun Microsystems, Inc.
 *    901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *    This library is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License version 2.1, as published by the Free Software Foundation.
 *
 *    This library is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *    Lesser General Public License for more details.
 *
 *    You should have received a copy of the GNU Lesser General Public
 *    License along with this library; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *    MA  02111-1307  USA
 *
 ************************************************************************/

// MARKER(update_precomp.py): autogen include statement, do not remove
#include "precompiled_sw.hxx"

#ifndef _FMTWRAPINFLUENCEONOBJPOS_HXX
#include <fmtwrapinfluenceonobjpos.hxx>
#endif

#ifndef _UNOMID_H
#include <unomid.h>
#endif

using namespace ::com::sun::star::uno;

TYPEINIT1(SwFmtWrapInfluenceOnObjPos, SfxPoolItem);

SwFmtWrapInfluenceOnObjPos::SwFmtWrapInfluenceOnObjPos(
                                            sal_Int16 _nWrapInfluenceOnPosition )
    : SfxPoolItem( RES_WRAP_INFLUENCE_ON_OBJPOS ),
    mnWrapInfluenceOnPosition( _nWrapInfluenceOnPosition )
{
}

SwFmtWrapInfluenceOnObjPos::SwFmtWrapInfluenceOnObjPos(
                                        const SwFmtWrapInfluenceOnObjPos& _rCpy )
    : SfxPoolItem( RES_WRAP_INFLUENCE_ON_OBJPOS ),
    mnWrapInfluenceOnPosition( _rCpy.GetWrapInfluenceOnObjPos() )
{
}

SwFmtWrapInfluenceOnObjPos::~SwFmtWrapInfluenceOnObjPos()
{
}

SwFmtWrapInfluenceOnObjPos& SwFmtWrapInfluenceOnObjPos::operator=(
                                    const SwFmtWrapInfluenceOnObjPos& _rSource )
{
    mnWrapInfluenceOnPosition = _rSource.GetWrapInfluenceOnObjPos();

    return *this;
}

int SwFmtWrapInfluenceOnObjPos::operator==( const SfxPoolItem& _rAttr ) const
{
    ASSERT( SfxPoolItem::operator==( _rAttr ), "keine gleichen Attribute" );
    return ( mnWrapInfluenceOnPosition ==
                    static_cast<const SwFmtWrapInfluenceOnObjPos&>(_rAttr).
                                                GetWrapInfluenceOnObjPos() );
}

SfxPoolItem* SwFmtWrapInfluenceOnObjPos::Clone( SfxItemPool *pPool ) const
{
    return new SwFmtWrapInfluenceOnObjPos(*this);
}

BOOL SwFmtWrapInfluenceOnObjPos::QueryValue( Any& rVal, BYTE nMemberId ) const
{
    nMemberId &= ~CONVERT_TWIPS;
    sal_Bool bRet = sal_True;
    switch ( nMemberId )
    {
        case MID_WRAP_INFLUENCE:
        {
            rVal <<= GetWrapInfluenceOnObjPos();
        }
        break;
        default:
            ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::QueryValue()> - unknown MemberId" );
            bRet = sal_False;
    }

    return bRet;
}

BOOL SwFmtWrapInfluenceOnObjPos::PutValue( const Any& rVal, BYTE nMemberId )
{
    nMemberId &= ~CONVERT_TWIPS;
    sal_Bool bRet = sal_True;

    switch ( nMemberId )
    {
        case MID_WRAP_INFLUENCE:
        {
            sal_Int16 nNewWrapInfluence;
            rVal >>= nNewWrapInfluence;
            // --> OD 2004-10-18 #i35017# - constant names have changed and
            // <ITERATIVE> has been added
            if ( nNewWrapInfluence == text::WrapInfluenceOnPosition::ONCE_SUCCESSIVE ||
                 nNewWrapInfluence == text::WrapInfluenceOnPosition::ONCE_CONCURRENT ||
                 nNewWrapInfluence == text::WrapInfluenceOnPosition::ITERATIVE )
            // <--
            {
                SetWrapInfluenceOnObjPos( nNewWrapInfluence );
            }
            else
            {
                ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::PutValue(..)> - invalid attribute value" );
                bRet = sal_False;
            }
        }
        break;
        default:
            ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::QueryValue()> - unknown MemberId" );
            bRet = sal_False;
    }

    return bRet;
}

void SwFmtWrapInfluenceOnObjPos::SetWrapInfluenceOnObjPos( sal_Int16 _nWrapInfluenceOnPosition )
{
    // --> OD 2004-10-18 #i35017# - constant names have changed and consider
    // new value <ITERATIVE>
    if ( _nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ONCE_SUCCESSIVE ||
         _nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ONCE_CONCURRENT ||
         _nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ITERATIVE )
    // <--
    {
        mnWrapInfluenceOnPosition = _nWrapInfluenceOnPosition;
    }
    else
    {
        ASSERT( false, "<SwFmtWrapInfluenceOnObjPos::SetWrapInfluenceOnObjPos(..)> - invalid attribute value" );
    }
}

// --> OD 2004-10-18 #i35017# - add parameter <_bIterativeAsOnceConcurrent>
// to control, if value <ITERATIVE> has to be treated as <ONCE_CONCURRENT>
sal_Int16 SwFmtWrapInfluenceOnObjPos::GetWrapInfluenceOnObjPos(
                                const bool _bIterativeAsOnceConcurrent ) const
{
    sal_Int16 nWrapInfluenceOnPosition( mnWrapInfluenceOnPosition );

    if ( _bIterativeAsOnceConcurrent &&
         nWrapInfluenceOnPosition == text::WrapInfluenceOnPosition::ITERATIVE )
    {
        nWrapInfluenceOnPosition = text::WrapInfluenceOnPosition::ONCE_CONCURRENT;
    }

    return nWrapInfluenceOnPosition;
}
// <--
<|endoftext|>


<commit_before>/* mbed Microcontroller Library
 * Copyright (c) 2016 Realtek Semiconductor Corp.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#if DEVICE_EMAC

#include <stdio.h>
#include "mbed_assert.h"
#include "mbed_events.h"

#include "emac_api.h"
#include "rtos.h"

#include "lwip/pbuf.h"
#include "netif/etharp.h"

#include "lwip_intf.h"
#include "wifi_constants.h"
#include "wifi_conf.h"

#define RTW_EMAC_MTU_SIZE  (1500U)

static emac_interface_t *_emac;
static emac_link_input_fn link_input_cb;
static emac_link_state_change_fn link_state_cb;
static void *link_input_data;
static void *link_state_data;

static uint32_t wlan_get_mtu_size(emac_interface_t *emac)
{
    return RTW_EMAC_MTU_SIZE;
}

static void wlan_get_ifname(emac_interface_t *emac, char *name, uint8_t size)
{
    MBED_ASSERT(name != NULL);
    strncpy(name, "r0", size);
}

static uint8_t wlan_get_hwaddr_size(emac_interface_t *emac)
{
    return ETHARP_HWADDR_LEN;
}

static void wlan_get_hwaddr(emac_interface_t *emac, uint8_t *addr)
{
    char mac[20];
    if(RTW_SUCCESS == wifi_get_mac_address(mac))
    {
        if (sscanf(mac, "%x:%x:%x:%x:%x:%x", &addr[0], &addr[1], &addr[2], &addr[3], &addr[4], &addr[5]) != 6)
            printf("Get HW address failed\r\n");
    }else{
            printf("Get HW address failed\r\n");
    }
}

static void wlan_set_hwaddr(emac_interface_t *emac, uint8_t *addr)
{

}

static bool wlan_link_out(emac_interface_t *emac, emac_stack_mem_t *buf)
{
    struct eth_drv_sg * sg_list=0;
    int sg_len = 0;
    int tot_len;
    struct pbuf *p;
    bool ret = true;

    if (!rltk_wlan_running(0)) {
        return false;
    }

    sg_list = (struct eth_drv_sg *)malloc(sizeof(struct eth_drv_sg)*MAX_ETH_DRV_SG);
    if(sg_list == 0){//malloc fail
        return false;
    }
    emac_stack_mem_ref(emac, buf);

    p = (struct pbuf *)buf;
    tot_len = p->tot_len;
    for (; p != NULL && sg_len < MAX_ETH_DRV_SG; p = p->next) {
        sg_list[sg_len].buf = (uint32_t) p->payload;
        sg_list[sg_len].len = p->len;
    sg_len++;
    }

    if (sg_len) {
        if (rltk_wlan_send(0, sg_list, sg_len, tot_len) != 0) {
            ret = false;
        }
    }

    emac_stack_mem_free(emac, buf);
    free(sg_list);
    return ret;
}

static bool wlan_power_up(emac_interface_t *emac)
{
    wifi_on(RTW_MODE_STA);
    wait_ms(1000);
    return true;
}

static void wlan_power_down(emac_interface_t *emac)
{
    wifi_off();
}

static void wlan_set_link_input_cb(emac_interface_t *emac, emac_link_input_fn cb, void *data)
{
    link_input_cb = cb;
    link_input_data = data;
}

static void wlan_set_link_state_cb(emac_interface_t *emac, emac_link_state_change_fn cb, void *data)
{
    link_state_cb = cb;
    link_state_data = data;
}

void wlan_emac_recv(struct netif *netif, int len)
{
    struct eth_drv_sg sg_list[MAX_ETH_DRV_SG];
    emac_stack_mem_t *buf;
    struct pbuf *p;
    int sg_len = 0;

    if (!rltk_wlan_running(0)) {
        return;
    }

    if (len > MAX_ETH_MSG || len < 0) {
        len = MAX_ETH_MSG;
    }

    buf = emac_stack_mem_alloc(NULL, len, 0);
    if (buf == NULL) {
        return;
    }

    p = (struct pbuf *)buf;
    for (; p != NULL && sg_len < MAX_ETH_DRV_SG; p = p->next) {
        sg_list[sg_len].buf = (uint32_t) p->payload;
        sg_list[sg_len].len = p->len;
        sg_len++;
    }
    rltk_wlan_recv(0, sg_list, sg_len);

    if (link_input_cb) {
        link_input_cb(link_input_data, buf);
    }
    return;
}

const emac_interface_ops_t wlan_emac_interface = {
    .get_mtu_size = wlan_get_mtu_size,
    .get_ifname = wlan_get_ifname,
    .get_hwaddr_size = wlan_get_hwaddr_size,
    .get_hwaddr = wlan_get_hwaddr,
    .set_hwaddr = wlan_set_hwaddr,
    .link_out = wlan_link_out,
    .power_up = wlan_power_up,
    .power_down = wlan_power_down,
    .set_link_input_cb = wlan_set_link_input_cb,
    .set_link_state_cb = wlan_set_link_state_cb
};

void mbed_default_mac_address(char *mac) {
    unsigned char RTK_mac_addr[3] = {0x00, 0xE0, 0x4C}; // default Realtek mac address

    mac[0] = RTK_mac_addr[0];
    mac[1] = RTK_mac_addr[1];
    mac[2] = RTK_mac_addr[2];
    mac[3] = 0x87;
    mac[4] = 0x00;
    mac[5] = 0x01;

    return;
}

void mbed_mac_address(char *mac)
{
    char hwaddr[20];
    if(RTW_SUCCESS == wifi_get_mac_address(hwaddr))
    {
        if (sscanf(hwaddr, "%x:%x:%x:%x:%x:%x", &mac[0], &mac[1], &mac[2], &mac[3], &mac[4], &mac[5]) != 6)
            printf("Get HW address failed\r\n");
    }else{
        printf("Get HW address failed\r\n");
        mbed_default_mac_address(mac);
    }
}

void wlan_emac_link_change(bool up)
{
    if (link_state_cb) {
        link_state_cb(link_state_data, up);
    }
}

emac_interface_t *wlan_emac_init_interface()
{

    if (_emac == NULL) {
        _emac = (emac_interface_t*) malloc(sizeof(emac_interface_t));
    	if (_emac == NULL) {//new emac_interface_t fail
        	printf("emac initialization failed\r\n");
		return NULL;
    	}
        _emac->hw = NULL;
        memcpy((void*)&_emac->ops, &wlan_emac_interface, sizeof(wlan_emac_interface));
    }
    return _emac;
}
#endif
<commit_msg>Added assert for malloc<commit_after>/* mbed Microcontroller Library
 * Copyright (c) 2016 Realtek Semiconductor Corp.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#if DEVICE_EMAC

#include <stdio.h>
#include "mbed_assert.h"
#include "mbed_events.h"

#include "emac_api.h"
#include "rtos.h"

#include "lwip/pbuf.h"
#include "netif/etharp.h"

#include "lwip_intf.h"
#include "wifi_constants.h"
#include "wifi_conf.h"

#define RTW_EMAC_MTU_SIZE  (1500U)

static emac_interface_t *_emac;
static emac_link_input_fn link_input_cb;
static emac_link_state_change_fn link_state_cb;
static void *link_input_data;
static void *link_state_data;

static uint32_t wlan_get_mtu_size(emac_interface_t *emac)
{
    return RTW_EMAC_MTU_SIZE;
}

static void wlan_get_ifname(emac_interface_t *emac, char *name, uint8_t size)
{
    MBED_ASSERT(name != NULL);
    strncpy(name, "r0", size);
}

static uint8_t wlan_get_hwaddr_size(emac_interface_t *emac)
{
    return ETHARP_HWADDR_LEN;
}

static void wlan_get_hwaddr(emac_interface_t *emac, uint8_t *addr)
{
    char mac[20];
    if(RTW_SUCCESS == wifi_get_mac_address(mac))
    {
        if (sscanf(mac, "%x:%x:%x:%x:%x:%x", &addr[0], &addr[1], &addr[2], &addr[3], &addr[4], &addr[5]) != 6)
            printf("Get HW address failed\r\n");
    }else{
            printf("Get HW address failed\r\n");
    }
}

static void wlan_set_hwaddr(emac_interface_t *emac, uint8_t *addr)
{

}

static bool wlan_link_out(emac_interface_t *emac, emac_stack_mem_t *buf)
{
    struct eth_drv_sg * sg_list=0;
    int sg_len = 0;
    int tot_len;
    struct pbuf *p;
    bool ret = true;

    if (!rltk_wlan_running(0)) {
        return false;
    }

    sg_list = (struct eth_drv_sg *)malloc(sizeof(struct eth_drv_sg)*MAX_ETH_DRV_SG);
    if(sg_list == 0){//malloc fail
        return false;
    }
    emac_stack_mem_ref(emac, buf);

    p = (struct pbuf *)buf;
    tot_len = p->tot_len;
    for (; p != NULL && sg_len < MAX_ETH_DRV_SG; p = p->next) {
        sg_list[sg_len].buf = (uint32_t) p->payload;
        sg_list[sg_len].len = p->len;
    sg_len++;
    }

    if (sg_len) {
        if (rltk_wlan_send(0, sg_list, sg_len, tot_len) != 0) {
            ret = false;
        }
    }

    emac_stack_mem_free(emac, buf);
    free(sg_list);
    return ret;
}

static bool wlan_power_up(emac_interface_t *emac)
{
    wifi_on(RTW_MODE_STA);
    wait_ms(1000);
    return true;
}

static void wlan_power_down(emac_interface_t *emac)
{
    wifi_off();
}

static void wlan_set_link_input_cb(emac_interface_t *emac, emac_link_input_fn cb, void *data)
{
    link_input_cb = cb;
    link_input_data = data;
}

static void wlan_set_link_state_cb(emac_interface_t *emac, emac_link_state_change_fn cb, void *data)
{
    link_state_cb = cb;
    link_state_data = data;
}

void wlan_emac_recv(struct netif *netif, int len)
{
    struct eth_drv_sg sg_list[MAX_ETH_DRV_SG];
    emac_stack_mem_t *buf;
    struct pbuf *p;
    int sg_len = 0;

    if (!rltk_wlan_running(0)) {
        return;
    }

    if (len > MAX_ETH_MSG || len < 0) {
        len = MAX_ETH_MSG;
    }

    buf = emac_stack_mem_alloc(NULL, len, 0);
    if (buf == NULL) {
        return;
    }

    p = (struct pbuf *)buf;
    for (; p != NULL && sg_len < MAX_ETH_DRV_SG; p = p->next) {
        sg_list[sg_len].buf = (uint32_t) p->payload;
        sg_list[sg_len].len = p->len;
        sg_len++;
    }
    rltk_wlan_recv(0, sg_list, sg_len);

    if (link_input_cb) {
        link_input_cb(link_input_data, buf);
    }
    return;
}

const emac_interface_ops_t wlan_emac_interface = {
    .get_mtu_size = wlan_get_mtu_size,
    .get_ifname = wlan_get_ifname,
    .get_hwaddr_size = wlan_get_hwaddr_size,
    .get_hwaddr = wlan_get_hwaddr,
    .set_hwaddr = wlan_set_hwaddr,
    .link_out = wlan_link_out,
    .power_up = wlan_power_up,
    .power_down = wlan_power_down,
    .set_link_input_cb = wlan_set_link_input_cb,
    .set_link_state_cb = wlan_set_link_state_cb
};

void mbed_default_mac_address(char *mac) {
    unsigned char RTK_mac_addr[3] = {0x00, 0xE0, 0x4C}; // default Realtek mac address

    mac[0] = RTK_mac_addr[0];
    mac[1] = RTK_mac_addr[1];
    mac[2] = RTK_mac_addr[2];
    mac[3] = 0x87;
    mac[4] = 0x00;
    mac[5] = 0x01;

    return;
}

void mbed_mac_address(char *mac)
{
    char hwaddr[20];
    if(RTW_SUCCESS == wifi_get_mac_address(hwaddr))
    {
        if (sscanf(hwaddr, "%x:%x:%x:%x:%x:%x", &mac[0], &mac[1], &mac[2], &mac[3], &mac[4], &mac[5]) != 6)
            printf("Get HW address failed\r\n");
    }else{
        printf("Get HW address failed\r\n");
        mbed_default_mac_address(mac);
    }
}

void wlan_emac_link_change(bool up)
{
    if (link_state_cb) {
        link_state_cb(link_state_data, up);
    }
}

emac_interface_t *wlan_emac_init_interface()
{

    if (_emac == NULL) {
        _emac = (emac_interface_t*) malloc(sizeof(emac_interface_t));
        MBED_ASSERT(_emac);
        _emac->hw = NULL;
        memcpy((void*)&_emac->ops, &wlan_emac_interface, sizeof(wlan_emac_interface));
    }
    return _emac;
}
#endif
<|endoftext|>


<commit_before>// Copyright 2009-2020 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

// ospray
#include "Device.h"
#include "rkcommon/os/library.h"
#include "rkcommon/tasking/tasking_system_init.h"
#include "rkcommon/utility/getEnvVar.h"

#include <map>

namespace ospray {
namespace api {

static FactoryMap<Device> g_devicesMap;

// Helper functions ///////////////////////////////////////////////////////////

template <typename OSTREAM_T>
static inline void installStatusMsgFunc(Device &device, OSTREAM_T &stream)
{
  device.msg_fcn = [&](const char *msg) { stream << msg; };
}

template <typename OSTREAM_T>
static inline void installErrorMsgFunc(Device &device, OSTREAM_T &stream)
{
  device.error_fcn = [&](OSPError e, const char *msg) {
    stream << "OSPRAY ERROR [" << e << "]: " << msg << std::endl;
  };
}

// Device definitions /////////////////////////////////////////////////////////

memory::IntrusivePtr<Device> Device::current;
uint32_t Device::logLevel = OSP_LOG_NONE;

Device *Device::createDevice(const char *type)
{
  // NOTE(jda) - If a user is manually creating the device (i.e. not using
  //             ospInit() to do it), then we need to check if there's a
  //             valid library for core ospray in our main symbol lookup
  //             table.
  auto &repo = *LibraryRepository::getInstance();
  if (!repo.libraryExists("ospray"))
    repo.addDefaultLibrary();

  if (!repo.libraryExists("ospray_module_ispc") && type == std::string("cpu"))
    repo.add("ospray_module_ispc");

  return createInstanceHelper(type, g_devicesMap[type]);
}

void Device::registerType(const char *type, FactoryFcn<Device> f)
{
  g_devicesMap[type] = f;
}

void Device::commit()
{
  auto OSPRAY_DEBUG = utility::getEnvVar<int>("OSPRAY_DEBUG");
  debugMode = OSPRAY_DEBUG.value_or(getParam<bool>("debug", 0));

  auto OSPRAY_WARN = utility::getEnvVar<int>("OSPRAY_WARN_AS_ERROR");
  warningsAreErrors =
      OSPRAY_WARN.value_or(getParam<bool>("warnAsError", false));

  auto OSPRAY_TRACE_API = utility::getEnvVar<int>("OSPRAY_TRACE_API");
  bool traceAPI = OSPRAY_TRACE_API.value_or(getParam<bool>("traceApi", false));

  if (traceAPI && !apiTraceEnabled) {
    auto streamPtr = std::make_shared<std::ofstream>("ospray_api_trace.txt");

    trace_fcn = [=](const char *message) {
      auto &stream = *streamPtr;
      stream << message << std::endl;
    };
  } else if (!traceAPI) {
    trace_fcn = [](const char *) {};
  }

  apiTraceEnabled = traceAPI;

  logLevel = getParam<int>("logLevel", OSP_LOG_NONE);

  auto logLevelStr =
      utility::getEnvVar<std::string>("OSPRAY_LOG_LEVEL").value_or("");

  logLevel = logLevelFromString(logLevelStr).value_or(logLevel);

  auto OSPRAY_NUM_THREADS = utility::getEnvVar<int>("OSPRAY_NUM_THREADS");
  numThreads = OSPRAY_NUM_THREADS.value_or(getParam<int>("numThreads", -1));

  auto OSPRAY_LOG_OUTPUT = utility::getEnvVar<std::string>("OSPRAY_LOG_OUTPUT");

  auto dst = OSPRAY_LOG_OUTPUT.value_or(getParam<std::string>("logOutput", ""));

  if (dst == "cout")
    installStatusMsgFunc(*this, std::cout);
  else if (dst == "cerr")
    installStatusMsgFunc(*this, std::cerr);
  else if (dst == "none")
    msg_fcn = [](const char *) {};

  auto OSPRAY_ERROR_OUTPUT =
      utility::getEnvVar<std::string>("OSPRAY_ERROR_OUTPUT");

  dst = OSPRAY_ERROR_OUTPUT.value_or(getParam<std::string>("errorOutput", ""));

  if (dst == "cout")
    installErrorMsgFunc(*this, std::cout);
  else if (dst == "cerr")
    installErrorMsgFunc(*this, std::cerr);
  else if (dst == "none")
    error_fcn = [](OSPError, const char *) {};

  if (debugMode) {
    logLevel = OSP_LOG_DEBUG;
    numThreads = 1;
    installStatusMsgFunc(*this, std::cout);
    installErrorMsgFunc(*this, std::cerr);
  }

  threadAffinity = AUTO_DETECT;

  auto OSPRAY_SET_AFFINITY = utility::getEnvVar<int>("OSPRAY_SET_AFFINITY");
  if (OSPRAY_SET_AFFINITY)
    threadAffinity = OSPRAY_SET_AFFINITY.value();

  if (hasParam("setAffinity"))
    threadAffinity = getParam<bool>("setAffinity", threadAffinity != 0);

  tasking::initTaskingSystem(numThreads);

  committed = true;
}

bool Device::isCommitted()
{
  return committed;
}

bool deviceIsSet()
{
  return Device::current.ptr != nullptr;
}

Device &currentDevice()
{
  return *Device::current;
}

std::string generateEmbreeDeviceCfg(const Device &device)
{
  std::stringstream embreeConfig;

  if (device.debugMode)
    embreeConfig << " verbose=2";

  if (device.threadAffinity == api::Device::AFFINITIZE)
    embreeConfig << ",set_affinity=1";
  else if (device.threadAffinity == api::Device::DEAFFINITIZE)
    embreeConfig << ",set_affinity=0";

  return embreeConfig.str();
}

int64_t Device::getProperty(const OSPDeviceProperty prop)
{
  /* documented properties */
  switch (prop) {
  case OSP_DEVICE_VERSION:
    return 10000 * OSPRAY_VERSION_MAJOR + 100 * OSPRAY_VERSION_MINOR
        + OSPRAY_VERSION_PATCH;
  case OSP_DEVICE_VERSION_MAJOR:
    return OSPRAY_VERSION_MAJOR;
  case OSP_DEVICE_VERSION_MINOR:
    return OSPRAY_VERSION_MINOR;
  case OSP_DEVICE_VERSION_PATCH:
    return OSPRAY_VERSION_PATCH;
  case OSP_DEVICE_SO_VERSION:
    return OSPRAY_SOVERSION;
  default:
    handleError(OSP_INVALID_ARGUMENT, "unknown readable property");
    return 0;
  }
}

utility::Optional<int> Device::logLevelFromString(const std::string &str)
{
  utility::Optional<int> retval;

  if (str == "none")
    retval = OSP_LOG_NONE;
  else if (str == "debug")
    retval = OSP_LOG_DEBUG;
  else if (str == "info")
    retval = OSP_LOG_INFO;
  else if (str == "warning")
    retval = OSP_LOG_WARNING;
  else if (str == "error")
    retval = OSP_LOG_ERROR;

  return retval;
}

} // namespace api

OSPTYPEFOR_DEFINITION(api::Device *);

} // namespace ospray
<commit_msg>OSPRAY_SET_AFFINITY has precedence<commit_after>// Copyright 2009-2020 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

// ospray
#include "Device.h"
#include "rkcommon/os/library.h"
#include "rkcommon/tasking/tasking_system_init.h"
#include "rkcommon/utility/getEnvVar.h"

#include <map>

namespace ospray {
namespace api {

static FactoryMap<Device> g_devicesMap;

// Helper functions ///////////////////////////////////////////////////////////

template <typename OSTREAM_T>
static inline void installStatusMsgFunc(Device &device, OSTREAM_T &stream)
{
  device.msg_fcn = [&](const char *msg) { stream << msg; };
}

template <typename OSTREAM_T>
static inline void installErrorMsgFunc(Device &device, OSTREAM_T &stream)
{
  device.error_fcn = [&](OSPError e, const char *msg) {
    stream << "OSPRAY ERROR [" << e << "]: " << msg << std::endl;
  };
}

// Device definitions /////////////////////////////////////////////////////////

memory::IntrusivePtr<Device> Device::current;
uint32_t Device::logLevel = OSP_LOG_NONE;

Device *Device::createDevice(const char *type)
{
  // NOTE(jda) - If a user is manually creating the device (i.e. not using
  //             ospInit() to do it), then we need to check if there's a
  //             valid library for core ospray in our main symbol lookup
  //             table.
  auto &repo = *LibraryRepository::getInstance();
  if (!repo.libraryExists("ospray"))
    repo.addDefaultLibrary();

  if (!repo.libraryExists("ospray_module_ispc") && type == std::string("cpu"))
    repo.add("ospray_module_ispc");

  return createInstanceHelper(type, g_devicesMap[type]);
}

void Device::registerType(const char *type, FactoryFcn<Device> f)
{
  g_devicesMap[type] = f;
}

void Device::commit()
{
  auto OSPRAY_DEBUG = utility::getEnvVar<int>("OSPRAY_DEBUG");
  debugMode = OSPRAY_DEBUG.value_or(getParam<bool>("debug", 0));

  auto OSPRAY_WARN = utility::getEnvVar<int>("OSPRAY_WARN_AS_ERROR");
  warningsAreErrors =
      OSPRAY_WARN.value_or(getParam<bool>("warnAsError", false));

  auto OSPRAY_TRACE_API = utility::getEnvVar<int>("OSPRAY_TRACE_API");
  bool traceAPI = OSPRAY_TRACE_API.value_or(getParam<bool>("traceApi", false));

  if (traceAPI && !apiTraceEnabled) {
    auto streamPtr = std::make_shared<std::ofstream>("ospray_api_trace.txt");

    trace_fcn = [=](const char *message) {
      auto &stream = *streamPtr;
      stream << message << std::endl;
    };
  } else if (!traceAPI) {
    trace_fcn = [](const char *) {};
  }

  apiTraceEnabled = traceAPI;

  logLevel = getParam<int>("logLevel", OSP_LOG_NONE);

  auto logLevelStr =
      utility::getEnvVar<std::string>("OSPRAY_LOG_LEVEL").value_or("");

  logLevel = logLevelFromString(logLevelStr).value_or(logLevel);

  auto OSPRAY_NUM_THREADS = utility::getEnvVar<int>("OSPRAY_NUM_THREADS");
  numThreads = OSPRAY_NUM_THREADS.value_or(getParam<int>("numThreads", -1));

  auto OSPRAY_LOG_OUTPUT = utility::getEnvVar<std::string>("OSPRAY_LOG_OUTPUT");

  auto dst = OSPRAY_LOG_OUTPUT.value_or(getParam<std::string>("logOutput", ""));

  if (dst == "cout")
    installStatusMsgFunc(*this, std::cout);
  else if (dst == "cerr")
    installStatusMsgFunc(*this, std::cerr);
  else if (dst == "none")
    msg_fcn = [](const char *) {};

  auto OSPRAY_ERROR_OUTPUT =
      utility::getEnvVar<std::string>("OSPRAY_ERROR_OUTPUT");

  dst = OSPRAY_ERROR_OUTPUT.value_or(getParam<std::string>("errorOutput", ""));

  if (dst == "cout")
    installErrorMsgFunc(*this, std::cout);
  else if (dst == "cerr")
    installErrorMsgFunc(*this, std::cerr);
  else if (dst == "none")
    error_fcn = [](OSPError, const char *) {};

  if (debugMode) {
    logLevel = OSP_LOG_DEBUG;
    numThreads = 1;
    installStatusMsgFunc(*this, std::cout);
    installErrorMsgFunc(*this, std::cerr);
  }

  threadAffinity = getParam<int>("setAffinity", AUTO_DETECT);

  auto OSPRAY_SET_AFFINITY = utility::getEnvVar<int>("OSPRAY_SET_AFFINITY");
  if (OSPRAY_SET_AFFINITY)
    threadAffinity = OSPRAY_SET_AFFINITY.value();

  tasking::initTaskingSystem(numThreads);

  committed = true;
}

bool Device::isCommitted()
{
  return committed;
}

bool deviceIsSet()
{
  return Device::current.ptr != nullptr;
}

Device &currentDevice()
{
  return *Device::current;
}

std::string generateEmbreeDeviceCfg(const Device &device)
{
  std::stringstream embreeConfig;

  if (device.debugMode)
    embreeConfig << " verbose=2";

  if (device.threadAffinity == api::Device::AFFINITIZE)
    embreeConfig << ",set_affinity=1";
  else if (device.threadAffinity == api::Device::DEAFFINITIZE)
    embreeConfig << ",set_affinity=0";

  return embreeConfig.str();
}

int64_t Device::getProperty(const OSPDeviceProperty prop)
{
  /* documented properties */
  switch (prop) {
  case OSP_DEVICE_VERSION:
    return 10000 * OSPRAY_VERSION_MAJOR + 100 * OSPRAY_VERSION_MINOR
        + OSPRAY_VERSION_PATCH;
  case OSP_DEVICE_VERSION_MAJOR:
    return OSPRAY_VERSION_MAJOR;
  case OSP_DEVICE_VERSION_MINOR:
    return OSPRAY_VERSION_MINOR;
  case OSP_DEVICE_VERSION_PATCH:
    return OSPRAY_VERSION_PATCH;
  case OSP_DEVICE_SO_VERSION:
    return OSPRAY_SOVERSION;
  default:
    handleError(OSP_INVALID_ARGUMENT, "unknown readable property");
    return 0;
  }
}

utility::Optional<int> Device::logLevelFromString(const std::string &str)
{
  utility::Optional<int> retval;

  if (str == "none")
    retval = OSP_LOG_NONE;
  else if (str == "debug")
    retval = OSP_LOG_DEBUG;
  else if (str == "info")
    retval = OSP_LOG_INFO;
  else if (str == "warning")
    retval = OSP_LOG_WARNING;
  else if (str == "error")
    retval = OSP_LOG_ERROR;

  return retval;
}

} // namespace api

OSPTYPEFOR_DEFINITION(api::Device *);

} // namespace ospray
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/c/experimental/gradients/math_grad.h"

#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/experimental/ops/array_ops.h"
#include "tensorflow/c/experimental/ops/math_ops.h"
#include "tensorflow/c/experimental/ops/nn_ops.h"

using std::vector;
using tensorflow::ops::Conj;
using tensorflow::ops::MatMul;
using tensorflow::ops::Mul;
using tensorflow::ops::Neg;
using tensorflow::ops::SqrtGrad;

namespace tensorflow {
namespace gradients {
namespace {

class AddGradientFunction : public GradientFunction {
 public:
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    grad_outputs->resize(2);
    // TODO(b/161805092): Support broadcasting.

    DCHECK(grad_inputs[0]);
    (*grad_outputs)[0] = grad_inputs[0];
    (*grad_outputs)[1] = grad_inputs[0];

    (*grad_outputs)[0]->Ref();
    (*grad_outputs)[1]->Ref();
    return Status::OK();
  }
  ~AddGradientFunction() override {}
};

class ExpGradientFunction : public GradientFunction {
 public:
  explicit ExpGradientFunction(AbstractTensorHandle* exp) : exp_(exp) {
    exp->Ref();
  }
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    vector<AbstractTensorHandle*> conj_outputs(1);
    std::string name = "Conj_Exp_Grad";
    TF_RETURN_IF_ERROR(Conj(ctx->ctx, {exp_.get()},
                            absl::MakeSpan(conj_outputs), name.c_str()));
    AbstractTensorHandlePtr conj_output_releaser(conj_outputs[0]);
    grad_outputs->resize(1);

    name = "Mul_Exp_Grad";
    TF_RETURN_IF_ERROR(Mul(ctx->ctx, {conj_outputs[0], grad_inputs[0]},
                           absl::MakeSpan(*grad_outputs), name.c_str()));
    return Status::OK();
  }
  ~ExpGradientFunction() override {}

 private:
  AbstractTensorHandlePtr exp_;
};

class SqrtGradientFunction : public GradientFunction {
 public:
  explicit SqrtGradientFunction(AbstractTensorHandle* sqrt) : sqrt_(sqrt) {
    sqrt->Ref();
  }
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    std::string name = "Sqrt_Grad";
    grad_outputs->resize(1);
    TF_RETURN_IF_ERROR(SqrtGrad(ctx->ctx, {sqrt_.get(), grad_inputs[0]},
                                absl::MakeSpan(*grad_outputs), name.c_str()));
    return Status::OK();
  }
  ~SqrtGradientFunction() override {}

 private:
  AbstractTensorHandlePtr sqrt_;
};

class MatMulGradientFunction : public GradientFunction {
 public:
  explicit MatMulGradientFunction(vector<AbstractTensorHandle*> f_inputs,
                                  AttrBuilder f_attrs)
      : forward_inputs(f_inputs), forward_attrs(f_attrs) {}

  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    /* Given upstream grad U and a matmul op A*B, the gradients are:
     *
     *    dA = U * B.T
     *    dB = A.T * U
     *
     *    where A.T means `transpose(A)`
     */
    AbstractTensorHandle* upstream_grad = grad_inputs[0];
    grad_outputs->resize(2);

    // Get transpose attrs
    bool t_a;
    TF_RETURN_IF_ERROR(forward_attrs.Get("transpose_a", &t_a));

    bool t_b;
    TF_RETURN_IF_ERROR(forward_attrs.Get("transpose_b", &t_b));

    // Conj each input
    vector<AbstractTensorHandle*> conj_outputs(1);
    std::string name = "Conj_A_MatMul_Grad";
    TF_RETURN_IF_ERROR(Conj(ctx->ctx, {forward_inputs[0]},
                            absl::MakeSpan(conj_outputs), name.c_str()));

    AbstractTensorHandle* A = conj_outputs[0];

    name = "Conj_B_MatMul_Grad";
    TF_RETURN_IF_ERROR(Conj(ctx->ctx, {forward_inputs[1]},
                            absl::MakeSpan(conj_outputs), name.c_str()));

    AbstractTensorHandle* B = conj_outputs[0];

    // Calc Grad
    vector<AbstractTensorHandle*> matmul_A_outputs(1);
    vector<AbstractTensorHandle*> matmul_B_outputs(1);
    std::string name_grad_A = "MatMul_Grad_A";
    std::string name_grad_B = "MatMul_Grad_B";
    if (!t_a && !t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, B},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ true));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {A, upstream_grad},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ false));
    } else if (!t_a && t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, B},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ false));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, A},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ false));

    } else if (t_a && !t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {B, upstream_grad},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ true));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {A, upstream_grad},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ false));
    } else {  // t_a && t_b
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {B, upstream_grad},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ true));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, A},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ true));
    }

    // Gradient for A
    (*grad_outputs)[0] = matmul_A_outputs[0];

    // Gradient for B
    (*grad_outputs)[1] = matmul_B_outputs[0];
    return Status::OK();
  }
  ~MatMulGradientFunction() override {}

 private:
  vector<AbstractTensorHandle*> forward_inputs;
  AttrBuilder forward_attrs;
};

class NegGradientFunction : public GradientFunction {
 public:
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    /* Given upstream grad U and a Neg op Y = -X, the gradients are:
     *
     *    dX =  -U
     *
     */

    grad_outputs->resize(1);
    std::string name = "Neg_Grad";
    TF_RETURN_IF_ERROR(ops::Neg(ctx->ctx, {grad_inputs[0]},
                                absl::MakeSpan(*grad_outputs), name.c_str()));
    return Status::OK();
  }
  ~NegGradientFunction() override {}
};

class SubGradientFunction : public GradientFunction {
 public:
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    /* Given upstream grad U and a Sub op A-B, the gradients are:
     *
     *    dA =  U
     *    dB = -U
     *
     */

    grad_outputs->resize(2);

    // Grad for A
    DCHECK(grad_inputs[0]);
    (*grad_outputs)[0] = grad_inputs[0];

    // Grad for B
    // negate the upstream grad
    std::vector<AbstractTensorHandle*> neg_outputs(1);
    std::string name = "Neg_Sub_Grad_B";
    TF_RETURN_IF_ERROR(ops::Neg(ctx->ctx, {grad_inputs[0]},
                                absl::MakeSpan(neg_outputs), name.c_str()));
    (*grad_outputs)[1] = neg_outputs[0];

    return Status::OK();
  }
  ~SubGradientFunction() override {}
};

}  // namespace

BackwardFunction* AddRegisterer(const ForwardOperation& op) {
  auto gradient_function = new AddGradientFunction;
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* ExpRegisterer(const ForwardOperation& op) {
  auto gradient_function = new ExpGradientFunction(op.outputs[0]);
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* MatMulRegisterer(const ForwardOperation& op) {
  auto gradient_function = new MatMulGradientFunction(op.inputs, op.attrs);
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* SqrtRegisterer(const ForwardOperation& op) {
  auto gradient_function = new SqrtGradientFunction(op.outputs[0]);
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* NegRegisterer(const ForwardOperation& op) {
  auto gradient_function = new NegGradientFunction;
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* SubRegisterer(const ForwardOperation& op) {
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto gradient_function = new SubGradientFunction;
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

}  // namespace gradients
}  // namespace tensorflow
<commit_msg>Call Ref on reference<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/c/experimental/gradients/math_grad.h"

#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/experimental/ops/array_ops.h"
#include "tensorflow/c/experimental/ops/math_ops.h"
#include "tensorflow/c/experimental/ops/nn_ops.h"

using std::vector;
using tensorflow::ops::Conj;
using tensorflow::ops::MatMul;
using tensorflow::ops::Mul;
using tensorflow::ops::Neg;
using tensorflow::ops::SqrtGrad;

namespace tensorflow {
namespace gradients {
namespace {

class AddGradientFunction : public GradientFunction {
 public:
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    grad_outputs->resize(2);
    // TODO(b/161805092): Support broadcasting.

    DCHECK(grad_inputs[0]);
    (*grad_outputs)[0] = grad_inputs[0];
    (*grad_outputs)[1] = grad_inputs[0];

    (*grad_outputs)[0]->Ref();
    (*grad_outputs)[1]->Ref();
    return Status::OK();
  }
  ~AddGradientFunction() override {}
};

class ExpGradientFunction : public GradientFunction {
 public:
  explicit ExpGradientFunction(AbstractTensorHandle* exp) : exp_(exp) {
    exp->Ref();
  }
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    vector<AbstractTensorHandle*> conj_outputs(1);
    std::string name = "Conj_Exp_Grad";
    TF_RETURN_IF_ERROR(Conj(ctx->ctx, {exp_.get()},
                            absl::MakeSpan(conj_outputs), name.c_str()));
    AbstractTensorHandlePtr conj_output_releaser(conj_outputs[0]);
    grad_outputs->resize(1);

    name = "Mul_Exp_Grad";
    TF_RETURN_IF_ERROR(Mul(ctx->ctx, {conj_outputs[0], grad_inputs[0]},
                           absl::MakeSpan(*grad_outputs), name.c_str()));
    return Status::OK();
  }
  ~ExpGradientFunction() override {}

 private:
  AbstractTensorHandlePtr exp_;
};

class SqrtGradientFunction : public GradientFunction {
 public:
  explicit SqrtGradientFunction(AbstractTensorHandle* sqrt) : sqrt_(sqrt) {
    sqrt->Ref();
  }
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    std::string name = "Sqrt_Grad";
    grad_outputs->resize(1);
    TF_RETURN_IF_ERROR(SqrtGrad(ctx->ctx, {sqrt_.get(), grad_inputs[0]},
                                absl::MakeSpan(*grad_outputs), name.c_str()));
    return Status::OK();
  }
  ~SqrtGradientFunction() override {}

 private:
  AbstractTensorHandlePtr sqrt_;
};

class MatMulGradientFunction : public GradientFunction {
 public:
  explicit MatMulGradientFunction(vector<AbstractTensorHandle*> f_inputs,
                                  AttrBuilder f_attrs)
      : forward_inputs(f_inputs), forward_attrs(f_attrs) {}

  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    /* Given upstream grad U and a matmul op A*B, the gradients are:
     *
     *    dA = U * B.T
     *    dB = A.T * U
     *
     *    where A.T means `transpose(A)`
     */
    AbstractTensorHandle* upstream_grad = grad_inputs[0];
    grad_outputs->resize(2);

    // Get transpose attrs
    bool t_a;
    TF_RETURN_IF_ERROR(forward_attrs.Get("transpose_a", &t_a));

    bool t_b;
    TF_RETURN_IF_ERROR(forward_attrs.Get("transpose_b", &t_b));

    // Conj each input
    vector<AbstractTensorHandle*> conj_outputs(1);
    std::string name = "Conj_A_MatMul_Grad";
    TF_RETURN_IF_ERROR(Conj(ctx->ctx, {forward_inputs[0]},
                            absl::MakeSpan(conj_outputs), name.c_str()));

    AbstractTensorHandle* A = conj_outputs[0];

    name = "Conj_B_MatMul_Grad";
    TF_RETURN_IF_ERROR(Conj(ctx->ctx, {forward_inputs[1]},
                            absl::MakeSpan(conj_outputs), name.c_str()));

    AbstractTensorHandle* B = conj_outputs[0];

    // Calc Grad
    vector<AbstractTensorHandle*> matmul_A_outputs(1);
    vector<AbstractTensorHandle*> matmul_B_outputs(1);
    std::string name_grad_A = "MatMul_Grad_A";
    std::string name_grad_B = "MatMul_Grad_B";
    if (!t_a && !t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, B},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ true));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {A, upstream_grad},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ false));
    } else if (!t_a && t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, B},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ false));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, A},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ false));

    } else if (t_a && !t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {B, upstream_grad},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ true));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {A, upstream_grad},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ false,
                                /*transpose_b = */ false));
    } else {  // t_a && t_b
      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {B, upstream_grad},
                                absl::MakeSpan(matmul_A_outputs),
                                name_grad_A.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ true));

      TF_RETURN_IF_ERROR(MatMul(ctx->ctx, {upstream_grad, A},
                                absl::MakeSpan(matmul_B_outputs),
                                name_grad_B.c_str(),
                                /*transpose_a = */ true,
                                /*transpose_b = */ true));
    }

    // Gradient for A
    (*grad_outputs)[0] = matmul_A_outputs[0];

    // Gradient for B
    (*grad_outputs)[1] = matmul_B_outputs[0];
    return Status::OK();
  }
  ~MatMulGradientFunction() override {}

 private:
  vector<AbstractTensorHandle*> forward_inputs;
  AttrBuilder forward_attrs;
};

class NegGradientFunction : public GradientFunction {
 public:
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    /* Given upstream grad U and a Neg op Y = -X, the gradients are:
     *
     *    dX =  -U
     *
     */

    grad_outputs->resize(1);
    std::string name = "Neg_Grad";
    TF_RETURN_IF_ERROR(ops::Neg(ctx->ctx, {grad_inputs[0]},
                                absl::MakeSpan(*grad_outputs), name.c_str()));
    return Status::OK();
  }
  ~NegGradientFunction() override {}
};

class SubGradientFunction : public GradientFunction {
 public:
  Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
                 vector<AbstractTensorHandle*>* grad_outputs) override {
    /* Given upstream grad U and a Sub op A-B, the gradients are:
     *
     *    dA =  U
     *    dB = -U
     *
     */

    grad_outputs->resize(2);

    // Grad for A
    DCHECK(grad_inputs[0]);
    (*grad_outputs)[0] = grad_inputs[0];
    (*grad_outputs)[0]->Ref();

    // Grad for B
    // negate the upstream grad
    std::vector<AbstractTensorHandle*> neg_outputs(1);
    std::string name = "Neg_Sub_Grad_B";
    TF_RETURN_IF_ERROR(ops::Neg(ctx->ctx, {grad_inputs[0]},
                                absl::MakeSpan(neg_outputs), name.c_str()));
    (*grad_outputs)[1] = neg_outputs[0];

    return Status::OK();
  }
  ~SubGradientFunction() override {}
};

}  // namespace

BackwardFunction* AddRegisterer(const ForwardOperation& op) {
  auto gradient_function = new AddGradientFunction;
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* ExpRegisterer(const ForwardOperation& op) {
  auto gradient_function = new ExpGradientFunction(op.outputs[0]);
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* MatMulRegisterer(const ForwardOperation& op) {
  auto gradient_function = new MatMulGradientFunction(op.inputs, op.attrs);
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* SqrtRegisterer(const ForwardOperation& op) {
  auto gradient_function = new SqrtGradientFunction(op.outputs[0]);
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* NegRegisterer(const ForwardOperation& op) {
  auto gradient_function = new NegGradientFunction;
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

BackwardFunction* SubRegisterer(const ForwardOperation& op) {
  // For ops with a single output, the gradient function is not called if there
  // is no incoming gradient. So we do not need to worry about creating zeros
  // grads in this case.
  auto gradient_function = new SubGradientFunction;
  auto default_gradients = new PassThroughDefaultGradients(op);
  return new BackwardFunction(gradient_function, default_gradients);
}

}  // namespace gradients
}  // namespace tensorflow
<|endoftext|>


<commit_before>// RUN: %python %S/check_clang_tidy.py %s modernize-use-default %t

class A {
public:
  A();
  ~A();
};

A::A() {}
// CHECK-MESSAGES: :[[@LINE-1]]:1: warning: use '= default' to define a trivial default constructor [modernize-use-default]
// CHECK-FIXES: A::A() = default;
A::~A() {}
// CHECK-MESSAGES: :[[@LINE-1]]:1: warning: use '= default' to define a trivial destructor [modernize-use-default]
// CHECK-FIXES: A::~A() = default;

// Inline definitions.
class B {
public:
  B() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: B() = default;
  ~B() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~B() = default;
};

void f();

class C {
public:
  // Non-empty constructor body.
  C() { f(); }
  // Non-empty destructor body.
  ~C() { f(); }
};

class D {
public:
  // Constructor with initializer.
  D() : Field(5) {}
  // Constructor with arguments.
  D(int Arg1, int Arg2) {}
  int Field;
};

// Private constructor/destructor.
class E {
  E() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: E() = default;
  ~E() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~E() = default;
};

// struct.
struct F {
  F() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: F() = default;
  ~F() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: F() = default;
};

// Deleted constructor/destructor.
class G {
public:
  G() = delete;
  ~G() = delete;
};

// Do not remove other keywords.
class H {
public:
  explicit H() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: explicit H() = default;
  virtual ~H() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: virtual ~H() = default;
};

// Nested class.
struct I {
  struct II {
    II() {}
    // CHECK-MESSAGES: :[[@LINE-1]]:5: warning: use '= default'
    // CHECK-FIXES: II() = default;
    ~II() {}
    // CHECK-MESSAGES: :[[@LINE-1]]:5: warning: use '= default'
    // CHECK-FIXES: ~II() = default;
  };
  int Int;
};

// Class template.
template <class T>
class J {
public:
  J() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: J() = default;
  ~J() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~J() = default;
};

// Non user-provided constructor/destructor.
struct K {
  int Int;
};
void g() {
  K *PtrK = new K();
  PtrK->~K();
  delete PtrK;
}

// Already using default.
struct L {
  L() = default;
  ~L() = default;
};
struct M {
  M();
  ~M();
};
M::M() = default;
M::~M() = default;

// Delegating constructor and overriden destructor.
struct N : H {
  N() : H() {}
  ~N() override {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~N() override = default;
};
<commit_msg>clang-tools-extra/test/clang-tidy/modernize-use-default.cpp: Appease MS mode.<commit_after>// RUN: %python %S/check_clang_tidy.py %s modernize-use-default %t -- -std=c++11 -fno-delayed-template-parsing

class A {
public:
  A();
  ~A();
};

A::A() {}
// CHECK-MESSAGES: :[[@LINE-1]]:1: warning: use '= default' to define a trivial default constructor [modernize-use-default]
// CHECK-FIXES: A::A() = default;
A::~A() {}
// CHECK-MESSAGES: :[[@LINE-1]]:1: warning: use '= default' to define a trivial destructor [modernize-use-default]
// CHECK-FIXES: A::~A() = default;

// Inline definitions.
class B {
public:
  B() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: B() = default;
  ~B() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~B() = default;
};

void f();

class C {
public:
  // Non-empty constructor body.
  C() { f(); }
  // Non-empty destructor body.
  ~C() { f(); }
};

class D {
public:
  // Constructor with initializer.
  D() : Field(5) {}
  // Constructor with arguments.
  D(int Arg1, int Arg2) {}
  int Field;
};

// Private constructor/destructor.
class E {
  E() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: E() = default;
  ~E() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~E() = default;
};

// struct.
struct F {
  F() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: F() = default;
  ~F() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: F() = default;
};

// Deleted constructor/destructor.
class G {
public:
  G() = delete;
  ~G() = delete;
};

// Do not remove other keywords.
class H {
public:
  explicit H() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: explicit H() = default;
  virtual ~H() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: virtual ~H() = default;
};

// Nested class.
struct I {
  struct II {
    II() {}
    // CHECK-MESSAGES: :[[@LINE-1]]:5: warning: use '= default'
    // CHECK-FIXES: II() = default;
    ~II() {}
    // CHECK-MESSAGES: :[[@LINE-1]]:5: warning: use '= default'
    // CHECK-FIXES: ~II() = default;
  };
  int Int;
};

// Class template.
template <class T>
class J {
public:
  J() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: J() = default;
  ~J() {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~J() = default;
};

// Non user-provided constructor/destructor.
struct K {
  int Int;
};
void g() {
  K *PtrK = new K();
  PtrK->~K();
  delete PtrK;
}

// Already using default.
struct L {
  L() = default;
  ~L() = default;
};
struct M {
  M();
  ~M();
};
M::M() = default;
M::~M() = default;

// Delegating constructor and overriden destructor.
struct N : H {
  N() : H() {}
  ~N() override {}
  // CHECK-MESSAGES: :[[@LINE-1]]:3: warning: use '= default'
  // CHECK-FIXES: ~N() override = default;
};
<|endoftext|>


<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Native XLA implementations of simple binary Ops

#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/compiler/xla/client/client_library.h"
#include "tensorflow/compiler/xla/client/lib/constants.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/core/framework/kernel_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"

namespace tensorflow {
namespace {

// A subclass of a XlaBinaryOp must build the computation that
// describes the (tensor,tensor)->tensor function to apply to each element of
// the input.
#define XLA_MAKE_BINARY(NAME, HLO)                                       \
  class NAME##Op : public XlaBinaryOp {                                  \
   public:                                                               \
    explicit NAME##Op(OpKernelConstruction* ctx) : XlaBinaryOp(ctx) {}   \
    xla::XlaOp Computation(                                              \
        XlaOpKernelContext* ctx, const xla::XlaOp& lhs,                  \
        const absl::Span<const int64>& lhs_shape, const xla::XlaOp& rhs, \
        const absl::Span<const int64>& rhs_shape,                        \
        const BCast& broadcast_helper,                                   \
        const std::vector<int64>& extend_dimensions) override {          \
      xla::XlaBuilder* b = ctx->builder();                               \
      (void)b;                                                           \
      return HLO;                                                        \
    }                                                                    \
  };                                                                     \
  REGISTER_XLA_OP(Name(#NAME), NAME##Op)

XLA_MAKE_BINARY(Add, xla::Add(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Sub, xla::Sub(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Mul, xla::Mul(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Div, xla::Div(lhs, rhs, extend_dimensions));

XLA_MAKE_BINARY(Atan2, xla::Atan2(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions));

// Implementation of DivNoNan. Pseudo-code:
// if (y == 0) {
//   return 0
// } else {
//   return x / y;
// }
static xla::XlaOp DivNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto y_equals_0 = xla::Eq(y, zero);
  auto zeros = xla::ZerosLike(x);
  auto result = xla::Select(y_equals_0, zeros, xla::Div(x, y));
  return result;
}
XLA_MAKE_BINARY(DivNoNan,
                DivNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper));

// Implementation of FloorDiv. Pseudo-code:
// if ((x < 0) != (y < 0)) {
//   T abs_x = std::abs(x);
//   T abs_y = std::abs(y);
//   return -(abs_x + abs_y - 1) / abs_y;
// } else {
//   return x / y;
// }
static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  if (DataTypeIsUnsigned(dtype)) {
    return xla::Div(x, y);
  }
  auto zero = XlaHelpers::Zero(b, dtype);
  auto one = XlaHelpers::One(b, dtype);
  auto different_sign = xla::Ne(xla::Lt(x, zero), xla::Lt(y, zero));
  auto abs_x = xla::Abs(x);
  auto abs_y = xla::Abs(y);
  auto t = xla::Neg(xla::Sub(xla::Add(abs_x, abs_y), one));
  auto result = xla::Select(different_sign, xla::Div(t, abs_y), xla::Div(x, y));
  if (DataTypeIsFloating(dtype)) {
    result = xla::Floor(result);
  }
  return result;
}
XLA_MAKE_BINARY(FloorDiv,
                FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));

static xla::XlaOp XlogyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                            xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto is_zero = xla::Eq(x, zero);
  return xla::Select(is_zero, zero, xla::Mul(x, xla::Log(y)));
}
XLA_MAKE_BINARY(Xlogy, XlogyImpl(b, input_type(0), lhs, rhs, broadcast_helper));

static xla::XlaOp XdivyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                            xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto is_zero = xla::Eq(x, zero);
  return xla::Select(is_zero, zero, xla::Div(x, y));
}
XLA_MAKE_BINARY(Xdivy, XdivyImpl(b, input_type(0), lhs, rhs, broadcast_helper));

// Implementation of FloorMod. Pseudo-code:
// T trunc_mod = std::fmod(x, y);
// return (x < T(0)) == (y < T(0)) ? trunc_mod : std::fmod(trunc_mod + y, y);
static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto same_sign = xla::Eq(xla::Lt(x, zero), xla::Lt(y, zero));
  auto trunc_mod = xla::Rem(x, y);
  return xla::Select(same_sign, trunc_mod, xla::Rem(xla::Add(trunc_mod, y), y));
}
XLA_MAKE_BINARY(FloorMod,
                FloorModImpl(b, input_type(0), lhs, rhs, broadcast_helper));

XLA_MAKE_BINARY(BitwiseAnd, xla::And(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(BitwiseOr, xla::Or(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(BitwiseXor, xla::Xor(lhs, rhs, extend_dimensions));

XLA_MAKE_BINARY(LeftShift, xla::ShiftLeft(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(RightShift,
                (DataTypeIsUnsigned(ctx->input_type(0))
                     ? xla::ShiftRightLogical(lhs, rhs, extend_dimensions)
                     : xla::ShiftRightArithmetic(lhs, rhs, extend_dimensions)));

XLA_MAKE_BINARY(LogicalAnd, xla::And(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(LogicalOr, xla::Or(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Mod, xla::Rem(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Maximum, xla::Max(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Minimum, xla::Min(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(RealDiv, xla::Div(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(ReciprocalGrad, xla::Neg(xla::Mul(rhs, xla::Mul(lhs, lhs))));
XLA_MAKE_BINARY(
    RsqrtGrad,
    xla::Mul(xla::Pow(lhs, XlaHelpers::IntegerLiteral(b, input_type(0), 3)),
             xla::Div(rhs, XlaHelpers::IntegerLiteral(b, input_type(0), -2)),
             extend_dimensions));
XLA_MAKE_BINARY(
    SqrtGrad,
    xla::Div(xla::Mul(rhs, XlaHelpers::FloatLiteral(b, input_type(0), 0.5)),
             lhs, extend_dimensions));

static xla::XlaOp Square(xla::XlaBuilder* builder, const xla::XlaOp& x) {
  return xla::Mul(x, x);
}

XLA_MAKE_BINARY(SquaredDifference,
                Square(b, xla::Sub(lhs, rhs, extend_dimensions)));

XLA_MAKE_BINARY(TruncateDiv, xla::Div(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions));

// Comparison ops
XLA_MAKE_BINARY(Equal, xla::Eq(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(NotEqual, xla::Ne(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Greater, xla::Gt(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(GreaterEqual, xla::Ge(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Less, xla::Lt(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(LessEqual, xla::Le(lhs, rhs, extend_dimensions));

// Non-linear ops
XLA_MAKE_BINARY(SigmoidGrad,
                xla::Mul(xla::Mul(rhs, lhs),
                         xla::Sub(XlaHelpers::One(b, input_type(0)), lhs)));

XLA_MAKE_BINARY(SoftplusGrad,
                xla::Div(lhs, xla::Add(xla::Exp(xla::Neg(rhs)),
                                       XlaHelpers::One(b, input_type(1)))));

// softsigngrad(gradients, features) = gradients / (1 + abs(features)) ** 2
XLA_MAKE_BINARY(SoftsignGrad,
                xla::Div(lhs,
                         Square(b, xla::Add(XlaHelpers::One(b, input_type(0)),
                                            xla::Abs(rhs)))));

XLA_MAKE_BINARY(TanhGrad,
                xla::Mul(rhs, xla::Sub(XlaHelpers::One(b, input_type(0)),
                                       xla::Mul(lhs, lhs))));

XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions));

#undef XLA_MAKE_BINARY

class ApproximateEqualOp : public XlaOpKernel {
 public:
  explicit ApproximateEqualOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("tolerance", &tolerance_));
  }

  // Computes the max of the scalar input x and 0.
  void Compile(XlaOpKernelContext* ctx) override {
    xla::XlaBuilder* b = ctx->builder();
    auto abs = xla::Abs(xla::Sub(ctx->Input(0), ctx->Input(1)));
    auto abs_shape = b->GetShape(abs);
    OP_REQUIRES_OK(ctx, abs_shape.status());
    auto abs_type = abs_shape.ValueOrDie().element_type();
    auto result =
        xla::Lt(abs, xla::ConvertElementType(
                         xla::ConstantR0<float>(b, tolerance_), abs_type));
    ctx->SetOutput(0, result);
  }

 private:
  float tolerance_;
};
REGISTER_XLA_OP(Name("ApproximateEqual"), ApproximateEqualOp);

}  // namespace
}  // namespace tensorflow
<commit_msg>Avoid unused variable warnings from XLA_MAKE_BINARY<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Native XLA implementations of simple binary Ops

#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/compiler/xla/client/client_library.h"
#include "tensorflow/compiler/xla/client/lib/constants.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/core/framework/kernel_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"

namespace tensorflow {
namespace {

// A subclass of a XlaBinaryOp must build the computation that
// describes the (tensor,tensor)->tensor function to apply to each element of
// the input.
#define XLA_MAKE_BINARY(NAME, HLO)                                       \
  class NAME##Op : public XlaBinaryOp {                                  \
   public:                                                               \
    explicit NAME##Op(OpKernelConstruction* ctx) : XlaBinaryOp(ctx) {}   \
    xla::XlaOp Computation(                                              \
        XlaOpKernelContext* ctx, const xla::XlaOp& lhs,                  \
        const absl::Span<const int64>& lhs_shape, const xla::XlaOp& rhs, \
        const absl::Span<const int64>& rhs_shape,                        \
        const BCast& broadcast_helper,                                   \
        const std::vector<int64>& extend_dimensions) override {          \
      xla::XlaBuilder* b = ctx->builder();                               \
      (void)b;                                                           \
      (void)lhs_shape;                                                   \
      (void)rhs_shape;                                                   \
      (void)extend_dimensions;                                           \
      return HLO;                                                        \
    }                                                                    \
  };                                                                     \
  REGISTER_XLA_OP(Name(#NAME), NAME##Op)

XLA_MAKE_BINARY(Add, xla::Add(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Sub, xla::Sub(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Mul, xla::Mul(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Div, xla::Div(lhs, rhs, extend_dimensions));

XLA_MAKE_BINARY(Atan2, xla::Atan2(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions));

// Implementation of DivNoNan. Pseudo-code:
// if (y == 0) {
//   return 0
// } else {
//   return x / y;
// }
static xla::XlaOp DivNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto y_equals_0 = xla::Eq(y, zero);
  auto zeros = xla::ZerosLike(x);
  auto result = xla::Select(y_equals_0, zeros, xla::Div(x, y));
  return result;
}
XLA_MAKE_BINARY(DivNoNan,
                DivNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper));

// Implementation of FloorDiv. Pseudo-code:
// if ((x < 0) != (y < 0)) {
//   T abs_x = std::abs(x);
//   T abs_y = std::abs(y);
//   return -(abs_x + abs_y - 1) / abs_y;
// } else {
//   return x / y;
// }
static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  if (DataTypeIsUnsigned(dtype)) {
    return xla::Div(x, y);
  }
  auto zero = XlaHelpers::Zero(b, dtype);
  auto one = XlaHelpers::One(b, dtype);
  auto different_sign = xla::Ne(xla::Lt(x, zero), xla::Lt(y, zero));
  auto abs_x = xla::Abs(x);
  auto abs_y = xla::Abs(y);
  auto t = xla::Neg(xla::Sub(xla::Add(abs_x, abs_y), one));
  auto result = xla::Select(different_sign, xla::Div(t, abs_y), xla::Div(x, y));
  if (DataTypeIsFloating(dtype)) {
    result = xla::Floor(result);
  }
  return result;
}
XLA_MAKE_BINARY(FloorDiv,
                FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));

static xla::XlaOp XlogyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                            xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto is_zero = xla::Eq(x, zero);
  return xla::Select(is_zero, zero, xla::Mul(x, xla::Log(y)));
}
XLA_MAKE_BINARY(Xlogy, XlogyImpl(b, input_type(0), lhs, rhs, broadcast_helper));

static xla::XlaOp XdivyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                            xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto is_zero = xla::Eq(x, zero);
  return xla::Select(is_zero, zero, xla::Div(x, y));
}
XLA_MAKE_BINARY(Xdivy, XdivyImpl(b, input_type(0), lhs, rhs, broadcast_helper));

// Implementation of FloorMod. Pseudo-code:
// T trunc_mod = std::fmod(x, y);
// return (x < T(0)) == (y < T(0)) ? trunc_mod : std::fmod(trunc_mod + y, y);
static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto same_sign = xla::Eq(xla::Lt(x, zero), xla::Lt(y, zero));
  auto trunc_mod = xla::Rem(x, y);
  return xla::Select(same_sign, trunc_mod, xla::Rem(xla::Add(trunc_mod, y), y));
}
XLA_MAKE_BINARY(FloorMod,
                FloorModImpl(b, input_type(0), lhs, rhs, broadcast_helper));

XLA_MAKE_BINARY(BitwiseAnd, xla::And(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(BitwiseOr, xla::Or(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(BitwiseXor, xla::Xor(lhs, rhs, extend_dimensions));

XLA_MAKE_BINARY(LeftShift, xla::ShiftLeft(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(RightShift,
                (DataTypeIsUnsigned(ctx->input_type(0))
                     ? xla::ShiftRightLogical(lhs, rhs, extend_dimensions)
                     : xla::ShiftRightArithmetic(lhs, rhs, extend_dimensions)));

XLA_MAKE_BINARY(LogicalAnd, xla::And(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(LogicalOr, xla::Or(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Mod, xla::Rem(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Maximum, xla::Max(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Minimum, xla::Min(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(RealDiv, xla::Div(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(ReciprocalGrad, xla::Neg(xla::Mul(rhs, xla::Mul(lhs, lhs))));
XLA_MAKE_BINARY(
    RsqrtGrad,
    xla::Mul(xla::Pow(lhs, XlaHelpers::IntegerLiteral(b, input_type(0), 3)),
             xla::Div(rhs, XlaHelpers::IntegerLiteral(b, input_type(0), -2)),
             extend_dimensions));
XLA_MAKE_BINARY(
    SqrtGrad,
    xla::Div(xla::Mul(rhs, XlaHelpers::FloatLiteral(b, input_type(0), 0.5)),
             lhs, extend_dimensions));

static xla::XlaOp Square(xla::XlaBuilder* builder, const xla::XlaOp& x) {
  return xla::Mul(x, x);
}

XLA_MAKE_BINARY(SquaredDifference,
                Square(b, xla::Sub(lhs, rhs, extend_dimensions)));

XLA_MAKE_BINARY(TruncateDiv, xla::Div(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions));

// Comparison ops
XLA_MAKE_BINARY(Equal, xla::Eq(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(NotEqual, xla::Ne(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Greater, xla::Gt(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(GreaterEqual, xla::Ge(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(Less, xla::Lt(lhs, rhs, extend_dimensions));
XLA_MAKE_BINARY(LessEqual, xla::Le(lhs, rhs, extend_dimensions));

// Non-linear ops
XLA_MAKE_BINARY(SigmoidGrad,
                xla::Mul(xla::Mul(rhs, lhs),
                         xla::Sub(XlaHelpers::One(b, input_type(0)), lhs)));

XLA_MAKE_BINARY(SoftplusGrad,
                xla::Div(lhs, xla::Add(xla::Exp(xla::Neg(rhs)),
                                       XlaHelpers::One(b, input_type(1)))));

// softsigngrad(gradients, features) = gradients / (1 + abs(features)) ** 2
XLA_MAKE_BINARY(SoftsignGrad,
                xla::Div(lhs,
                         Square(b, xla::Add(XlaHelpers::One(b, input_type(0)),
                                            xla::Abs(rhs)))));

XLA_MAKE_BINARY(TanhGrad,
                xla::Mul(rhs, xla::Sub(XlaHelpers::One(b, input_type(0)),
                                       xla::Mul(lhs, lhs))));

XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions));

#undef XLA_MAKE_BINARY

class ApproximateEqualOp : public XlaOpKernel {
 public:
  explicit ApproximateEqualOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("tolerance", &tolerance_));
  }

  // Computes the max of the scalar input x and 0.
  void Compile(XlaOpKernelContext* ctx) override {
    xla::XlaBuilder* b = ctx->builder();
    auto abs = xla::Abs(xla::Sub(ctx->Input(0), ctx->Input(1)));
    auto abs_shape = b->GetShape(abs);
    OP_REQUIRES_OK(ctx, abs_shape.status());
    auto abs_type = abs_shape.ValueOrDie().element_type();
    auto result =
        xla::Lt(abs, xla::ConvertElementType(
                         xla::ConstantR0<float>(b, tolerance_), abs_type));
    ctx->SetOutput(0, result);
  }

 private:
  float tolerance_;
};
REGISTER_XLA_OP(Name("ApproximateEqual"), ApproximateEqualOp);

}  // namespace
}  // namespace tensorflow
<|endoftext|>


<commit_before>/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/bef_thunk.h"

#include "mlir-hlo/Dialect/mhlo/IR/lhlo_gpu_ops.h"
#include "tensorflow/core/platform/errors.h"

#if BEF_THUNKS
#include "tfrt/bef/bef_buffer.h"
#include "tfrt/bef_converter/mlir_to_bef_translate.h"
#include "tfrt/core_runtime/core_runtime.h"
#include "tfrt/gpu/gpu_types.h"
#include "tfrt/host_context/async_value_ref.h"
#include "tfrt/host_context/chain.h"
#include "tfrt/host_context/execution_context.h"
#include "tfrt/host_context/host_context.h"
#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h"
#include "tensorflow/stream_executor/cuda/cuda_driver.h"
#include "tensorflow/stream_executor/device_memory.h"
#include "tensorflow/stream_executor/gpu/gpu_executor.h"
#include "tensorflow/stream_executor/gpu/gpu_stream.h"
#endif  // BEF_THUNKS

namespace xla {
namespace gpu {

namespace {

#if BEF_THUNKS
StatusOr<std::unique_ptr<tfrt::gpu::Program>> ConvertToGpuProgram(
    mlir::Operation* op, tfrt::HostContext* host) {
  mlir::OwningModuleRef module =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(op->getContext()));
  const std::string func_name = "main";
  // TODO(hanbinyoon): Merge with the async lhlo->tfrt_gpu lowering pipeline.
  if (tensorflow::LhloGpuOpToTfrtCudaModule(op, module.get(), func_name)
          .failed()) {
    return tensorflow::errors::Internal(
        "Failed to lower lmhlo_gpu op to tfrt_gpu dialect.");
  }

  std::string bef;
  llvm::raw_string_ostream bef_ostream(bef);
  if (tfrt::MLIRToBEFTranslate(*module, bef_ostream).failed()) {
    return tensorflow::errors::Internal("Failed to translate MLIR to BEF.");
  }

  auto buffer = tfrt::BefBuffer(bef.data(), bef.data() + bef.size());
  return absl::make_unique<tfrt::gpu::Program>(std::move(buffer), func_name,
                                               host);
}
#endif  // BEF_THUNKS

}  // namespace

bool BefThunk::SupportsOp(mlir::Operation* op) {
#if !BEF_THUNKS
  return false;
#else
  return GetThunkKind(op).ok();
#endif  // BEF_THUNKS
}

StatusOr<std::unique_ptr<BefThunk>> BefThunk::Create(
    ThunkInfo thunk_info, mlir::Operation* op,
    std::vector<BufferAllocation::Slice> inputs,
    std::vector<BufferAllocation::Slice> outputs) {
#if !BEF_THUNKS
  LOG(FATAL) << "BefThunk is disabled.";
#else
  TF_ASSIGN_OR_RETURN(auto kind, GetThunkKind(op));
  auto thunk = absl::WrapUnique(
      new BefThunk(thunk_info, kind, std::move(inputs), std::move(outputs)));
  TF_ASSIGN_OR_RETURN(
      thunk->gpu_program_,
      ConvertToGpuProgram(op, runtime().core_runtime()->GetHostContext()));
  return thunk;
#endif  // BEF_THUNKS
}

BefThunk::BefThunk(ThunkInfo thunk_info, Thunk::Kind kind,
                   std::vector<BufferAllocation::Slice> inputs,
                   std::vector<BufferAllocation::Slice> outputs)
    : Thunk(kind, thunk_info)
#if BEF_THUNKS
      ,
      inputs_(std::move(inputs)),
      outputs_(std::move(outputs))
#endif  // BEF_THUNKS
{
}

Status BefThunk::ExecuteOnStream(const ExecuteParams& params) {
#if !BEF_THUNKS
  LOG(FATAL) << "BefThunk is disabled.";
#else
  VLOG(2) << "Executing BEF thunk.";

  auto* host = runtime().core_runtime()->GetHostContext();
  auto gpu_system = tfrt::gpu::System::Instantiate(host);
  host->Await({gpu_system.CopyRCRef()});

  // Use the GPU resources specified in ExecuteParams (initialized by the
  // StreamExecutor).
  auto se_gpu_executor = down_cast<stream_executor::gpu::GpuExecutor*>(
      params.stream->parent()->implementation());
  auto se_gpu_stream = down_cast<stream_executor::gpu::GpuStream*>(
      params.stream->implementation());
  auto context =
      tfrt::gpu::wrapper::Context(se_gpu_executor->gpu_context()->context());
  auto stream = tfrt::gpu::wrapper::Stream(se_gpu_stream->gpu_stream());
  tfrt::gpu::BorrowedGpuStream gpu_stream(host, context, stream);

  // Prepare arguments for BEF execution.
  auto get_async_value_ref = [&](const BufferAllocation::Slice& slice) {
    se::DeviceMemoryBase data =
        params.buffer_allocations->GetDeviceAddress(slice);
    tfrt::gpu::wrapper::Pointer<void> pointer(
        data.opaque(), tfrt::gpu::wrapper::Platform::CUDA);
    auto buffer = tfrt::gpu::GpuBuffer::Borrow(pointer, data.size());
    return tfrt::MakeAvailableAsyncValueRef<tfrt::gpu::GpuBuffer>(
        host, std::move(buffer));
  };

  llvm::SmallVector<tfrt::AsyncValueRef<tfrt::gpu::GpuBuffer>, 4> inputs;
  llvm::SmallVector<tfrt::AsyncValueRef<tfrt::gpu::GpuBuffer>, 4> outputs;
  for (auto input : inputs_) {
    inputs.push_back(get_async_value_ref(input));
  }
  for (auto output : outputs_) {
    outputs.push_back(get_async_value_ref(output));
  }

  tfrt::AsyncValueRef<tfrt::Chain> chain = tfrt::GetReadyChain(host);
  chain = gpu_system->Execute(exec_ctx(), *gpu_program_, gpu_stream, inputs,
                              outputs, std::move(chain));

  // TODO(hanbinyoon): BorrowedGpuStream releases context/stream information
  // upon destruction. Change it to not use wrapper::OwningContext and
  // wrapper::OwningStream.
  // TODO(b/184696034): Remove this. Right now we need this to ensure kernels
  // in bef function have been dispatched to stream.
  host->Quiesce();
  return Status::OK();
#endif  // BEF_THUNKS
}

StatusOr<Thunk::Kind> BefThunk::GetThunkKind(mlir::Operation* op) {
  if (mlir::isa<mlir::lmhlo_gpu::GEMMOp, mlir::lmhlo_gpu::GEMM_BiasOp>(op)) {
    return Thunk::Kind::kGemm;
  }
  return tensorflow::errors::Unimplemented(
      "Operation is not supported by BefThunk.");
}

#if BEF_THUNKS
tensorflow::tfrt_stub::Runtime& BefThunk::runtime() {
  static auto runtime = tensorflow::tfrt_stub::Runtime::Create().release();
  return *runtime;
}

tfrt::ExecutionContext& BefThunk::exec_ctx() {
  static tfrt::ExecutionContext* exec_ctx = [] {
    // Create request context and prepare deadline tracker.
    tfrt::RequestContextBuilder request_context_builder(
        runtime().core_runtime()->GetHostContext(),
        /*resource_context=*/nullptr);
    tensorflow::thread::ThreadPoolInterface* intra_op_threadpool = nullptr;
    DCHECK(
        runtime()
            .work_queue()
            ->InitializeRequest(&request_context_builder, &intra_op_threadpool)
            .ok());
    auto req_ctx = std::move(request_context_builder).build();
    DCHECK(req_ctx);
    return new tfrt::ExecutionContext(std::move(*req_ctx));
  }();
  return *exec_ctx;
}
#endif  // BEF_THUNKS

}  // namespace gpu
}  // namespace xla
<commit_msg>Remove GpuBuffer::Borrow(). Use GpuOneShotAllocator<void> instead.<commit_after>/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/bef_thunk.h"

#include "mlir-hlo/Dialect/mhlo/IR/lhlo_gpu_ops.h"
#include "tensorflow/core/platform/errors.h"

#if BEF_THUNKS
#include "tfrt/bef/bef_buffer.h"
#include "tfrt/bef_converter/mlir_to_bef_translate.h"
#include "tfrt/core_runtime/core_runtime.h"
#include "tfrt/gpu/gpu_types.h"
#include "tfrt/host_context/async_value_ref.h"
#include "tfrt/host_context/chain.h"
#include "tfrt/host_context/execution_context.h"
#include "tfrt/host_context/host_context.h"
#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h"
#include "tensorflow/stream_executor/cuda/cuda_driver.h"
#include "tensorflow/stream_executor/device_memory.h"
#include "tensorflow/stream_executor/gpu/gpu_executor.h"
#include "tensorflow/stream_executor/gpu/gpu_stream.h"
#endif  // BEF_THUNKS

namespace xla {
namespace gpu {

namespace {

#if BEF_THUNKS
StatusOr<std::unique_ptr<tfrt::gpu::Program>> ConvertToGpuProgram(
    mlir::Operation* op, tfrt::HostContext* host) {
  mlir::OwningModuleRef module =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(op->getContext()));
  const std::string func_name = "main";
  // TODO(hanbinyoon): Merge with the async lhlo->tfrt_gpu lowering pipeline.
  if (tensorflow::LhloGpuOpToTfrtCudaModule(op, module.get(), func_name)
          .failed()) {
    return tensorflow::errors::Internal(
        "Failed to lower lmhlo_gpu op to tfrt_gpu dialect.");
  }

  std::string bef;
  llvm::raw_string_ostream bef_ostream(bef);
  if (tfrt::MLIRToBEFTranslate(*module, bef_ostream).failed()) {
    return tensorflow::errors::Internal("Failed to translate MLIR to BEF.");
  }

  auto buffer = tfrt::BefBuffer(bef.data(), bef.data() + bef.size());
  return absl::make_unique<tfrt::gpu::Program>(std::move(buffer), func_name,
                                               host);
}
#endif  // BEF_THUNKS

}  // namespace

bool BefThunk::SupportsOp(mlir::Operation* op) {
#if !BEF_THUNKS
  return false;
#else
  return GetThunkKind(op).ok();
#endif  // BEF_THUNKS
}

StatusOr<std::unique_ptr<BefThunk>> BefThunk::Create(
    ThunkInfo thunk_info, mlir::Operation* op,
    std::vector<BufferAllocation::Slice> inputs,
    std::vector<BufferAllocation::Slice> outputs) {
#if !BEF_THUNKS
  LOG(FATAL) << "BefThunk is disabled.";
#else
  TF_ASSIGN_OR_RETURN(auto kind, GetThunkKind(op));
  auto thunk = absl::WrapUnique(
      new BefThunk(thunk_info, kind, std::move(inputs), std::move(outputs)));
  TF_ASSIGN_OR_RETURN(
      thunk->gpu_program_,
      ConvertToGpuProgram(op, runtime().core_runtime()->GetHostContext()));
  return thunk;
#endif  // BEF_THUNKS
}

BefThunk::BefThunk(ThunkInfo thunk_info, Thunk::Kind kind,
                   std::vector<BufferAllocation::Slice> inputs,
                   std::vector<BufferAllocation::Slice> outputs)
    : Thunk(kind, thunk_info)
#if BEF_THUNKS
      ,
      inputs_(std::move(inputs)),
      outputs_(std::move(outputs))
#endif  // BEF_THUNKS
{
}

Status BefThunk::ExecuteOnStream(const ExecuteParams& params) {
#if !BEF_THUNKS
  LOG(FATAL) << "BefThunk is disabled.";
#else
  VLOG(2) << "Executing BEF thunk.";

  auto* host = runtime().core_runtime()->GetHostContext();
  auto gpu_system = tfrt::gpu::System::Instantiate(host);
  host->Await({gpu_system.CopyRCRef()});

  // Use the GPU resources specified in ExecuteParams (initialized by the
  // StreamExecutor).
  auto se_gpu_executor = down_cast<stream_executor::gpu::GpuExecutor*>(
      params.stream->parent()->implementation());
  auto se_gpu_stream = down_cast<stream_executor::gpu::GpuStream*>(
      params.stream->implementation());
  auto context =
      tfrt::gpu::wrapper::Context(se_gpu_executor->gpu_context()->context());
  auto stream = tfrt::gpu::wrapper::Stream(se_gpu_stream->gpu_stream());
  tfrt::gpu::BorrowedGpuStream gpu_stream(host, context, stream);

  // Prepare arguments for BEF execution.
  auto get_async_value_ref = [&](const BufferAllocation::Slice& slice) {
    se::DeviceMemoryBase data =
        params.buffer_allocations->GetDeviceAddress(slice);
    tfrt::gpu::wrapper::Pointer<void> pointer(
        data.opaque(), tfrt::gpu::wrapper::Platform::CUDA);
    auto allocator =
        tfrt::MakeAvailableAsyncValueRef<tfrt::gpu::GpuOneShotAllocator<void>>(
            host, pointer);
    return tfrt::MakeAvailableAsyncValueRef<tfrt::gpu::GpuBuffer>(
        host, std::move(allocator), pointer, data.size());
  };

  llvm::SmallVector<tfrt::AsyncValueRef<tfrt::gpu::GpuBuffer>, 4> inputs;
  llvm::SmallVector<tfrt::AsyncValueRef<tfrt::gpu::GpuBuffer>, 4> outputs;
  for (auto input : inputs_) {
    inputs.push_back(get_async_value_ref(input));
  }
  for (auto output : outputs_) {
    outputs.push_back(get_async_value_ref(output));
  }

  tfrt::AsyncValueRef<tfrt::Chain> chain = tfrt::GetReadyChain(host);
  chain = gpu_system->Execute(exec_ctx(), *gpu_program_, gpu_stream, inputs,
                              outputs, std::move(chain));

  // TODO(hanbinyoon): BorrowedGpuStream releases context/stream information
  // upon destruction. Change it to not use wrapper::OwningContext and
  // wrapper::OwningStream.
  // TODO(b/184696034): Remove this. Right now we need this to ensure kernels
  // in bef function have been dispatched to stream.
  host->Quiesce();
  return Status::OK();
#endif  // BEF_THUNKS
}

StatusOr<Thunk::Kind> BefThunk::GetThunkKind(mlir::Operation* op) {
  if (mlir::isa<mlir::lmhlo_gpu::GEMMOp, mlir::lmhlo_gpu::GEMM_BiasOp>(op)) {
    return Thunk::Kind::kGemm;
  }
  return tensorflow::errors::Unimplemented(
      "Operation is not supported by BefThunk.");
}

#if BEF_THUNKS
tensorflow::tfrt_stub::Runtime& BefThunk::runtime() {
  static auto runtime = tensorflow::tfrt_stub::Runtime::Create().release();
  return *runtime;
}

tfrt::ExecutionContext& BefThunk::exec_ctx() {
  static tfrt::ExecutionContext* exec_ctx = [] {
    // Create request context and prepare deadline tracker.
    tfrt::RequestContextBuilder request_context_builder(
        runtime().core_runtime()->GetHostContext(),
        /*resource_context=*/nullptr);
    tensorflow::thread::ThreadPoolInterface* intra_op_threadpool = nullptr;
    DCHECK(
        runtime()
            .work_queue()
            ->InitializeRequest(&request_context_builder, &intra_op_threadpool)
            .ok());
    auto req_ctx = std::move(request_context_builder).build();
    DCHECK(req_ctx);
    return new tfrt::ExecutionContext(std::move(*req_ctx));
  }();
  return *exec_ctx;
}
#endif  // BEF_THUNKS

}  // namespace gpu
}  // namespace xla
<|endoftext|>