<commit_before>/*
 * Copyright 2011 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "gm.h"
#include "sk_tool_utils.h"
#include "SkCanvas.h"
#include "SkColorFilter.h"
#include "SkColorPriv.h"
#include "SkImageFilterPriv.h"
#include "SkShader.h"

#include "SkBlurImageFilter.h"
#include "SkColorFilterImageFilter.h"
#include "SkDropShadowImageFilter.h"
#include "SkSpecialImage.h"

class FailImageFilter : public SkImageFilter {
public:
    class Registrar {
    public:
        Registrar() {
            SkFlattenable::Register("FailImageFilter",
                                    FailImageFilter::CreateProc,
                                    FailImageFilter::GetFlattenableType());
        }
    };
    static sk_sp<SkImageFilter> Make() {
        return sk_sp<SkImageFilter>(new FailImageFilter);
    }

    SK_TO_STRING_OVERRIDE()
    SK_DECLARE_PUBLIC_FLATTENABLE_DESERIALIZATION_PROCS(FailImageFilter)

protected:
    FailImageFilter() : INHERITED(nullptr, 0, nullptr) {}

    sk_sp<SkSpecialImage> onFilterImage(SkSpecialImage* source, const Context&,
                                        SkIPoint* offset) const override {
        return nullptr;
    }
    sk_sp<SkImageFilter> onMakeColorSpace(SkColorSpaceXformer*) const override {
        return nullptr;
    }

private:
    typedef SkImageFilter INHERITED;
};

static FailImageFilter::Registrar gReg0;

sk_sp<SkFlattenable> FailImageFilter::CreateProc(SkReadBuffer& buffer) {
    SK_IMAGEFILTER_UNFLATTEN_COMMON(common, 0);
    return FailImageFilter::Make();
}

#ifndef SK_IGNORE_TO_STRING
void FailImageFilter::toString(SkString* str) const {
    str->appendf("FailImageFilter: (");
    str->append(")");
}
#endif

class IdentityImageFilter : public SkImageFilter {
public:
    class Registrar {
    public:
        Registrar() {
            SkFlattenable::Register("IdentityImageFilter",
                                    IdentityImageFilter::CreateProc,
                                    IdentityImageFilter::GetFlattenableType());
        }
    };
    static sk_sp<SkImageFilter> Make(sk_sp<SkImageFilter> input) {
        return sk_sp<SkImageFilter>(new IdentityImageFilter(std::move(input)));
    }

    SK_TO_STRING_OVERRIDE()
    SK_DECLARE_PUBLIC_FLATTENABLE_DESERIALIZATION_PROCS(IdentityImageFilter)

protected:
    sk_sp<SkSpecialImage> onFilterImage(SkSpecialImage* source, const Context&,
                                        SkIPoint* offset) const override {
        offset->set(0, 0);
        return sk_ref_sp<SkSpecialImage>(source);
    }
    sk_sp<SkImageFilter> onMakeColorSpace(SkColorSpaceXformer*) const override {
        return sk_ref_sp(const_cast<IdentityImageFilter*>(this));
    }

private:
    IdentityImageFilter(sk_sp<SkImageFilter> input) : INHERITED(&input, 1, nullptr) {}

    typedef SkImageFilter INHERITED;
};

static IdentityImageFilter::Registrar gReg1;

sk_sp<SkFlattenable> IdentityImageFilter::CreateProc(SkReadBuffer& buffer) {
    SK_IMAGEFILTER_UNFLATTEN_COMMON(common, 1);
    return IdentityImageFilter::Make(common.getInput(0));
}

#ifndef SK_IGNORE_TO_STRING
void IdentityImageFilter::toString(SkString* str) const {
    str->appendf("IdentityImageFilter: (");
    str->append(")");
}
#endif

///////////////////////////////////////////////////////////////////////////////

static void draw_paint(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setImageFilter(std::move(imf));
    paint.setColor(SK_ColorGREEN);
    canvas->save();
    canvas->clipRect(r);
    canvas->drawPaint(paint);
    canvas->restore();
}

static void draw_line(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setColor(SK_ColorBLUE);
    paint.setImageFilter(imf);
    paint.setStrokeWidth(r.width()/10);
    canvas->drawLine(r.fLeft, r.fTop, r.fRight, r.fBottom, paint);
}

static void draw_rect(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setColor(SK_ColorYELLOW);
    paint.setImageFilter(imf);
    SkRect rr(r);
    rr.inset(r.width()/10, r.height()/10);
    canvas->drawRect(rr, paint);
}

static void draw_path(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setColor(SK_ColorMAGENTA);
    paint.setImageFilter(imf);
    paint.setAntiAlias(true);
    canvas->drawCircle(r.centerX(), r.centerY(), r.width()*2/5, paint);
}

static void draw_text(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setImageFilter(imf);
    paint.setColor(SK_ColorCYAN);
    paint.setAntiAlias(true);
    sk_tool_utils::set_portable_typeface(&paint);
    paint.setTextSize(r.height()/2);
    paint.setTextAlign(SkPaint::kCenter_Align);
    canvas->drawString("Text", r.centerX(), r.centerY(), paint);
}

static void draw_bitmap(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setImageFilter(std::move(imf));

    SkIRect bounds;
    r.roundOut(&bounds);

    SkBitmap bm;
    bm.allocN32Pixels(bounds.width(), bounds.height());
    bm.eraseColor(SK_ColorTRANSPARENT);
    SkCanvas c(bm);
    draw_path(&c, r, nullptr);

    canvas->drawBitmap(bm, 0, 0, &paint);
}

///////////////////////////////////////////////////////////////////////////////

class ImageFiltersBaseGM : public skiagm::GM {
public:
    ImageFiltersBaseGM () {}

protected:
    SkString onShortName() override {
        return SkString("imagefiltersbase");
    }

    SkISize onISize() override { return SkISize::Make(700, 500); }

    void draw_frame(SkCanvas* canvas, const SkRect& r) {
        SkPaint paint;
        paint.setStyle(SkPaint::kStroke_Style);
        paint.setColor(SK_ColorRED);
        canvas->drawRect(r, paint);
    }

    void onDraw(SkCanvas* canvas) override {
        void (*drawProc[])(SkCanvas*, const SkRect&, sk_sp<SkImageFilter>) = {
            draw_paint,
            draw_line, draw_rect, draw_path, draw_text,
            draw_bitmap,
        };

        auto cf = SkColorFilter::MakeModeFilter(SK_ColorRED, SkBlendMode::kSrcIn);
        sk_sp<SkImageFilter> filters[] = {
            nullptr,
            IdentityImageFilter::Make(nullptr),
            FailImageFilter::Make(),
            SkColorFilterImageFilter::Make(std::move(cf), nullptr),
            SkBlurImageFilter::Make(12.0f, 0.0f, nullptr),
            SkDropShadowImageFilter::Make(
                                    10.0f, 5.0f, 3.0f, 3.0f, SK_ColorBLUE,
                                    SkDropShadowImageFilter::kDrawShadowAndForeground_ShadowMode,
                                    nullptr),
        };

        SkRect r = SkRect::MakeWH(SkIntToScalar(64), SkIntToScalar(64));
        SkScalar MARGIN = SkIntToScalar(16);
        SkScalar DX = r.width() + MARGIN;
        SkScalar DY = r.height() + MARGIN;

        canvas->translate(MARGIN, MARGIN);
        for (size_t i = 0; i < SK_ARRAY_COUNT(drawProc); ++i) {
            canvas->save();
            for (size_t j = 0; j < SK_ARRAY_COUNT(filters); ++j) {
                drawProc[i](canvas, r, filters[j]);

                draw_frame(canvas, r);
                canvas->translate(0, DY);
            }
            canvas->restore();
            canvas->translate(DX, 0);
        }
    }

private:
    typedef GM INHERITED;
};
DEF_GM( return new ImageFiltersBaseGM; )

///////////////////////////////////////////////////////////////////////////////

/*
 *  Want to test combos of filter and LCD text, to be sure we disable LCD in the presence of
 *  a filter.
 */
class ImageFiltersTextBaseGM : public skiagm::GM {
    SkString fSuffix;
public:
    ImageFiltersTextBaseGM(const char suffix[]) : fSuffix(suffix) {}

protected:
    SkString onShortName() override {
        SkString name;
        name.printf("%s_%s", "textfilter", fSuffix.c_str());
        return name;
    }

    SkISize onISize() override { return SkISize::Make(512, 342); }

    void drawWaterfall(SkCanvas* canvas, const SkPaint& origPaint) {
        const uint32_t flags[] = {
            0,
            SkPaint::kAntiAlias_Flag,
            SkPaint::kAntiAlias_Flag | SkPaint::kLCDRenderText_Flag,
        };
        SkPaint paint(origPaint);
        sk_tool_utils::set_portable_typeface(&paint);
        paint.setTextSize(30);

        SkAutoCanvasRestore acr(canvas, true);
        for (size_t i = 0; i < SK_ARRAY_COUNT(flags); ++i) {
            paint.setFlags(flags[i]);
            canvas->drawString("Hamburgefon", 0, 0, paint);
            canvas->translate(0, 40);
        }
    }

    virtual void installFilter(SkPaint* paint) = 0;

    void onDraw(SkCanvas* canvas) override {
        SkPaint paint;

        canvas->translate(20, 40);

        for (int doSaveLayer = 0; doSaveLayer <= 1; ++doSaveLayer) {
            SkAutoCanvasRestore acr(canvas, true);
            for (int useFilter = 0; useFilter <= 1; ++useFilter) {
                SkAutoCanvasRestore acr2(canvas, true);

                SkPaint paint;
                if (useFilter) {
                    this->installFilter(&paint);
                }
                if (doSaveLayer) {
                    canvas->saveLayer(nullptr, &paint);
                    paint.setImageFilter(nullptr);
                }
                this->drawWaterfall(canvas, paint);

                acr2.restore();
                canvas->translate(250, 0);
            }
            acr.restore();
            canvas->translate(0, 200);
        }
    }

private:
    typedef GM INHERITED;
};

class ImageFiltersText_IF : public ImageFiltersTextBaseGM {
public:
    ImageFiltersText_IF() : ImageFiltersTextBaseGM("image") {}

    void installFilter(SkPaint* paint) override {
        paint->setImageFilter(SkBlurImageFilter::Make(1.5f, 1.5f, nullptr));
    }
};
DEF_GM( return new ImageFiltersText_IF; )

class ImageFiltersText_CF : public ImageFiltersTextBaseGM {
public:
    ImageFiltersText_CF() : ImageFiltersTextBaseGM("color") {}

