<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 <string>

#include "llvm/Support/FormatVariadic.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"  // from @llvm-project
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"

namespace mlir {
namespace kernel_gen {
namespace tf_framework {
namespace {

using LLVM::LLVMFuncOp;

static constexpr StringRef kCInterfaceAlloc = "_mlir_ciface_tf_alloc";
static constexpr StringRef kCInterfaceDealloc = "_mlir_ciface_tf_dealloc";
static constexpr StringRef kCInterfaceReportError =
    "_mlir_ciface_tf_report_error";
static constexpr StringRef kCInterfaceJITCompile =
    "_mlir_ciface_tf_jit_compile";
static constexpr StringRef kCInterfaceJITExecute =
    "_mlir_ciface_tf_jit_execute";
static constexpr StringRef kJITCodeGlobalBaseName = "jit_module_code";
static constexpr StringRef kErrorMessageGlobalBaseName = "error_message";

Value CreateOrFindGlobalStringConstant(Location loc, OpBuilder &builder,
                                       StringRef base_name, StringRef str) {
  auto module =
      builder.getInsertionBlock()->getParentOp()->getParentOfType<ModuleOp>();
  std::string global_name =
      llvm::formatv("{0}_{1}", base_name, llvm::hash_value(str));
  Operation *global_constant =
      SymbolTable::lookupNearestSymbolFrom(module, global_name);
  if (global_constant) {
    Value global_ptr = builder.create<LLVM::AddressOfOp>(
        loc, cast<LLVM::GlobalOp>(global_constant));
    Value c0 = builder.create<LLVM::ConstantOp>(loc, builder.getI64Type(),
                                                builder.getIndexAttr(0));
    return builder.create<LLVM::GEPOp>(
        loc, LLVM::LLVMPointerType::get(builder.getIntegerType(8)), global_ptr,
        ValueRange{c0, c0});
  }
  return LLVM::createGlobalString(loc, builder, global_name, str,
                                  LLVM::Linkage::Internal);
}

/// Base class for patterns converting TF Framework ops to function calls.
template <typename OpTy>
class ConvertToLLVMCallOpPattern : public ConvertOpToLLVMPattern<OpTy> {
 public:
  using ConvertOpToLLVMPattern<OpTy>::ConvertOpToLLVMPattern;

  // Attempts to find function symbol in the module, adds it if not found.
  FlatSymbolRefAttr getOrInsertTFFunction(PatternRewriter &rewriter,
                                          Operation *op) const {
    ModuleOp module = op->getParentOfType<ModuleOp>();
    StringRef tf_func_name = GetFuncName();
    auto tf_func = module.lookupSymbol<LLVMFuncOp>(tf_func_name);
    if (!tf_func) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(module.getBody());
      auto func_type = GetFuncType();
      tf_func = rewriter.create<LLVMFuncOp>(rewriter.getUnknownLoc(),
                                            tf_func_name, func_type);
    }
    return SymbolRefAttr::get(rewriter.getContext(), tf_func_name);
  }

 protected:
  virtual StringRef GetFuncName() const = 0;
  virtual Type GetFuncType() const = 0;

