AnnotateInternalFunctionsPass.h

Go to the documentation of this file.

/*
 * Copyright 2022 HEAVY.AI, Inc.
 *
 * 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 <llvm/IR/Function.h>

#include <llvm/Analysis/CallGraph.h>
#include <llvm/Analysis/CallGraphSCCPass.h>

#include <llvm/Pass.h>
#include <llvm/Support/raw_ostream.h>

#include "Logger/Logger.h"

class AnnotateInternalFunctionsPass : public llvm::CallGraphSCCPass {
 public:
  static char ID;
  AnnotateInternalFunctionsPass() : CallGraphSCCPass(ID) {}

  bool runOnSCC(llvm::CallGraphSCC& SCC) override {
    bool updated_function_defs = false;

    // iterate the call graph
    for (auto& node : SCC) {
      CHECK(node);
      auto fcn = node->getFunction();
      if (!fcn) {
        continue;
      }
      if (isInternalStatelessFunction(fcn->getName()) ||
          isInternalMathFunction(fcn->getName())) {
        updated_function_defs = true;
        std::vector<llvm::Attribute::AttrKind> attrs{llvm::Attribute::NoFree,
                                                     llvm::Attribute::NoSync,
                                                     llvm::Attribute::NoUnwind,
                                                     llvm::Attribute::WillReturn,
                                                     llvm::Attribute::ReadNone,
                                                     llvm::Attribute::Speculatable};
        // WriteOnly is automatically added to all math functions in llvm 14.0
        // https://reviews.llvm.org/D116426 which is incompatible with ReadNone.
        fcn->removeFnAttr(llvm::Attribute::WriteOnly);
        for (const auto& attr : attrs) {
          fcn->addFnAttr(attr);
        }
      }
    }

    return updated_function_defs;
  }

  llvm::StringRef getPassName() const override { return "AnnotateInternalFunctionsPass"; }

 private:
  static const std::set<std::string> extension_functions;

  static bool isInternalStatelessFunction(const llvm::StringRef& func_name) {
    // extension functions or non-inlined builtins which do not modify any state
    return extension_functions.count(func_name.str()) > 0;
  }

  static const std::set<std::string> math_builtins;

  static bool isInternalMathFunction(const llvm::StringRef& func_name) {
    // include all math functions from ExtensionFunctions.hpp
    return math_builtins.count(func_name.str()) > 0;
  }
};

char AnnotateInternalFunctionsPass::ID = 0;

const std::set<std::string> AnnotateInternalFunctionsPass::extension_functions =
    std::set<std::string>{"point_coord_array_is_null",
                          "decompress_x_coord_geoint",
                          "decompress_y_coord_geoint",
                          "compress_x_coord_geoint",
                          "compress_y_coord_geoint",
                          // GeoOpsRuntime.cpp
                          "transform_4326_900913_x",
                          "transform_4326_900913_y",
                          "transform_900913_4326_x",
                          "transform_900913_4326_y",
                          // ExtensionFunctions.hpp
                          "conv_4326_900913_x",
                          "conv_4326_900913_y",
                          "distance_in_meters",
                          "approx_distance_in_meters",
                          "rect_pixel_bin_x",
                          "rect_pixel_bin_y",
                          "rect_pixel_bin_packed",
                          "reg_hex_horiz_pixel_bin_x",
                          "reg_hex_horiz_pixel_bin_y",
                          "reg_hex_horiz_pixel_bin_packed",
                          "reg_hex_vert_pixel_bin_x",
                          "reg_hex_vert_pixel_bin_y",
                          "reg_hex_vert_pixel_bin_packed",
                          "convert_meters_to_merc_pixel_width",
                          "convert_meters_to_merc_pixel_height",
                          "is_point_in_merc_view",
                          "is_point_size_in_merc_view"};

// TODO: consider either adding specializations here for the `__X` versions (for different
// types), or just truncate the function name removing the underscores in
// `isInternalMathFunction`.
const std::set<std::string> AnnotateInternalFunctionsPass::math_builtins =
    std::set<std::string>{"Acos",  "Asin",    "Atan", "Atan2",    "Ceil",    "Cos",
                          "Cot",   "degrees", "Exp",  "Floor",    "ln",      "Log",
                          "Log10", "log",     "pi",   "power",    "radians", "Round",
                          "Sin",   "Tan",     "tan",  "Truncate", "is_nan",  "is_inf"};