    void installFilter(SkPaint* paint) override {
        paint->setColorFilter(SkColorFilter::MakeModeFilter(SK_ColorBLUE, SkBlendMode::kSrcIn));
    }
};
DEF_GM( return new ImageFiltersText_CF; )
<commit_msg>Make unit test tickle msan bug.<commit_after>/*
 * Copyright 2011 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "gm.h"
#include "sk_tool_utils.h"
#include "SkCanvas.h"
#include "SkColorFilter.h"
#include "SkColorPriv.h"
#include "SkImageFilterPriv.h"
#include "SkShader.h"

#include "SkBlurImageFilter.h"
#include "SkColorFilterImageFilter.h"
#include "SkDropShadowImageFilter.h"
#include "SkSpecialImage.h"

class FailImageFilter : public SkImageFilter {
public:
    class Registrar {
    public:
        Registrar() {
            SkFlattenable::Register("FailImageFilter",
                                    FailImageFilter::CreateProc,
                                    FailImageFilter::GetFlattenableType());
        }
    };
    static sk_sp<SkImageFilter> Make() {
        return sk_sp<SkImageFilter>(new FailImageFilter);
    }

    SK_TO_STRING_OVERRIDE()
    SK_DECLARE_PUBLIC_FLATTENABLE_DESERIALIZATION_PROCS(FailImageFilter)

protected:
    FailImageFilter() : INHERITED(nullptr, 0, nullptr) {}

    sk_sp<SkSpecialImage> onFilterImage(SkSpecialImage* source, const Context&,
                                        SkIPoint* offset) const override {
        return nullptr;
    }
    sk_sp<SkImageFilter> onMakeColorSpace(SkColorSpaceXformer*) const override {
        return nullptr;
    }

private:
    typedef SkImageFilter INHERITED;
};

static FailImageFilter::Registrar gReg0;

sk_sp<SkFlattenable> FailImageFilter::CreateProc(SkReadBuffer& buffer) {
    SK_IMAGEFILTER_UNFLATTEN_COMMON(common, 0);
    return FailImageFilter::Make();
}

#ifndef SK_IGNORE_TO_STRING
void FailImageFilter::toString(SkString* str) const {
    str->appendf("FailImageFilter: (");
    str->append(")");
}
#endif

class IdentityImageFilter : public SkImageFilter {
public:
    class Registrar {
    public:
        Registrar() {
            SkFlattenable::Register("IdentityImageFilter",
                                    IdentityImageFilter::CreateProc,
                                    IdentityImageFilter::GetFlattenableType());
        }
    };
    static sk_sp<SkImageFilter> Make(sk_sp<SkImageFilter> input) {
        return sk_sp<SkImageFilter>(new IdentityImageFilter(std::move(input)));
    }

    SK_TO_STRING_OVERRIDE()
    SK_DECLARE_PUBLIC_FLATTENABLE_DESERIALIZATION_PROCS(IdentityImageFilter)

protected:
    sk_sp<SkSpecialImage> onFilterImage(SkSpecialImage* source, const Context&,
                                        SkIPoint* offset) const override {
        offset->set(0, 0);
        return sk_ref_sp<SkSpecialImage>(source);
    }
    sk_sp<SkImageFilter> onMakeColorSpace(SkColorSpaceXformer*) const override {
        return sk_ref_sp(const_cast<IdentityImageFilter*>(this));
    }

private:
    IdentityImageFilter(sk_sp<SkImageFilter> input) : INHERITED(&input, 1, nullptr) {}

    typedef SkImageFilter INHERITED;
};

static IdentityImageFilter::Registrar gReg1;

sk_sp<SkFlattenable> IdentityImageFilter::CreateProc(SkReadBuffer& buffer) {
    SK_IMAGEFILTER_UNFLATTEN_COMMON(common, 1);
    return IdentityImageFilter::Make(common.getInput(0));
}

#ifndef SK_IGNORE_TO_STRING
void IdentityImageFilter::toString(SkString* str) const {
    str->appendf("IdentityImageFilter: (");
    str->append(")");
}
#endif

///////////////////////////////////////////////////////////////////////////////

static void draw_paint(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setImageFilter(std::move(imf));
    paint.setColor(SK_ColorGREEN);
    canvas->save();
    canvas->clipRect(r);
    canvas->drawPaint(paint);
    canvas->restore();
}

static void draw_line(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setColor(SK_ColorBLUE);
    paint.setImageFilter(imf);
    paint.setStrokeWidth(r.width()/10);
    canvas->drawLine(r.fLeft, r.fTop, r.fRight, r.fBottom, paint);
}

static void draw_rect(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setColor(SK_ColorYELLOW);
    paint.setImageFilter(imf);
    SkRect rr(r);
    rr.inset(r.width()/10, r.height()/10);
    canvas->drawRect(rr, paint);
}

static void draw_path(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setColor(SK_ColorMAGENTA);
    paint.setImageFilter(imf);
    paint.setAntiAlias(true);
    canvas->drawCircle(r.centerX(), r.centerY(), r.width()*2/5, paint);
}

static void draw_text(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setImageFilter(imf);
    paint.setColor(SK_ColorCYAN);
    paint.setAntiAlias(true);
    sk_tool_utils::set_portable_typeface(&paint);
    paint.setTextSize(r.height()/2);
    paint.setTextAlign(SkPaint::kCenter_Align);
    canvas->drawString("Text", r.centerX(), r.centerY(), paint);
}

static void draw_bitmap(SkCanvas* canvas, const SkRect& r, sk_sp<SkImageFilter> imf) {
    SkPaint paint;
    paint.setImageFilter(std::move(imf));

    SkIRect bounds;
    r.roundOut(&bounds);

    SkBitmap bm;
    bm.allocN32Pixels(bounds.width(), bounds.height());
    bm.eraseColor(SK_ColorTRANSPARENT);
    SkCanvas c(bm);
    draw_path(&c, r, nullptr);

    canvas->drawBitmap(bm, 0, 0, &paint);
}

///////////////////////////////////////////////////////////////////////////////

class ImageFiltersBaseGM : public skiagm::GM {
public:
    ImageFiltersBaseGM () {}

protected:
    SkString onShortName() override {
        return SkString("imagefiltersbase");
    }

    SkISize onISize() override { return SkISize::Make(700, 500); }

    void draw_frame(SkCanvas* canvas, const SkRect& r) {
        SkPaint paint;
        paint.setStyle(SkPaint::kStroke_Style);
        paint.setColor(SK_ColorRED);
        canvas->drawRect(r, paint);
    }

    void onDraw(SkCanvas* canvas) override {
        void (*drawProc[])(SkCanvas*, const SkRect&, sk_sp<SkImageFilter>) = {
            draw_paint,
            draw_line, draw_rect, draw_path, draw_text,
            draw_bitmap,
        };

        auto cf = SkColorFilter::MakeModeFilter(SK_ColorRED, SkBlendMode::kSrcIn);
        sk_sp<SkImageFilter> filters[] = {
            nullptr,
            IdentityImageFilter::Make(nullptr),
            FailImageFilter::Make(),
            SkColorFilterImageFilter::Make(std::move(cf), nullptr),
            // The strage 0.29 value tickles an edge case where crop rect calculates
            // a small border, but the blur really needs no border. This tickels
            // an msan uninitialized value bug.
            SkBlurImageFilter::Make(12.0f, 0.29f, nullptr),
            SkDropShadowImageFilter::Make(
                                    10.0f, 5.0f, 3.0f, 3.0f, SK_ColorBLUE,
                                    SkDropShadowImageFilter::kDrawShadowAndForeground_ShadowMode,
                                    nullptr),
        };

        SkRect r = SkRect::MakeWH(SkIntToScalar(64), SkIntToScalar(64));
        SkScalar MARGIN = SkIntToScalar(16);
        SkScalar DX = r.width() + MARGIN;
        SkScalar DY = r.height() + MARGIN;

        canvas->translate(MARGIN, MARGIN);
        for (size_t i = 0; i < SK_ARRAY_COUNT(drawProc); ++i) {
            canvas->save();
            for (size_t j = 0; j < SK_ARRAY_COUNT(filters); ++j) {
                drawProc[i](canvas, r, filters[j]);

                draw_frame(canvas, r);
                canvas->translate(0, DY);
            }
            canvas->restore();
            canvas->translate(DX, 0);
        }
    }

private:
    typedef GM INHERITED;
};
DEF_GM( return new ImageFiltersBaseGM; )

///////////////////////////////////////////////////////////////////////////////

/*
 *  Want to test combos of filter and LCD text, to be sure we disable LCD in the presence of
 *  a filter.
 */
class ImageFiltersTextBaseGM : public skiagm::GM {
    SkString fSuffix;
public:
    ImageFiltersTextBaseGM(const char suffix[]) : fSuffix(suffix) {}

protected:
    SkString onShortName() override {
        SkString name;
        name.printf("%s_%s", "textfilter", fSuffix.c_str());
        return name;
    }

    SkISize onISize() override { return SkISize::Make(512, 342); }

    void drawWaterfall(SkCanvas* canvas, const SkPaint& origPaint) {
        const uint32_t flags[] = {
            0,
            SkPaint::kAntiAlias_Flag,
            SkPaint::kAntiAlias_Flag | SkPaint::kLCDRenderText_Flag,
        };
        SkPaint paint(origPaint);
        sk_tool_utils::set_portable_typeface(&paint);
        paint.setTextSize(30);

        SkAutoCanvasRestore acr(canvas, true);
        for (size_t i = 0; i < SK_ARRAY_COUNT(flags); ++i) {
            paint.setFlags(flags[i]);
            canvas->drawString("Hamburgefon", 0, 0, paint);
            canvas->translate(0, 40);
        }
    }

    virtual void installFilter(SkPaint* paint) = 0;

    void onDraw(SkCanvas* canvas) override {
        SkPaint paint;

        canvas->translate(20, 40);

        for (int doSaveLayer = 0; doSaveLayer <= 1; ++doSaveLayer) {
            SkAutoCanvasRestore acr(canvas, true);
            for (int useFilter = 0; useFilter <= 1; ++useFilter) {
                SkAutoCanvasRestore acr2(canvas, true);

                SkPaint paint;
                if (useFilter) {
                    this->installFilter(&paint);
                }
                if (doSaveLayer) {
                    canvas->saveLayer(nullptr, &paint);
                    paint.setImageFilter(nullptr);
                }
                this->drawWaterfall(canvas, paint);

                acr2.restore();
                canvas->translate(250, 0);
            }
            acr.restore();
            canvas->translate(0, 200);
        }
    }

private:
    typedef GM INHERITED;
};

class ImageFiltersText_IF : public ImageFiltersTextBaseGM {
public:
    ImageFiltersText_IF() : ImageFiltersTextBaseGM("image") {}

    void installFilter(SkPaint* paint) override {
        paint->setImageFilter(SkBlurImageFilter::Make(1.5f, 1.5f, nullptr));
    }
};
DEF_GM( return new ImageFiltersText_IF; )

class ImageFiltersText_CF : public ImageFiltersTextBaseGM {
public:
    ImageFiltersText_CF() : ImageFiltersTextBaseGM("color") {}

    void installFilter(SkPaint* paint) override {
        paint->setColorFilter(SkColorFilter::MakeModeFilter(SK_ColorBLUE, SkBlendMode::kSrcIn));
    }
};
DEF_GM( return new ImageFiltersText_CF; )
<|endoftext|>


<commit_before>/***********************************************************************************************************************
 *
 * Copyright (c) 2015, ABB Schweiz AG
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with
 * or without modification, are permitted provided that
 * the following conditions are met:
 *
 *    * Redistributions of source code must retain the
 *      above copyright notice, this list of conditions
 *      and the following disclaimer.
 *    * Redistributions in binary form must reproduce the
 *      above copyright notice, this list of conditions
 *      and the following disclaimer in the documentation
 *      and/or other materials provided with the
 *      distribution.
 *    * Neither the name of ABB nor the names of its
 *      contributors may be used to endorse or promote
 *      products derived from this software without
 *      specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 ***********************************************************************************************************************
 */

#include <sstream>

#include "abb_libegm/egm_common.h"
#include "abb_libegm/egm_logger.h"