  std::pair<Value, Value> ConvertIntegerArrayAttrToStackAllocatedArray(
      Location loc, Type size_ty, Type element_ty,
      llvm::Optional<ArrayAttr> attr,
      ConversionPatternRewriter *rewriter) const {
    assert(size_ty.isa<IntegerType>() && "expect integer size type");
    assert(element_ty.isa<IntegerType>() && "expect integer element type");
    Type element_ptr_ty = LLVM::LLVMPointerType::get(element_ty);

    // If the attribute is missing or empty, set the element count to 0 and
    // return NULL.
    if (!attr.hasValue() || attr.getValue().empty()) {
      Value zero = rewriter->create<LLVM::ConstantOp>(
          loc, size_ty, rewriter->getIntegerAttr(size_ty, 0));
      Value null_ptr = rewriter->create<LLVM::NullOp>(loc, element_ptr_ty);
      return std::make_pair(zero, null_ptr);
    }

    // Allocate array to store the elements.
    auto &array_attr = attr.getValue();<commit_msg>[KernelGen][JIT] Move CreateOrFindGlobalStringConstant helper into the pattern<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 <string>

#include "llvm/Support/FormatVariadic.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"  // from @llvm-project
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"

namespace mlir {
namespace kernel_gen {
namespace tf_framework {
namespace {

using LLVM::LLVMFuncOp;

static constexpr StringRef kCInterfaceAlloc = "_mlir_ciface_tf_alloc";
static constexpr StringRef kCInterfaceDealloc = "_mlir_ciface_tf_dealloc";
static constexpr StringRef kCInterfaceReportError =
    "_mlir_ciface_tf_report_error";
static constexpr StringRef kCInterfaceJITCompile =
    "_mlir_ciface_tf_jit_compile";
static constexpr StringRef kCInterfaceJITExecute =
    "_mlir_ciface_tf_jit_execute";
static constexpr StringRef kJITCodeGlobalBaseName = "jit_module_code";
static constexpr StringRef kErrorMessageGlobalBaseName = "error_message";

/// Base class for patterns converting TF Framework ops to function calls.
template <typename OpTy>
class ConvertToLLVMCallOpPattern : public ConvertOpToLLVMPattern<OpTy> {
 public:
  using ConvertOpToLLVMPattern<OpTy>::ConvertOpToLLVMPattern;

  // Attempts to find function symbol in the module, adds it if not found.
  FlatSymbolRefAttr getOrInsertTFFunction(PatternRewriter &rewriter,
                                          Operation *op) const {
    ModuleOp module = op->getParentOfType<ModuleOp>();
    StringRef tf_func_name = GetFuncName();
    auto tf_func = module.lookupSymbol<LLVMFuncOp>(tf_func_name);
    if (!tf_func) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(module.getBody());
      auto func_type = GetFuncType();
      tf_func = rewriter.create<LLVMFuncOp>(rewriter.getUnknownLoc(),
                                            tf_func_name, func_type);
    }
    return SymbolRefAttr::get(rewriter.getContext(), tf_func_name);
  }

 protected:
  virtual StringRef GetFuncName() const = 0;
  virtual Type GetFuncType() const = 0;

  Value CreateOrFindGlobalStringConstant(Location loc, OpBuilder &builder,
                                         StringRef base_name,
                                         StringRef str) const {
    auto module =
        builder.getInsertionBlock()->getParentOp()->getParentOfType<ModuleOp>();
    std::string global_name =
        llvm::formatv("{0}_{1}", base_name, llvm::hash_value(str));
    Operation *global_constant =
        SymbolTable::lookupNearestSymbolFrom(module, global_name);
    if (global_constant) {
      Value global_ptr = builder.create<LLVM::AddressOfOp>(
          loc, cast<LLVM::GlobalOp>(global_constant));
      Value c0 = builder.create<LLVM::ConstantOp>(loc, builder.getI64Type(),
                                                  builder.getIndexAttr(0));
      return builder.create<LLVM::GEPOp>(
          loc, LLVM::LLVMPointerType::get(builder.getIntegerType(8)),
          global_ptr, ValueRange{c0, c0});
    }
    return LLVM::createGlobalString(loc, builder, global_name, str,
                                    LLVM::Linkage::Internal);
  }

  std::pair<Value, Value> ConvertIntegerArrayAttrToStackAllocatedArray(
      Location loc, Type size_ty, Type element_ty,
      llvm::Optional<ArrayAttr> attr,
      ConversionPatternRewriter *rewriter) const {
    assert(size_ty.isa<IntegerType>() && "expect integer size type");
    assert(element_ty.isa<IntegerType>() && "expect integer element type");
    Type element_ptr_ty = LLVM::LLVMPointerType::get(element_ty);

    // If the attribute is missing or empty, set the element count to 0 and
    // return NULL.
    if (!attr.hasValue() || attr.getValue().empty()) {
      Value zero = rewriter->create<LLVM::ConstantOp>(
          loc, size_ty, rewriter->getIntegerAttr(size_ty, 0));
      Value null_ptr = rewriter->create<LLVM::NullOp>(loc, element_ptr_ty);
      return std::make_pair(zero, null_ptr);
    }

    // Allocate array to store the elements.
    auto &array_attr = attr.getValue();
{'fraction_non_alphanumeric': 0.11577311577311578, 'fraction_numerical': 0.004079254079254079, 'mean_word_length': 3.92915270464337, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 51, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 19, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3232588', 'n_tokens_mistral': 3043, 'n_tokens_neox': 2967, 'n_words': 728}