namespace abb
{
namespace egm
{
/***********************************************************************************************************************
 * Class definitions: EGMLogger
 */

/************************************************************
 * Primary methods
 */

EGMLogger::EGMLogger(const std::string& filename, const bool use_default_headers)
:
number_of_logged_messages_(0)
{
  log_stream_.open(filename, std::ios::trunc);

  if (use_default_headers)
  {
    std::stringstream ss;
    ss << "TIMESTAMP,"
       // Robot feedback.
       << "R_FB_POS_RJ1,R_FB_POS_RJ2,R_FB_POS_RJ3,R_FB_POS_RJ4,R_FB_POS_RJ5,R_FB_POS_RJ6,"
       << "R_FB_POS_EJ1,R_FB_POS_EJ2,R_FB_POS_EJ3,R_FB_POS_EJ4,R_FB_POS_EJ5,R_FB_POS_EJ6,"
       << "R_FB_VEL_RJ1,R_FB_VEL_RJ2,R_FB_VEL_RJ3,R_FB_VEL_RJ4,R_FB_VEL_RJ5,R_FB_VEL_RJ6,"
       << "R_FB_VEL_EJ1,R_FB_VEL_EJ2,R_FB_VEL_EJ3,R_FB_VEL_EJ4,R_FB_VEL_EJ5,R_FB_VEL_EJ6,"
       << "R_FB_POS_X,R_FB_POS_Y,R_FB_POS_Z,"
       << "R_FB_EULER_X,R_FB_EULER_Y,R_FB_EULER_Z,"
       << "R_FB_QUAT_U0,R_FB_QUAT_U1,R_FB_QUAT_U2,R_FB_QUAT_U3,"
       << "R_FB_VEL_LX,R_FB_VEL_LY,R_FB_VEL_LZ,"
       << "R_FB_VEL_AX,R_FB_VEL_AY,R_FB_VEL_AZ,"
       // Robot planned.
       << "R_PL_POS_RJ1,R_PL_POS_RJ2,R_PL_POS_RJ3,R_PL_POS_RJ4,R_PL_POS_RJ5,R_PL_POS_RJ6,"
       << "R_PL_POS_EJ1,R_PL_POS_EJ2,R_PL_POS_EJ3,R_PL_POS_EJ4,R_PL_POS_EJ5,R_PL_POS_EJ6,"
       << "R_PL_VEL_RJ1,R_PL_VEL_RJ2,R_PL_VEL_RJ3,R_PL_VEL_RJ4,R_PL_VEL_RJ5,R_PL_VEL_RJ6,"
       << "R_PL_VEL_EJ1,R_PL_VEL_EJ2,R_PL_VEL_EJ3,R_PL_VEL_EJ4,R_PL_VEL_EJ5,R_PL_VEL_EJ6,"
       << "R_PL_POS_X,R_PL_POS_Y,R_PL_POS_Z,"
       << "R_PL_EULER_X,R_PL_EULER_Y,R_PL_EULER_Z,"
       << "R_PL_QUAT_U0,R_PL_QUAT_U1,R_PL_QUAT_U2,R_PL_QUAT_U3,"
       << "R_PL_VEL_LX,R_PL_VEL_LY,R_PL_VEL_LZ,"
       << "R_PL_VEL_AX,R_PL_VEL_AY,R_PL_VEL_AZ,"
       // Sensor references.
       << "S_REF_POS_RJ1,S_REF_POS_RJ2,S_REF_POS_RJ3,S_REF_POS_RJ4,S_REF_POS_RJ5,S_REF_POS_RJ6,"
       << "S_REF_POS_EJ1,S_REF_POS_EJ2,S_REF_POS_EJ3,S_REF_POS_EJ4,S_REF_POS_EJ5,S_REF_POS_EJ6,"
       << "S_REF_VEL_RJ1,S_REF_VEL_RJ2,S_REF_VEL_RJ3,S_REF_VEL_RJ4,S_REF_VEL_RJ5,S_REF_VEL_RJ6,"
       << "S_REF_VEL_EJ1,S_REF_VEL_EJ2,S_REF_VEL_EJ3,S_REF_VEL_EJ4,S_REF_VEL_EJ5,S_REF_VEL_EJ6,"
       << "S_REF_POS_X,S_REF_POS_Y,S_REF_POS_Z,"
       << "S_REF_EULER_X,S_REF_EULER_Y,S_REF_EULER_Z,"
       << "S_REF_QUAT_U0,S_REF_QUAT_U1,S_REF_QUAT_U2,S_REF_QUAT_U3,"
       << "S_REF_VEL_LX,S_REF_VEL_LY,S_REF_VEL_LZ,"
       << "S_REF_VEL_AX,S_REF_VEL_AY,S_REF_VEL_AZ\n";

    log_stream_ << ss.str();
    log_stream_.flush();
  }
}

EGMLogger::~EGMLogger()
{
  log_stream_.close();
}

void EGMLogger::flush()
{
  log_stream_ << "\n";
  log_stream_.flush();
  ++number_of_logged_messages_;
}

void EGMLogger::add(const wrapper::Header& header)
{
  log_stream_ << header.time_stamp() << ",";
}

void EGMLogger::add(const wrapper::Joints& robot, const wrapper::Joints& external)
{
  google::protobuf::RepeatedField<double>::const_iterator i;

  for (i = robot.values().begin(); i != robot.values().end(); ++i)
  {
    log_stream_ << *i << ",";
  }
  addMockJoints(true, robot.values_size(), external.values_size());

  for (i = external.values().begin(); i != external.values().end(); ++i)
  {
    log_stream_ << *i << ",";
  }
  addMockJoints(false, robot.values_size(), external.values_size());
}

void EGMLogger::add(const wrapper::CartesianPose& pose)
{
  log_stream_ << pose.position().x() << ","
              << pose.position().y() << ","
              << pose.position().z() << ",";

  log_stream_ << pose.euler().x() << ","
              << pose.euler().y() << ","
              << pose.euler().z() << ",";

  log_stream_ << pose.quaternion().u0() << ","
              << pose.quaternion().u1() << ","
              << pose.quaternion().u2() << ","
              << pose.quaternion().u3() << ",";
}

void EGMLogger::add(const wrapper::CartesianVelocity& velocity, const bool last)
{
  log_stream_ << velocity.linear().x() << ","
              << velocity.linear().y() << ","
              << velocity.linear().z() << ",";

  log_stream_ << velocity.angular().x() << ","
              << velocity.angular().y() << ","
              << velocity.angular().z() << (last ? "" : ",");
}

/************************************************************
 * Auxiliary methods
 */

void EGMLogger::addMockJoints(const bool robot, const size_t robot_size, const size_t external_size)
{
  if (robot)
  {
    // Add mock values for the missing robot joint data.
    size_t condition = Constants::RobotController::DEFAULT_NUMBER_OF_ROBOT_JOINTS;
    for (size_t i = robot_size; i < condition; ++i)
    {
      log_stream_ << 0.0 << ",";
    }
  }
  else
  {
    // Add mock values for the missing external joint data.
    size_t condition = Constants::RobotController::MAX_NUMBER_OF_JOINTS - robot_size;
    for (size_t i = external_size; i < condition; ++i)
    {
      log_stream_ << 0.0 << ",";
    }
  }
}

double EGMLogger::calculateTimeLogged(const double sample_time)
{
  return (double)number_of_logged_messages_*sample_time;
}

} // end namespace egm
} // end namespace abb<commit_msg>Changed to ofstream method open(char*, ...) instead of open(string&, ...)<commit_after>/***********************************************************************************************************************
 *
 * Copyright (c) 2015, ABB Schweiz AG
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with
 * or without modification, are permitted provided that
 * the following conditions are met:
 *
 *    * Redistributions of source code must retain the
 *      above copyright notice, this list of conditions
 *      and the following disclaimer.
 *    * Redistributions in binary form must reproduce the
 *      above copyright notice, this list of conditions
 *      and the following disclaimer in the documentation
 *      and/or other materials provided with the
 *      distribution.
 *    * Neither the name of ABB nor the names of its
 *      contributors may be used to endorse or promote
 *      products derived from this software without
 *      specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 ***********************************************************************************************************************
 */

#include <sstream>

#include "abb_libegm/egm_common.h"
#include "abb_libegm/egm_logger.h"

namespace abb
{
namespace egm
{
/***********************************************************************************************************************
 * Class definitions: EGMLogger
 */

/************************************************************
 * Primary methods
 */

EGMLogger::EGMLogger(const std::string& filename, const bool use_default_headers)
:
number_of_logged_messages_(0)
{
  log_stream_.open(filename.c_str(), std::ios::trunc);

  if (use_default_headers)
  {
    std::stringstream ss;
    ss << "TIMESTAMP,"
       // Robot feedback.
       << "R_FB_POS_RJ1,R_FB_POS_RJ2,R_FB_POS_RJ3,R_FB_POS_RJ4,R_FB_POS_RJ5,R_FB_POS_RJ6,"
       << "R_FB_POS_EJ1,R_FB_POS_EJ2,R_FB_POS_EJ3,R_FB_POS_EJ4,R_FB_POS_EJ5,R_FB_POS_EJ6,"
       << "R_FB_VEL_RJ1,R_FB_VEL_RJ2,R_FB_VEL_RJ3,R_FB_VEL_RJ4,R_FB_VEL_RJ5,R_FB_VEL_RJ6,"
       << "R_FB_VEL_EJ1,R_FB_VEL_EJ2,R_FB_VEL_EJ3,R_FB_VEL_EJ4,R_FB_VEL_EJ5,R_FB_VEL_EJ6,"
       << "R_FB_POS_X,R_FB_POS_Y,R_FB_POS_Z,"
       << "R_FB_EULER_X,R_FB_EULER_Y,R_FB_EULER_Z,"
       << "R_FB_QUAT_U0,R_FB_QUAT_U1,R_FB_QUAT_U2,R_FB_QUAT_U3,"
       << "R_FB_VEL_LX,R_FB_VEL_LY,R_FB_VEL_LZ,"
       << "R_FB_VEL_AX,R_FB_VEL_AY,R_FB_VEL_AZ,"
       // Robot planned.
       << "R_PL_POS_RJ1,R_PL_POS_RJ2,R_PL_POS_RJ3,R_PL_POS_RJ4,R_PL_POS_RJ5,R_PL_POS_RJ6,"
       << "R_PL_POS_EJ1,R_PL_POS_EJ2,R_PL_POS_EJ3,R_PL_POS_EJ4,R_PL_POS_EJ5,R_PL_POS_EJ6,"
       << "R_PL_VEL_RJ1,R_PL_VEL_RJ2,R_PL_VEL_RJ3,R_PL_VEL_RJ4,R_PL_VEL_RJ5,R_PL_VEL_RJ6,"
       << "R_PL_VEL_EJ1,R_PL_VEL_EJ2,R_PL_VEL_EJ3,R_PL_VEL_EJ4,R_PL_VEL_EJ5,R_PL_VEL_EJ6,"
       << "R_PL_POS_X,R_PL_POS_Y,R_PL_POS_Z,"
       << "R_PL_EULER_X,R_PL_EULER_Y,R_PL_EULER_Z,"
       << "R_PL_QUAT_U0,R_PL_QUAT_U1,R_PL_QUAT_U2,R_PL_QUAT_U3,"
       << "R_PL_VEL_LX,R_PL_VEL_LY,R_PL_VEL_LZ,"
       << "R_PL_VEL_AX,R_PL_VEL_AY,R_PL_VEL_AZ,"
       // Sensor references.
       << "S_REF_POS_RJ1,S_REF_POS_RJ2,S_REF_POS_RJ3,S_REF_POS_RJ4,S_REF_POS_RJ5,S_REF_POS_RJ6,"
       << "S_REF_POS_EJ1,S_REF_POS_EJ2,S_REF_POS_EJ3,S_REF_POS_EJ4,S_REF_POS_EJ5,S_REF_POS_EJ6,"
       << "S_REF_VEL_RJ1,S_REF_VEL_RJ2,S_REF_VEL_RJ3,S_REF_VEL_RJ4,S_REF_VEL_RJ5,S_REF_VEL_RJ6,"
       << "S_REF_VEL_EJ1,S_REF_VEL_EJ2,S_REF_VEL_EJ3,S_REF_VEL_EJ4,S_REF_VEL_EJ5,S_REF_VEL_EJ6,"
       << "S_REF_POS_X,S_REF_POS_Y,S_REF_POS_Z,"
       << "S_REF_EULER_X,S_REF_EULER_Y,S_REF_EULER_Z,"
       << "S_REF_QUAT_U0,S_REF_QUAT_U1,S_REF_QUAT_U2,S_REF_QUAT_U3,"
       << "S_REF_VEL_LX,S_REF_VEL_LY,S_REF_VEL_LZ,"
       << "S_REF_VEL_AX,S_REF_VEL_AY,S_REF_VEL_AZ\n";

    log_stream_ << ss.str();
    log_stream_.flush();
  }
}

EGMLogger::~EGMLogger()
{
  log_stream_.close();
}

void EGMLogger::flush()
{
  log_stream_ << "\n";
  log_stream_.flush();
  ++number_of_logged_messages_;
}

void EGMLogger::add(const wrapper::Header& header)
{
  log_stream_ << header.time_stamp() << ",";
}

void EGMLogger::add(const wrapper::Joints& robot, const wrapper::Joints& external)
{
  google::protobuf::RepeatedField<double>::const_iterator i;

  for (i = robot.values().begin(); i != robot.values().end(); ++i)
  {
    log_stream_ << *i << ",";
  }
  addMockJoints(true, robot.values_size(), external.values_size());

  for (i = external.values().begin(); i != external.values().end(); ++i)
  {
    log_stream_ << *i << ",";
  }
  addMockJoints(false, robot.values_size(), external.values_size());
}

void EGMLogger::add(const wrapper::CartesianPose& pose)
{
  log_stream_ << pose.position().x() << ","
              << pose.position().y() << ","
              << pose.position().z() << ",";

  log_stream_ << pose.euler().x() << ","
              << pose.euler().y() << ","
              << pose.euler().z() << ",";

  log_stream_ << pose.quaternion().u0() << ","
              << pose.quaternion().u1() << ","
              << pose.quaternion().u2() << ","
              << pose.quaternion().u3() << ",";
}

void EGMLogger::add(const wrapper::CartesianVelocity& velocity, const bool last)
{
  log_stream_ << velocity.linear().x() << ","
              << velocity.linear().y() << ","
              << velocity.linear().z() << ",";

  log_stream_ << velocity.angular().x() << ","
              << velocity.angular().y() << ","
              << velocity.angular().z() << (last ? "" : ",");
}

/************************************************************
 * Auxiliary methods
 */

void EGMLogger::addMockJoints(const bool robot, const size_t robot_size, const size_t external_size)
{
  if (robot)
  {
    // Add mock values for the missing robot joint data.
    size_t condition = Constants::RobotController::DEFAULT_NUMBER_OF_ROBOT_JOINTS;
    for (size_t i = robot_size; i < condition; ++i)
    {
      log_stream_ << 0.0 << ",";
    }
  }
  else
  {
    // Add mock values for the missing external joint data.
    size_t condition = Constants::RobotController::MAX_NUMBER_OF_JOINTS - robot_size;
    for (size_t i = external_size; i < condition; ++i)
    {
      log_stream_ << 0.0 << ",";
    }
  }
}

double EGMLogger::calculateTimeLogged(const double sample_time)
{
  return (double)number_of_logged_messages_*sample_time;
}

} // end namespace egm
} // end namespace abb<|endoftext|>


<commit_before>// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/labs.h"

#include <algorithm>
#include <iterator>
#include <map>
#include <set>

#include "app/l10n_util.h"
#include "base/command_line.h"
#include "base/values.h"
#include "chrome/browser/platform_util.h"
#include "chrome/browser/prefs/pref_service.h"
#include "chrome/browser/profile.h"
#include "chrome/common/chrome_switches.h"
#include "chrome/common/pref_names.h"
#include "grit/generated_resources.h"

namespace about_labs {

enum { kOsMac = 1 << 0, kOsWin = 1 << 1, kOsLinux = 1 << 2 };

unsigned kOsAll = kOsMac | kOsWin | kOsLinux;

struct Experiment {
  // The internal name of the experiment. This is never shown to the user.
  // It _is_ however stored in the prefs file, so you shouldn't change the
  // name of existing labs.
  const char* internal_name;

  // String id of the message containing the experiment's name.
  int visible_name_id;

  // String id of the message containing the experiment's description.
  int visible_description_id;

  // The platforms the experiment is available on
  // Needs to be more than a compile-time #ifdef because of profile sync.
  unsigned supported_platforms;  // bitmask

  // The commandline parameter that's added when this lab is active. This is
  // different from |internal_name| so that the commandline flag can be
  // renamed without breaking the prefs file.
  const char* command_line;
};

const Experiment kExperiments[] = {
  {
    "expose-for-tabs",  // Do not change; see above.
    IDS_LABS_TABPOSE_NAME,
    IDS_LABS_TABPOSE_DESCRIPTION,
    kOsMac,
#if defined(OS_MACOSX)
    // The switch exists only on OS X.
    switches::kEnableExposeForTabs
#else
    ""
#endif
  },
  {
    "vertical-tabs",  // Do not change; see above.
    IDS_LABS_SIDE_TABS_NAME,
    IDS_LABS_SIDE_TABS_DESCRIPTION,
    kOsWin,
    switches::kEnableVerticalTabs
  },
  {
    "tabbed-options",  // Do not change; see above.
    IDS_LABS_TABBED_OPTIONS_NAME,
    IDS_LABS_TABBED_OPTIONS_DESCRIPTION,
    kOsAll,
    switches::kEnableTabbedOptions
  },
  {
    "match-preview",  // Do not change; see above.
    IDS_LABS_INSTANT_NAME,
    IDS_LABS_INSTANT_DESCRIPTION,
    kOsWin,
    switches::kEnableMatchPreview
  },
  {
    "remoting",  // Do not change; see above.
    IDS_LABS_REMOTING_NAME,
#if defined(OS_WIN)
    // Windows only supports host functionality at the moment.
    IDS_LABS_REMOTING_HOST_DESCRIPTION,
#elif defined(OS_LINUX)
    // Linux only supports client functionality at the moment.
    IDS_LABS_REMOTING_CLIENT_DESCRIPTION,
#else
    // On other platforms, this lab isn't available at all.
    0,
#endif
    kOsWin | kOsLinux,
    switches::kEnableRemoting
  },
  {
    "page-info-bubble",  // Do not change; see above.
    IDS_LABS_PAGE_INFO_BUBBLE_NAME,
    IDS_LABS_PAGE_INFO_BUBBLE_DESCRIPTION,
    kOsAll,
    switches::kEnableNewPageInfoBubble
  },
  {
    "disable-outdated-plugins",  // Do not change; see above.
    IDS_LABS_DISABLE_OUTDATED_PLUGINS_NAME,
    IDS_LABS_DISABLE_OUTDATED_PLUGINS_DESCRIPTION,
    kOsAll,
    switches::kDisableOutdatedPlugins
  },
  {
    "xss-auditor",  // Do not change; see above.
    IDS_LABS_XSS_AUDITOR_NAME,
    IDS_LABS_XSS_AUDITOR_DESCRIPTION,
    kOsAll,
    switches::kEnableXSSAuditor
  },
};

// Extracts the list of enabled lab experiments from a profile and stores them
// in a set.
void GetEnabledLabs(const PrefService* prefs, std::set<std::string>* result) {
  const ListValue* enabled_experiments = prefs->GetList(
      prefs::kEnabledLabsExperiments);
  if (!enabled_experiments)
    return;

  for (ListValue::const_iterator it = enabled_experiments->begin();
       it != enabled_experiments->end();
       ++it) {
    std::string experiment_name;
    if (!(*it)->GetAsString(&experiment_name)) {
      LOG(WARNING) << "Invalid entry in " << prefs::kEnabledLabsExperiments;
      continue;
    }
    result->insert(experiment_name);
  }
}

// Takes a set of enabled lab experiments
void SetEnabledLabs(
    PrefService* prefs, const std::set<std::string>& enabled_experiments) {
  ListValue* experiments_list = prefs->GetMutableList(
      prefs::kEnabledLabsExperiments);
  if (!experiments_list)
    return;

  experiments_list->Clear();
  for (std::set<std::string>::const_iterator it = enabled_experiments.begin();
       it != enabled_experiments.end();
       ++it) {
    experiments_list->Append(new StringValue(*it));
  }
}

// Removes all experiments from prefs::kEnabledLabsExperiments that are
// unknown, to prevent this list to become very long as experiments are added
// and removed.
void SanitizeList(PrefService* prefs) {
  std::set<std::string> known_experiments;
  for (size_t i = 0; i < arraysize(kExperiments); ++i)
    known_experiments.insert(kExperiments[i].internal_name);

  std::set<std::string> enabled_experiments;
  GetEnabledLabs(prefs, &enabled_experiments);

  std::set<std::string> new_enabled_experiments;
  std::set_intersection(
      known_experiments.begin(), known_experiments.end(),
      enabled_experiments.begin(), enabled_experiments.end(),
      std::inserter(new_enabled_experiments, new_enabled_experiments.begin()));

  SetEnabledLabs(prefs, new_enabled_experiments);
}

void GetSanitizedEnabledLabs(
    PrefService* prefs, std::set<std::string>* result) {
  SanitizeList(prefs);
  GetEnabledLabs(prefs, result);
}

int GetCurrentPlatform() {
#if defined(OS_MACOSX)
  return kOsMac;
#elif defined(OS_WIN)
  return kOsWin;
#elif defined(OS_LINUX)
  return kOsLinux;
#else
#error Unknown platform
#endif
}

bool IsEnabled() {
#if defined(OS_CHROMEOS)
  // ChromeOS uses a different mechanism for about:labs; integrated with their
  // dom ui options.
  return false;
#elif defined(GOOGLE_CHROME_BUILD)
  // Don't enable this on the stable channel.
  return !platform_util::GetVersionStringModifier().empty();
#else
  return true;
#endif
}

void ConvertLabsToSwitches(Profile* profile, CommandLine* command_line) {
  // Do not activate labs features on the stable channel.
  if (!IsEnabled())
    return;

  std::set<std::string> enabled_experiments;
  GetSanitizedEnabledLabs(profile->GetPrefs(), &enabled_experiments);

  std::map<std::string, const Experiment*> experiments;
  for (size_t i = 0; i < arraysize(kExperiments); ++i)
    experiments[kExperiments[i].internal_name] = &kExperiments[i];

  for (std::set<std::string>::iterator it = enabled_experiments.begin();
       it != enabled_experiments.end();
       ++it) {
    const std::string& experiment_name = *it;
    std::map<std::string, const Experiment*>::iterator experiment =
        experiments.find(experiment_name);
    DCHECK(experiment != experiments.end());
    if (experiment == experiments.end())
      continue;

    command_line->AppendSwitch(experiment->second->command_line);
  }
}

ListValue* GetLabsExperimentsData(Profile* profile) {
  std::set<std::string> enabled_experiments;
  GetSanitizedEnabledLabs(profile->GetPrefs(), &enabled_experiments);

  int current_platform = GetCurrentPlatform();

  ListValue* experiments_data = new ListValue();
  for (size_t i = 0; i < arraysize(kExperiments); ++i) {
    const Experiment& experiment = kExperiments[i];
    if (!(experiment.supported_platforms & current_platform))
      continue;

    DictionaryValue* data = new DictionaryValue();
    data->SetString("internal_name", experiment.internal_name);
    data->SetString("name",
                    l10n_util::GetStringUTF16(experiment.visible_name_id));
    data->SetString("description",
                    l10n_util::GetStringUTF16(
                        experiment.visible_description_id));
    data->SetBoolean("enabled",
                      enabled_experiments.count(experiment.internal_name) > 0);

    experiments_data->Append(data);
  }
  return experiments_data;
}

static bool needs_restart_ = false;

bool IsRestartNeededToCommitChanges() {
  return needs_restart_;
}

void SetExperimentEnabled(
    Profile* profile, const std::string& internal_name, bool enable) {
  needs_restart_ = true;

  std::set<std::string> enabled_experiments;
  GetSanitizedEnabledLabs(profile->GetPrefs(), &enabled_experiments);

  if (enable)
    enabled_experiments.insert(internal_name);
  else
    enabled_experiments.erase(internal_name);

  SetEnabledLabs(profile->GetPrefs(), enabled_experiments);
}

}  // namespace Labs
<commit_msg>Enable about:labs on the stable channel as well. It's now enabled on all channels.<commit_after>// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/labs.h"

#include <algorithm>
#include <iterator>
#include <map>
#include <set>

#include "app/l10n_util.h"
#include "base/command_line.h"
#include "base/values.h"
#include "chrome/browser/prefs/pref_service.h"
#include "chrome/browser/profile.h"
#include "chrome/common/chrome_switches.h"
#include "chrome/common/pref_names.h"
#include "grit/generated_resources.h"

namespace about_labs {

enum { kOsMac = 1 << 0, kOsWin = 1 << 1, kOsLinux = 1 << 2 };

unsigned kOsAll = kOsMac | kOsWin | kOsLinux;

struct Experiment {
  // The internal name of the experiment. This is never shown to the user.
  // It _is_ however stored in the prefs file, so you shouldn't change the
  // name of existing labs.
  const char* internal_name;

  // String id of the message containing the experiment's name.
  int visible_name_id;

  // String id of the message containing the experiment's description.
  int visible_description_id;

  // The platforms the experiment is available on
  // Needs to be more than a compile-time #ifdef because of profile sync.
  unsigned supported_platforms;  // bitmask

  // The commandline parameter that's added when this lab is active. This is
  // different from |internal_name| so that the commandline flag can be
  // renamed without breaking the prefs file.
  const char* command_line;
};

const Experiment kExperiments[] = {
  {
    "expose-for-tabs",  // Do not change; see above.
    IDS_LABS_TABPOSE_NAME,
    IDS_LABS_TABPOSE_DESCRIPTION,
    kOsMac,
#if defined(OS_MACOSX)
    // The switch exists only on OS X.
    switches::kEnableExposeForTabs
#else
    ""
#endif
  },
  {
    "vertical-tabs",  // Do not change; see above.
    IDS_LABS_SIDE_TABS_NAME,
    IDS_LABS_SIDE_TABS_DESCRIPTION,
    kOsWin,
    switches::kEnableVerticalTabs
  },
  {
    "tabbed-options",  // Do not change; see above.
    IDS_LABS_TABBED_OPTIONS_NAME,
    IDS_LABS_TABBED_OPTIONS_DESCRIPTION,
    kOsAll,
    switches::kEnableTabbedOptions
  },
  {
    "match-preview",  // Do not change; see above.
    IDS_LABS_INSTANT_NAME,
    IDS_LABS_INSTANT_DESCRIPTION,
    kOsWin,
    switches::kEnableMatchPreview
  },
  {
    "remoting",  // Do not change; see above.
    IDS_LABS_REMOTING_NAME,
#if defined(OS_WIN)
    // Windows only supports host functionality at the moment.
    IDS_LABS_REMOTING_HOST_DESCRIPTION,
#elif defined(OS_LINUX)
    // Linux only supports client functionality at the moment.
    IDS_LABS_REMOTING_CLIENT_DESCRIPTION,
#else
    // On other platforms, this lab isn't available at all.
    0,
#endif
    kOsWin | kOsLinux,
    switches::kEnableRemoting
  },
  {
    "page-info-bubble",  // Do not change; see above.
    IDS_LABS_PAGE_INFO_BUBBLE_NAME,
    IDS_LABS_PAGE_INFO_BUBBLE_DESCRIPTION,
    kOsAll,
    switches::kEnableNewPageInfoBubble
  },
  {
    "disable-outdated-plugins",  // Do not change; see above.
    IDS_LABS_DISABLE_OUTDATED_PLUGINS_NAME,
    IDS_LABS_DISABLE_OUTDATED_PLUGINS_DESCRIPTION,
    kOsAll,
    switches::kDisableOutdatedPlugins
  },
  {
    "xss-auditor",  // Do not change; see above.
    IDS_LABS_XSS_AUDITOR_NAME,
    IDS_LABS_XSS_AUDITOR_DESCRIPTION,
    kOsAll,
    switches::kEnableXSSAuditor
  },
};

// Extracts the list of enabled lab experiments from a profile and stores them
// in a set.
void GetEnabledLabs(const PrefService* prefs, std::set<std::string>* result) {
  const ListValue* enabled_experiments = prefs->GetList(
      prefs::kEnabledLabsExperiments);
  if (!enabled_experiments)
    return;

  for (ListValue::const_iterator it = enabled_experiments->begin();
       it != enabled_experiments->end();
       ++it) {
    std::string experiment_name;
    if (!(*it)->GetAsString(&experiment_name)) {
      LOG(WARNING) << "Invalid entry in " << prefs::kEnabledLabsExperiments;
      continue;
    }
    result->insert(experiment_name);
  }
}

// Takes a set of enabled lab experiments
void SetEnabledLabs(
    PrefService* prefs, const std::set<std::string>& enabled_experiments) {
  ListValue* experiments_list = prefs->GetMutableList(
      prefs::kEnabledLabsExperiments);
  if (!experiments_list)
    return;

  experiments_list->Clear();
  for (std::set<std::string>::const_iterator it = enabled_experiments.begin();
       it != enabled_experiments.end();
       ++it) {
    experiments_list->Append(new StringValue(*it));
  }
}

// Removes all experiments from prefs::kEnabledLabsExperiments that are
// unknown, to prevent this list to become very long as experiments are added
// and removed.
void SanitizeList(PrefService* prefs) {
  std::set<std::string> known_experiments;
  for (size_t i = 0; i < arraysize(kExperiments); ++i)
    known_experiments.insert(kExperiments[i].internal_name);

  std::set<std::string> enabled_experiments;
  GetEnabledLabs(prefs, &enabled_experiments);

  std::set<std::string> new_enabled_experiments;
  std::set_intersection(
      known_experiments.begin(), known_experiments.end(),
      enabled_experiments.begin(), enabled_experiments.end(),
      std::inserter(new_enabled_experiments, new_enabled_experiments.begin()));

  SetEnabledLabs(prefs, new_enabled_experiments);
}

void GetSanitizedEnabledLabs(
    PrefService* prefs, std::set<std::string>* result) {
  SanitizeList(prefs);
  GetEnabledLabs(prefs, result);
}

int GetCurrentPlatform() {
#if defined(OS_MACOSX)
  return kOsMac;
#elif defined(OS_WIN)
  return kOsWin;
#elif defined(OS_LINUX)
  return kOsLinux;
#else
#error Unknown platform
#endif
}

bool IsEnabled() {
#if defined(OS_CHROMEOS)
  // ChromeOS uses a different mechanism for about:labs; integrated with their
  // dom ui options.
  // TODO(thakis): Port about:labs to chromeos -- http://crbug.com/57634
  return false;
#else
  return true;
#endif
}

void ConvertLabsToSwitches(Profile* profile, CommandLine* command_line) {
  if (!IsEnabled())
    return;

  std::set<std::string> enabled_experiments;
  GetSanitizedEnabledLabs(profile->GetPrefs(), &enabled_experiments);

  std::map<std::string, const Experiment*> experiments;
  for (size_t i = 0; i < arraysize(kExperiments); ++i)
    experiments[kExperiments[i].internal_name] = &kExperiments[i];

  for (std::set<std::string>::iterator it = enabled_experiments.begin();
       it != enabled_experiments.end();
       ++it) {
    const std::string& experiment_name = *it;
    std::map<std::string, const Experiment*>::iterator experiment =
        experiments.find(experiment_name);
    DCHECK(experiment != experiments.end());
    if (experiment == experiments.end())
      continue;

    command_line->AppendSwitch(experiment->second->command_line);
  }
}

ListValue* GetLabsExperimentsData(Profile* profile) {
  std::set<std::string> enabled_experiments;
  GetSanitizedEnabledLabs(profile->GetPrefs(), &enabled_experiments);

  int current_platform = GetCurrentPlatform();

  ListValue* experiments_data = new ListValue();
  for (size_t i = 0; i < arraysize(kExperiments); ++i) {
    const Experiment& experiment = kExperiments[i];
    if (!(experiment.supported_platforms & current_platform))
      continue;

    DictionaryValue* data = new DictionaryValue();
    data->SetString("internal_name", experiment.internal_name);
    data->SetString("name",
                    l10n_util::GetStringUTF16(experiment.visible_name_id));
    data->SetString("description",
                    l10n_util::GetStringUTF16(
                        experiment.visible_description_id));
    data->SetBoolean("enabled",
                      enabled_experiments.count(experiment.internal_name) > 0);

    experiments_data->Append(data);
  }
  return experiments_data;
}

static bool needs_restart_ = false;

bool IsRestartNeededToCommitChanges() {
  return needs_restart_;
}

void SetExperimentEnabled(
    Profile* profile, const std::string& internal_name, bool enable) {
  needs_restart_ = true;

  std::set<std::string> enabled_experiments;
  GetSanitizedEnabledLabs(profile->GetPrefs(), &enabled_experiments);

  if (enable)
    enabled_experiments.insert(internal_name);
  else
    enabled_experiments.erase(internal_name);

  SetEnabledLabs(profile->GetPrefs(), enabled_experiments);
}

}  // namespace Labs
<|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/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.h"

#include <utility>

#include "absl/types/span.h"
#include "llvm/ADT/None.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/Support/ToolOutputFile.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Support/FileUtilities.h"  // from @llvm-project
#include "mlir/Transforms/ViewOpGraph.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
#include "tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h"
#include "tensorflow/compiler/mlir/lite/tf_tfl_passes.h"
#include "tensorflow/compiler/mlir/lite/tf_to_tfl_flatbuffer.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/graph_debug_info.pb.h"
#include "tensorflow/lite/toco/model_flags.pb.h"
#include "tensorflow/lite/toco/toco_flags.pb.h"
#include "tensorflow/lite/toco/types.pb.h"
#include "tensorflow/stream_executor/lib/statusor.h"

namespace tensorflow {

Status HandleInputOutputArraysWithModule(const toco::ModelFlags& model_flags,
                                         mlir::OwningModuleRef* module) {
  mlir::FuncOp entry_function = nullptr;
  for (auto func : module->get().getOps<mlir::FuncOp>()) {
    if (auto tf_attrs =
            func->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function")) {
      // TODO(jaesung): There could be multiple entry functions. Let's handle
      // such cases if there are any needs for that.
      if (entry_function != nullptr) {
        return errors::InvalidArgument(
            "There should be only one tf.entry_function");
      }
      entry_function = func;
    }
  }
  if (entry_function == nullptr) {
    return errors::InvalidArgument("no tf.entry_function found");
  }

  // Get the list of input Op names from the function attribute.
  mlir::DictionaryAttr tf_attrs =
      entry_function->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function");
  llvm::SmallVector<llvm::StringRef, 4> function_input_names;
  function_input_names.reserve(model_flags.input_arrays().size());
  auto input_attr = tf_attrs.get("inputs");
  if (!input_attr) {
    return errors::InvalidArgument("no inputs attribute found");
  }
  auto input_names = input_attr.cast<mlir::StringAttr>().getValue();
  input_names.split(function_input_names, ",");
  const int function_input_names_size = function_input_names.size();
  if (function_input_names_size != model_flags.input_arrays().size()) {
    return errors::InvalidArgument(
        "input array size mismatch: got ", function_input_names.size(),
        ", expected: ", model_flags.input_arrays().size());
  }
  llvm::StringSet<> function_input_names_set;
  function_input_names_set.insert(function_input_names.begin(),
                                  function_input_names.end());
  for (const auto& input_array : model_flags.input_arrays()) {
    if (function_input_names_set.count(input_array.name()) == 0) {
      return errors::InvalidArgument("input array name (", input_array.name(),
                                     ") does not exist in the given graph");
    }
  }

  // Get the list of output Op names from the function attribute.
  llvm::SmallVector<llvm::StringRef, 4> function_output_names;
  function_output_names.reserve(model_flags.output_arrays().size());
  auto output_attr = tf_attrs.get("outputs");
  if (!output_attr) {
    return errors::InvalidArgument("no outputs attribute found");
  }
  auto output_names = output_attr.cast<mlir::StringAttr>().getValue();
  output_names.split(function_output_names, ",");
  const int function_output_names_size = function_output_names.size();
  if (function_output_names_size != model_flags.output_arrays().size()) {
    return errors::InvalidArgument(
        "output array size mismatch: got ", function_output_names.size(),
        ", expected: ", model_flags.output_arrays().size());
  }
  llvm::StringSet<> function_output_names_set;
  function_output_names_set.insert(function_output_names.begin(),
                                   function_output_names.end());
  for (const auto& output_array : model_flags.output_arrays()) {
    if (function_output_names_set.count(output_array) == 0) {
      return errors::InvalidArgument("output array name (", output_array,
                                     ") does not exist in the given graph");
    }
  }
  return Status::OK();
}

Status ConvertSavedModelToTFLiteFlatBuffer(const toco::ModelFlags& model_flags,
                                           const toco::TocoFlags& toco_flags,
                                           string* result) {
  mlir::MLIRContext context;
  mlir::TFL::QuantizationSpecs quant_specs;

  // Parse input arrays.
  std::vector<string> node_names;
  std::vector<string> node_dtypes;
  std::vector<llvm::Optional<std::vector<int>>> node_shapes;
  std::vector<llvm::Optional<double>> node_mins;
  std::vector<llvm::Optional<double>> node_maxs;

  // Populate quantization specs.
  TF_RETURN_IF_ERROR(internal::PopulateQuantizationSpecs(
      model_flags, toco_flags, &quant_specs, &node_names, &node_dtypes,
      &node_shapes, &node_mins, &node_maxs));

  internal::WarningUnusedFlags(model_flags, toco_flags);

  // Register all custom ops, including user-specified custom ops.
  TF_RETURN_IF_ERROR(internal::RegisterAllCustomOps(toco_flags));

  auto& saved_model_tags = model_flags.saved_model_tags();
  auto& saved_model_exported_names = model_flags.saved_model_exported_names();
  std::unordered_set<std::string> tags(saved_model_tags.begin(),
                                       saved_model_tags.end());
  auto exported_names_in_vector = std::vector<std::string>(
      saved_model_exported_names.begin(), saved_model_exported_names.end());
  absl::Span<std::string> exported_names(exported_names_in_vector);

  if (exported_names.size() != 1) {
    return errors::Unimplemented("Only support a single exported name.");
  }

  tensorflow::GraphImportConfig specs;
  specs.upgrade_legacy = true;

  std::vector<std::string> custom_opdefs(toco_flags.custom_opdefs().begin(),
                                         toco_flags.custom_opdefs().end());
  TF_ASSIGN_OR_RETURN(auto module,
                      ImportSavedModel(model_flags.saved_model_dir(),
                                       model_flags.saved_model_version(), tags,
                                       absl::MakeSpan(custom_opdefs),
                                       exported_names, specs, &context));

  if (!model_flags.input_arrays().empty() ||
      !model_flags.output_arrays().empty()) {
    TF_RETURN_IF_ERROR(HandleInputOutputArraysWithModule(model_flags, &module));
  }

  mlir::TFL::PassConfig pass_config(quant_specs);
  bool emit_builtin_tflite_ops = !toco_flags.force_select_tf_ops();
  pass_config.emit_builtin_tflite_ops = emit_builtin_tflite_ops;
  pass_config.lower_tensor_list_ops = true;
  // Disable the unfolding of the 16x16 TF::BatchMatMulOp to avoid the
  // conversion to an unsupported 16x16 TFL::FullyConnectedOp.
  if (toco_flags.inference_type() == toco::IODataType::QUANTIZED_INT16) {
    pass_config.unfold_batch_matmul = false;
  }

  // TODO(b/153507667): Pass the session object when importing logic is removed.
  auto status = internal::ConvertMLIRToTFLiteFlatBuffer(
      toco_flags, std::move(module), pass_config, tags, result,
      /*session=*/llvm::None);
  return status;
}

}  // namespace tensorflow
<commit_msg>[lite] Update saved model conversion to set the enable_tflite_variable in pass config based on flags.<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/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.h"

#include <utility>

#include "absl/types/span.h"
#include "llvm/ADT/None.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/Support/ToolOutputFile.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Support/FileUtilities.h"  // from @llvm-project
#include "mlir/Transforms/ViewOpGraph.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
#include "tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h"
#include "tensorflow/compiler/mlir/lite/tf_tfl_passes.h"
#include "tensorflow/compiler/mlir/lite/tf_to_tfl_flatbuffer.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/graph_debug_info.pb.h"
#include "tensorflow/lite/toco/model_flags.pb.h"
#include "tensorflow/lite/toco/toco_flags.pb.h"
#include "tensorflow/lite/toco/types.pb.h"
#include "tensorflow/stream_executor/lib/statusor.h"

namespace tensorflow {

Status HandleInputOutputArraysWithModule(const toco::ModelFlags& model_flags,
                                         mlir::OwningModuleRef* module) {
  mlir::FuncOp entry_function = nullptr;
  for (auto func : module->get().getOps<mlir::FuncOp>()) {
    if (auto tf_attrs =
            func->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function")) {
      // TODO(jaesung): There could be multiple entry functions. Let's handle
      // such cases if there are any needs for that.
      if (entry_function != nullptr) {
        return errors::InvalidArgument(
            "There should be only one tf.entry_function");
      }
      entry_function = func;
    }
  }
  if (entry_function == nullptr) {
    return errors::InvalidArgument("no tf.entry_function found");
  }

  // Get the list of input Op names from the function attribute.
  mlir::DictionaryAttr tf_attrs =
      entry_function->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function");
  llvm::SmallVector<llvm::StringRef, 4> function_input_names;
  function_input_names.reserve(model_flags.input_arrays().size());
  auto input_attr = tf_attrs.get("inputs");
  if (!input_attr) {
    return errors::InvalidArgument("no inputs attribute found");
  }
  auto input_names = input_attr.cast<mlir::StringAttr>().getValue();
  input_names.split(function_input_names, ",");
  const int function_input_names_size = function_input_names.size();
  if (function_input_names_size != model_flags.input_arrays().size()) {
    return errors::InvalidArgument(
        "input array size mismatch: got ", function_input_names.size(),
        ", expected: ", model_flags.input_arrays().size());
  }
  llvm::StringSet<> function_input_names_set;
  function_input_names_set.insert(function_input_names.begin(),
                                  function_input_names.end());
  for (const auto& input_array : model_flags.input_arrays()) {
    if (function_input_names_set.count(input_array.name()) == 0) {
      return errors::InvalidArgument("input array name (", input_array.name(),
                                     ") does not exist in the given graph");
    }
  }

  // Get the list of output Op names from the function attribute.
  llvm::SmallVector<llvm::StringRef, 4> function_output_names;
  function_output_names.reserve(model_flags.output_arrays().size());
  auto output_attr = tf_attrs.get("outputs");
  if (!output_attr) {
    return errors::InvalidArgument("no outputs attribute found");
  }
  auto output_names = output_attr.cast<mlir::StringAttr>().getValue();
  output_names.split(function_output_names, ",");
  const int function_output_names_size = function_output_names.size();
  if (function_output_names_size != model_flags.output_arrays().size()) {
    return errors::InvalidArgument(
        "output array size mismatch: got ", function_output_names.size(),
        ", expected: ", model_flags.output_arrays().size());
  }
  llvm::StringSet<> function_output_names_set;
  function_output_names_set.insert(function_output_names.begin(),
                                   function_output_names.end());
  for (const auto& output_array : model_flags.output_arrays()) {
    if (function_output_names_set.count(output_array) == 0) {
      return errors::InvalidArgument("output array name (", output_array,
                                     ") does not exist in the given graph");
    }
  }
  return Status::OK();
}

Status ConvertSavedModelToTFLiteFlatBuffer(const toco::ModelFlags& model_flags,
                                           const toco::TocoFlags& toco_flags,
                                           string* result) {
  mlir::MLIRContext context;
  mlir::TFL::QuantizationSpecs quant_specs;

  // Parse input arrays.
  std::vector<string> node_names;
  std::vector<string> node_dtypes;
  std::vector<llvm::Optional<std::vector<int>>> node_shapes;
  std::vector<llvm::Optional<double>> node_mins;
  std::vector<llvm::Optional<double>> node_maxs;

  // Populate quantization specs.
  TF_RETURN_IF_ERROR(internal::PopulateQuantizationSpecs(
      model_flags, toco_flags, &quant_specs, &node_names, &node_dtypes,
      &node_shapes, &node_mins, &node_maxs));

  internal::WarningUnusedFlags(model_flags, toco_flags);

  // Register all custom ops, including user-specified custom ops.
  TF_RETURN_IF_ERROR(internal::RegisterAllCustomOps(toco_flags));

  auto& saved_model_tags = model_flags.saved_model_tags();
  auto& saved_model_exported_names = model_flags.saved_model_exported_names();
  std::unordered_set<std::string> tags(saved_model_tags.begin(),
                                       saved_model_tags.end());
  auto exported_names_in_vector = std::vector<std::string>(
      saved_model_exported_names.begin(), saved_model_exported_names.end());
  absl::Span<std::string> exported_names(exported_names_in_vector);

  if (exported_names.size() != 1) {
    return errors::Unimplemented("Only support a single exported name.");
  }

  tensorflow::GraphImportConfig specs;
  specs.upgrade_legacy = true;

  std::vector<std::string> custom_opdefs(toco_flags.custom_opdefs().begin(),
                                         toco_flags.custom_opdefs().end());
  TF_ASSIGN_OR_RETURN(auto module,
                      ImportSavedModel(model_flags.saved_model_dir(),
                                       model_flags.saved_model_version(), tags,
                                       absl::MakeSpan(custom_opdefs),
                                       exported_names, specs, &context));

  if (!model_flags.input_arrays().empty() ||
      !model_flags.output_arrays().empty()) {
    TF_RETURN_IF_ERROR(HandleInputOutputArraysWithModule(model_flags, &module));
  }

  mlir::TFL::PassConfig pass_config(quant_specs);
  bool emit_builtin_tflite_ops = !toco_flags.force_select_tf_ops();
  pass_config.emit_builtin_tflite_ops = emit_builtin_tflite_ops;
  pass_config.lower_tensor_list_ops = true;
  pass_config.enable_tflite_variables =
      toco_flags.enable_tflite_resource_variables();
  // Disable the unfolding of the 16x16 TF::BatchMatMulOp to avoid the
  // conversion to an unsupported 16x16 TFL::FullyConnectedOp.
  if (toco_flags.inference_type() == toco::IODataType::QUANTIZED_INT16) {
    pass_config.unfold_batch_matmul = false;
  }

  // TODO(b/153507667): Pass the session object when importing logic is removed.
  auto status = internal::ConvertMLIRToTFLiteFlatBuffer(
      toco_flags, std::move(module), pass_config, tags, result,
      /*session=*/llvm::None);
  return status;
}

}  // 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/stream_executor/tpu/tpu_initializer_helper.h"

#include <dirent.h>
#include <dlfcn.h>
#include <fcntl.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <fstream>
#include <string>
#include <utility>

#include "absl/strings/str_cat.h"
#include "absl/strings/str_split.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/compiler/xla/pjrt/c/pjrt_c_api_tpu.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/libtftpu.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/pjrt_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_api_dlsym_set_fn.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_ops_c_api.h"
#include "tensorflow/tsl/platform/errors.h"
#include "tensorflow/tsl/platform/logging.h"

#if !defined(PLATFORM_GOOGLE)
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform.h"
#include "tensorflow/tsl/platform/cloud/gcs_file_system.h"
#include "tensorflow/tsl/platform/env.h"
#elif defined(LIBTPU_STATIC)
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform.h"
#endif  // PLATFORM_GOOGLE

namespace tensorflow {
namespace tpu {
namespace {

static std::string GetEnvVar(const char* name) {
  // Constructing a std::string directly from nullptr is undefined behavior so
  // we can return empty string in that case
  const char* env_value = getenv(name);
  if (!env_value) return "";
  return std::string(env_value);
}

bool GetEnvBool(const char* name, bool defval) {
  const char* env = getenv(name);
  if (env == nullptr) {
    return defval;
  }
  if (std::strcmp(env, "true") == 0) {
    return true;
  }
  if (std::strcmp(env, "false") == 0) {
    return false;
  }
  int int_env;
  bool has_int = absl::SimpleAtoi(env, &int_env);
  return has_int && int_env != 0;
}

}  // namespace

// This function gets pid of a process and checks if that process is using tpu.
// It is not able to check processes that are owned by another user.
bool IsTpuUsed(int64_t pid) {
  std::string path = absl::StrCat("/proc/", pid, "/fd");
  DIR* raw_fd_dir = opendir(path.c_str());
  if (!raw_fd_dir) {
    return false;
  }
  std::unique_ptr<DIR, int (*)(DIR*)> fd_dir(raw_fd_dir, closedir);
  struct dirent* ent;
  std::string line;
  std::string tpu_dev_path = "/dev/accel0";
  line.resize(tpu_dev_path.size());
  while ((ent = readdir(raw_fd_dir))) {
    if (!isdigit(*ent->d_name)) continue;
    int64_t fd = strtol(ent->d_name, nullptr, 10);
    path = absl::StrCat("/proc/", pid, "/fd/", fd);
    if (!readlink(path.c_str(), &line[0], line.size())) continue;
    if (line != tpu_dev_path) continue;
    return true;
  }
  return false;
}

// This function iterates through all the processes in /proc and finds out if
// any process it was able to check is using the TPU. It does not have
// permission to processes owned by another user.
// TODO (shahrokhi) use tensorflow/core/platform/filesystem (GetChildren) for
// this.
stream_executor::port::StatusOr<int64_t> FindLibtpuProcess() {
  DIR* proc = opendir("/proc");

  if (proc == nullptr) {
    return tsl::errors::Unavailable("was not able to open /proc");
  }
  std::unique_ptr<DIR, int (*)(DIR*)> proc_dir(proc, closedir);
  struct dirent* ent;
  int64_t pid;
  while ((ent = readdir(proc))) {
    if (!isdigit(*ent->d_name)) continue;

    pid = strtol(ent->d_name, nullptr, 10);
    if (IsTpuUsed(pid)) {
      return pid;
    }
  }
  return tsl::errors::NotFound("did not find which pid uses the libtpu.so");
}

stream_executor::port::Status TryAcquireTpuLock() {
  static absl::Mutex* mu = new absl::Mutex();
  absl::MutexLock l(mu);

  // TODO(skyewm): use `absl::StrCat(getenv(name))` once we build with the
  // fix for https://github.com/abseil/abseil-cpp/issues/1167.
  std::string load_library_override;
  const char* env_value = getenv("TPU_LOAD_LIBRARY");
  if (env_value != nullptr) {
    load_library_override = std::string(env_value);
  }

  if (load_library_override == "1") {
    return ::tsl::OkStatus();
  } else if (load_library_override == "0") {
    return tsl::errors::FailedPrecondition(
        "TPU_LOAD_LIBRARY=0, not loading libtpu");
  }

  // If TPU_CHIPS_PER_PROCESS_BOUNDS doesn't include all chips, we assume
  // we're using different chips in different processes and thus multiple
  // libtpu loads are ok.
  // TODO(skyewm): we could make per-chip lock files and look at
  // TPU_VISIBLE_DEVICES if we wanted to make this really precise.
  std::string chips_per_process_bounds =
      GetEnvVar("TPU_CHIPS_PER_PROCESS_BOUNDS");
  bool allow_multiple_libtpu_load =
      GetEnvBool("ALLOW_MULTIPLE_LIBTPU_LOAD", false);
  // TODO(skyewm): remove this when TPU_CHIPS_PER_HOST_BOUNDS is fully
  // deprecated
  if (chips_per_process_bounds.empty()) {
    chips_per_process_bounds = GetEnvVar("TPU_CHIPS_PER_HOST_BOUNDS");
  }
  if ((chips_per_process_bounds.empty() ||
       chips_per_process_bounds == "2,2,1") &&
      !allow_multiple_libtpu_load) {
    int fd = open("/tmp/libtpu_lockfile", O_CREAT | O_RDWR, 0644);

    // This lock is held until the process exits intentionally. The underlying
    // TPU device will be held on until it quits.
    if (lockf(fd, F_TLOCK, 0) != 0) {
      auto pid = FindLibtpuProcess();
      if (pid.ok()) {
        return tsl::errors::Aborted(absl::StrCat(
            "The TPU is already in use by process with pid ", pid.value(),
            ". Not attempting to load libtpu.so in this process."));
      } else {
        return tsl::errors::Aborted(
            "The TPU is already in use by another process probably owned by "
            "another user. Run \"$ sudo lsof -w /dev/accel0\" to figure out "
            "which process is using the TPU. If you still get this message, "
            "run \"$ sudo rm /tmp/libtpu_lockfile\". Not attempting to load "
            "libtpu.so in this process.");
      }
    } else {
      return ::tsl::OkStatus();
    }
  } else {
    VLOG(1) << "TPU_CHIPS_PER_PROCESS_BOUNDS is not empty or "
               "ALLOW_MULTIPLE_LIBTPU_LOAD is set to True, "
               "therefore allowing multiple libtpu.so loads.";
    return ::tsl::OkStatus();
  }
}
#if !defined(PLATFORM_GOOGLE)
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_library_init_fns.inc"

stream_executor::port::Status InitializeTpuLibrary(void* library_handle) {
  stream_executor::port::Status s = InitializeTpuStructFns(library_handle);

  // Retrieve arguments from environment if applicable
  std::pair<std::vector<std::string>, std::vector<const char*>> args =
      GetLibTpuInitArguments();

  // TPU platform registration must only be performed after the library is
  // loaded. We do not want to register a TPU platform in XLA without the
  // supporting library providing the necessary APIs.
  if (s.ok()) {
    void (*initialize_fn)(bool init_library, int num_args, const char** args);
    initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
        dlsym(library_handle, "TfTpu_Initialize"));
    (*initialize_fn)(/*init_library=*/true, args.second.size(),
                     args.second.data());

    RegisterTpuPlatform();
  }

  return s;
}

typedef const PJRT_Api* (*PjRtFuncPtr)();
void InitializePjRt(void* library_handle) {
  PjRtFuncPtr fptr = &GetTpuPjrtApi;
  *reinterpret_cast<void**>(&fptr) = dlsym(library_handle, "GetTpuPjrtApi");
  if (fptr == nullptr) {
    LOG(INFO) << "GetTpuPjrtApi not found";
  } else {
    LOG(INFO) << "GetTpuPjrtApi was found";
    tensorflow::tpu::SetPjrtApi("TPU", fptr());
  }
}

namespace {
void* CreateGcsFilesystemFn() { return new tsl::RetryingGcsFileSystem(); }

// This is a temporary fix for including GCS file system on TPU builds.
// Will be removed once b/176954917 is fully resolved with the build fix.
void InitializeCreateGcsFileSystemFnPtr() {
  int fd = shm_open(absl::StrCat("/tmp_tf_gcs_fs_pointer_", getpid()).data(),
                    O_CREAT | O_RDWR, S_IRUSR | S_IWUSR);
  if (fd == -1) {
    LOG(ERROR) << "Unable to open shared memory for GCS file system creator.";
    return;
  }

  if (ftruncate(fd, sizeof(tsl::FileSystem*)) == -1) {
    LOG(ERROR)
        << "Unable to allocate shared memory for GCS file system creator.";
    return;
  }

  void* (**fn)() = reinterpret_cast<void* (**)()>(mmap(
      NULL, sizeof(void* (*)()), PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0));
  if (fn == MAP_FAILED) {
    LOG(ERROR) << "Cannot mmap shared memory for GCS file system creator.";
    return;
  }

  *fn = &CreateGcsFilesystemFn;

  munmap(fn, sizeof(void* (*)()));
  close(fd);

  // Clean up shared memory on a clean exit.
  atexit([]() {
    shm_unlink(absl::StrCat("/tmp_tf_gcs_fs_pointer_", getpid()).data());
  });
}
}  // namespace
stream_executor::port::Status FindAndLoadTpuLibrary() {
  const char* env_value = getenv("TPU_LIBRARY_PATH");
  const char* libtpu_path =
      env_value && strlen(env_value) > 0 ? env_value : "libtpu.so";
  LOG(INFO) << "Libtpu path is: " << libtpu_path;
  void* library = dlopen(libtpu_path, RTLD_NOW);
  if (library) {
    // We can open the shared library which means we are in a TPU environment.
    // Try to acquire exclusive access.
    TF_RETURN_IF_ERROR(TryAcquireTpuLock());
    TF_RETURN_IF_ERROR(InitializeTpuLibrary(library));
    InitializePjRt(library);
  }

  InitializeCreateGcsFileSystemFnPtr();
  return ::tsl::OkStatus();
}

#elif defined(LIBTPU_STATIC)

#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_library_init_fns.inc"

stream_executor::port::Status InitializeTpuLibrary() {
  // Retrieve arguments from environment if applicable
  std::pair<std::vector<std::string>, std::vector<const char*>> args =
      GetLibTpuInitArguments();

  TfTpu_Initialize(/*init_library*/ true, args.second.size(),
                   args.second.data());

  RegisterTpuPlatform();
  return ::tsl::OkStatus();
}

stream_executor::port::Status FindAndLoadTpuLibrary() {
  // We can open the shared library which means we are in a TPU environment.
  // Try to acquire exclusive access.
  TF_RETURN_IF_ERROR(TryAcquireTpuLock());
  TF_RETURN_IF_ERROR(InitializeTpuLibrary());
  return ::tsl::OkStatus();
}

#else   // PLATFORM_GOOGLE
stream_executor::port::Status InitializeTpuLibrary(void* library_handle) {
  return tsl::errors::Unimplemented(
      "You must statically link in a TPU library.");
}
#endif  // PLATFORM_GOOGLE
std::pair<std::vector<std::string>, std::vector<const char*>>
GetLibTpuInitArguments() {
  // We make copies of the arguments returned by getenv because the memory
  // returned may be altered or invalidated by further calls to getenv.
  std::vector<std::string> args;
  std::vector<const char*> arg_ptrs;

  // Retrieve arguments from environment if applicable.
  char* env = getenv("LIBTPU_INIT_ARGS");
  if (env != nullptr) {
    // TODO(frankchn): Handles quotes properly if necessary.
    args = absl::StrSplit(env, ' ');
  }

  arg_ptrs.reserve(args.size());
  for (int i = 0; i < args.size(); ++i) {
    arg_ptrs.push_back(args[i].data());
  }

  return {std::move(args), std::move(arg_ptrs)};
}

}  // namespace tpu
}  // namespace tensorflow
<commit_msg>Rename InitializePjRt to MaybeInitializePjRt as it does not fail if PJRT is not loaded.<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/stream_executor/tpu/tpu_initializer_helper.h"

#include <dirent.h>
#include <dlfcn.h>
#include <fcntl.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <fstream>
#include <string>
#include <utility>

#include "absl/strings/str_cat.h"
#include "absl/strings/str_split.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/compiler/xla/pjrt/c/pjrt_c_api_tpu.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/libtftpu.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/pjrt_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_api_dlsym_set_fn.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_ops_c_api.h"
#include "tensorflow/tsl/platform/errors.h"
#include "tensorflow/tsl/platform/logging.h"

#if !defined(PLATFORM_GOOGLE)
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform.h"
#include "tensorflow/tsl/platform/cloud/gcs_file_system.h"
#include "tensorflow/tsl/platform/env.h"
#elif defined(LIBTPU_STATIC)
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform.h"
#endif  // PLATFORM_GOOGLE

namespace tensorflow {
namespace tpu {
namespace {

static std::string GetEnvVar(const char* name) {
  // Constructing a std::string directly from nullptr is undefined behavior so
  // we can return empty string in that case
  const char* env_value = getenv(name);
  if (!env_value) return "";
  return std::string(env_value);
}

bool GetEnvBool(const char* name, bool defval) {
  const char* env = getenv(name);
  if (env == nullptr) {
    return defval;
  }
  if (std::strcmp(env, "true") == 0) {
    return true;
  }
  if (std::strcmp(env, "false") == 0) {
    return false;
  }
  int int_env;
  bool has_int = absl::SimpleAtoi(env, &int_env);
  return has_int && int_env != 0;
}

}  // namespace

// This function gets pid of a process and checks if that process is using tpu.
// It is not able to check processes that are owned by another user.
bool IsTpuUsed(int64_t pid) {
  std::string path = absl::StrCat("/proc/", pid, "/fd");
  DIR* raw_fd_dir = opendir(path.c_str());
  if (!raw_fd_dir) {
    return false;
  }
  std::unique_ptr<DIR, int (*)(DIR*)> fd_dir(raw_fd_dir, closedir);
  struct dirent* ent;
  std::string line;
  std::string tpu_dev_path = "/dev/accel0";
  line.resize(tpu_dev_path.size());
  while ((ent = readdir(raw_fd_dir))) {
    if (!isdigit(*ent->d_name)) continue;
    int64_t fd = strtol(ent->d_name, nullptr, 10);
    path = absl::StrCat("/proc/", pid, "/fd/", fd);
    if (!readlink(path.c_str(), &line[0], line.size())) continue;
    if (line != tpu_dev_path) continue;
    return true;
  }
  return false;
}

// This function iterates through all the processes in /proc and finds out if
// any process it was able to check is using the TPU. It does not have
// permission to processes owned by another user.
// TODO (shahrokhi) use tensorflow/core/platform/filesystem (GetChildren) for
// this.
stream_executor::port::StatusOr<int64_t> FindLibtpuProcess() {
  DIR* proc = opendir("/proc");

  if (proc == nullptr) {
    return tsl::errors::Unavailable("was not able to open /proc");
  }
  std::unique_ptr<DIR, int (*)(DIR*)> proc_dir(proc, closedir);
  struct dirent* ent;
  int64_t pid;
  while ((ent = readdir(proc))) {
    if (!isdigit(*ent->d_name)) continue;

    pid = strtol(ent->d_name, nullptr, 10);
    if (IsTpuUsed(pid)) {
      return pid;
    }
  }
  return tsl::errors::NotFound("did not find which pid uses the libtpu.so");
}

stream_executor::port::Status TryAcquireTpuLock() {
  static absl::Mutex* mu = new absl::Mutex();
  absl::MutexLock l(mu);

  // TODO(skyewm): use `absl::StrCat(getenv(name))` once we build with the
  // fix for https://github.com/abseil/abseil-cpp/issues/1167.
  std::string load_library_override;
  const char* env_value = getenv("TPU_LOAD_LIBRARY");
  if (env_value != nullptr) {
    load_library_override = std::string(env_value);
  }

  if (load_library_override == "1") {
    return ::tsl::OkStatus();
  } else if (load_library_override == "0") {
    return tsl::errors::FailedPrecondition(
        "TPU_LOAD_LIBRARY=0, not loading libtpu");
  }

  // If TPU_CHIPS_PER_PROCESS_BOUNDS doesn't include all chips, we assume
  // we're using different chips in different processes and thus multiple
  // libtpu loads are ok.
  // TODO(skyewm): we could make per-chip lock files and look at
  // TPU_VISIBLE_DEVICES if we wanted to make this really precise.
  std::string chips_per_process_bounds =
      GetEnvVar("TPU_CHIPS_PER_PROCESS_BOUNDS");
  bool allow_multiple_libtpu_load =
      GetEnvBool("ALLOW_MULTIPLE_LIBTPU_LOAD", false);
  // TODO(skyewm): remove this when TPU_CHIPS_PER_HOST_BOUNDS is fully
  // deprecated
  if (chips_per_process_bounds.empty()) {
    chips_per_process_bounds = GetEnvVar("TPU_CHIPS_PER_HOST_BOUNDS");
  }
  if ((chips_per_process_bounds.empty() ||
       chips_per_process_bounds == "2,2,1") &&
      !allow_multiple_libtpu_load) {
    int fd = open("/tmp/libtpu_lockfile", O_CREAT | O_RDWR, 0644);

    // This lock is held until the process exits intentionally. The underlying
    // TPU device will be held on until it quits.
    if (lockf(fd, F_TLOCK, 0) != 0) {
      auto pid = FindLibtpuProcess();
      if (pid.ok()) {
        return tsl::errors::Aborted(absl::StrCat(
            "The TPU is already in use by process with pid ", pid.value(),
            ". Not attempting to load libtpu.so in this process."));
      } else {
        return tsl::errors::Aborted(
            "The TPU is already in use by another process probably owned by "
            "another user. Run \"$ sudo lsof -w /dev/accel0\" to figure out "
            "which process is using the TPU. If you still get this message, "
            "run \"$ sudo rm /tmp/libtpu_lockfile\". Not attempting to load "
            "libtpu.so in this process.");
      }
    } else {
      return ::tsl::OkStatus();
    }
  } else {
    VLOG(1) << "TPU_CHIPS_PER_PROCESS_BOUNDS is not empty or "
               "ALLOW_MULTIPLE_LIBTPU_LOAD is set to True, "
               "therefore allowing multiple libtpu.so loads.";
    return ::tsl::OkStatus();
  }
}
#if !defined(PLATFORM_GOOGLE)
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_library_init_fns.inc"

stream_executor::port::Status InitializeTpuLibrary(void* library_handle) {
  stream_executor::port::Status s = InitializeTpuStructFns(library_handle);

  // Retrieve arguments from environment if applicable
  std::pair<std::vector<std::string>, std::vector<const char*>> args =
      GetLibTpuInitArguments();

  // TPU platform registration must only be performed after the library is
  // loaded. We do not want to register a TPU platform in XLA without the
  // supporting library providing the necessary APIs.
  if (s.ok()) {
    void (*initialize_fn)(bool init_library, int num_args, const char** args);
    initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
        dlsym(library_handle, "TfTpu_Initialize"));
    (*initialize_fn)(/*init_library=*/true, args.second.size(),
                     args.second.data());

    RegisterTpuPlatform();
  }

  return s;
}

typedef const PJRT_Api* (*PjRtFuncPtr)();
void MaybeInitializePjRt(void* library_handle) {
  PjRtFuncPtr fptr = &GetTpuPjrtApi;
  *reinterpret_cast<void**>(&fptr) = dlsym(library_handle, "GetTpuPjrtApi");
  if (fptr == nullptr) {
    LOG(INFO) << "GetTpuPjrtApi not found. PjrtApi will not be used.";
  } else {
    LOG(INFO) << "GetTpuPjrtApi was found";
    tensorflow::tpu::SetPjrtApi("TPU", fptr());
  }
}

namespace {
void* CreateGcsFilesystemFn() { return new tsl::RetryingGcsFileSystem(); }

// This is a temporary fix for including GCS file system on TPU builds.
// Will be removed once b/176954917 is fully resolved with the build fix.
void InitializeCreateGcsFileSystemFnPtr() {
  int fd = shm_open(absl::StrCat("/tmp_tf_gcs_fs_pointer_", getpid()).data(),
                    O_CREAT | O_RDWR, S_IRUSR | S_IWUSR);
  if (fd == -1) {
    LOG(ERROR) << "Unable to open shared memory for GCS file system creator.";
    return;
  }

  if (ftruncate(fd, sizeof(tsl::FileSystem*)) == -1) {
    LOG(ERROR)
        << "Unable to allocate shared memory for GCS file system creator.";
    return;
  }

  void* (**fn)() = reinterpret_cast<void* (**)()>(mmap(
      NULL, sizeof(void* (*)()), PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0));
  if (fn == MAP_FAILED) {
    LOG(ERROR) << "Cannot mmap shared memory for GCS file system creator.";
    return;
  }

  *fn = &CreateGcsFilesystemFn;

  munmap(fn, sizeof(void* (*)()));
  close(fd);

  // Clean up shared memory on a clean exit.
  atexit([]() {
    shm_unlink(absl::StrCat("/tmp_tf_gcs_fs_pointer_", getpid()).data());
  });
}
}  // namespace
stream_executor::port::Status FindAndLoadTpuLibrary() {
  const char* env_value = getenv("TPU_LIBRARY_PATH");
  const char* libtpu_path =
      env_value && strlen(env_value) > 0 ? env_value : "libtpu.so";
  LOG(INFO) << "Libtpu path is: " << libtpu_path;
  void* library = dlopen(libtpu_path, RTLD_NOW);
  if (library) {
    // We can open the shared library which means we are in a TPU environment.
    // Try to acquire exclusive access.
    TF_RETURN_IF_ERROR(TryAcquireTpuLock());
    TF_RETURN_IF_ERROR(InitializeTpuLibrary(library));
    MaybeInitializePjRt(library);
  }

  InitializeCreateGcsFileSystemFnPtr();
  return ::tsl::OkStatus();
}

#elif defined(LIBTPU_STATIC)

#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_library_init_fns.inc"

stream_executor::port::Status InitializeTpuLibrary() {
  // Retrieve arguments from environment if applicable
  std::pair<std::vector<std::string>, std::vector<const char*>> args =
      GetLibTpuInitArguments();

  TfTpu_Initialize(/*init_library*/ true, args.second.size(),
                   args.second.data());

  RegisterTpuPlatform();
  return ::tsl::OkStatus();
}

stream_executor::port::Status FindAndLoadTpuLibrary() {
  // We can open the shared library which means we are in a TPU environment.
  // Try to acquire exclusive access.
  TF_RETURN_IF_ERROR(TryAcquireTpuLock());
  TF_RETURN_IF_ERROR(InitializeTpuLibrary());
  return ::tsl::OkStatus();
}

#else   // PLATFORM_GOOGLE
stream_executor::port::Status InitializeTpuLibrary(void* library_handle) {
  return tsl::errors::Unimplemented(
      "You must statically link in a TPU library.");
}
#endif  // PLATFORM_GOOGLE
std::pair<std::vector<std::string>, std::vector<const char*>>
GetLibTpuInitArguments() {
  // We make copies of the arguments returned by getenv because the memory
  // returned may be altered or invalidated by further calls to getenv.
  std::vector<std::string> args;
  std::vector<const char*> arg_ptrs;

  // Retrieve arguments from environment if applicable.
  char* env = getenv("LIBTPU_INIT_ARGS");
  if (env != nullptr) {
    // TODO(frankchn): Handles quotes properly if necessary.
    args = absl::StrSplit(env, ' ');
  }

  arg_ptrs.reserve(args.size());
  for (int i = 0; i < args.size(); ++i) {
    arg_ptrs.push_back(args[i].data());
  }

  return {std::move(args), std::move(arg_ptrs)};
}

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