Generate datalayout in module. (llvm#1021)

* Generate datalayout in module.

Signed-off-by: Ettore Tiotto <[email protected]>

* Generate datalayout in module.

Signed-off-by: Ettore Tiotto <[email protected]>

MLIR.cmake

@@ -92,7 +92,7 @@ endfunction()
 
 # add_onnx_mlir_library(name sources...
 #   This function (generally) has the same semantic as add_library. In
-#   addition is supports the arguments below and it does the following
+#   addition it supports the arguments below and it does the following
 #   by default (unless an argument overrides this):
 #   1. Add the library
 #   2. Add the default target_include_directories
@@ -108,12 +108,15 @@ endfunction()
 #     Same semantics as target_include_directories().
 #   LINK_LIBS lib_targets...
 #     Same semantics as target_link_libraries().
+#   LINK_COMPONENTS llvm_components...
+#     Link the specified LLVM components.
+#     Note: only one linkage mode can be specified.
 #   )
 function(add_onnx_mlir_library name)
   cmake_parse_arguments(ARG
     "EXCLUDE_FROM_OM_LIBS;NO_INSTALL"
     ""
-    "DEPENDS;INCLUDE_DIRS;LINK_LIBS"
+    "DEPENDS;INCLUDE_DIRS;LINK_LIBS;LINK_COMPONENTS"
     ${ARGN}
     )
 
@@ -142,6 +145,21 @@ function(add_onnx_mlir_library name)
     target_link_libraries(${name} ${ARG_LINK_LIBS})
   endif()
 
+  if (ARG_LINK_COMPONENTS)
+    set(LinkageMode)
+    if (ARG_LINK_COMPONENTS MATCHES "^(PUBLIC|PRIVATE|INTERFACE)")
+      list(POP_FRONT ARG_LINK_COMPONENTS LinkageMode)
+    endif()
+
+    llvm_map_components_to_libnames(COMPONENT_LIBS ${ARG_LINK_COMPONENTS})
+
+    if (LinkageMode)
+      target_link_libraries(${name} ${LinkageMode} ${COMPONENT_LIBS})
+    else()
+      target_link_libraries(${name} PRIVATE ${COMPONENT_LIBS})
+    endif()
+  endif()
+
   if (NOT ARG_NO_INSTALL)
     install(TARGETS ${name}
       ARCHIVE DESTINATION lib

src/Compiler/CMakeLists.txt

@@ -43,6 +43,14 @@ add_onnx_mlir_library(CompilerUtils
   MLIRAffineTransforms
   MLIRLinalgTransforms
   MLIRLLVMToLLVMIRTranslation
+
+  # Link LLVM libraries necessary to query which target architectures are configured.
+  LINK_COMPONENTS PRIVATE
+  AllTargetsAsmParsers
+  AllTargetsCodeGens
+  AllTargetsDescs
+  AllTargetsInfos
+  MC
   )
 
 # CompilerUtils does not require cruntime or jniruntime to build, however, they are

src/Compiler/CompilerUtils.cpp

@@ -2,7 +2,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-//===--------------------------- MainUtils.cpp ---------------------------===//
+//===-------------------------- CompilerUtils.cpp -------------------------===//
 //
 // Copyright 2019-2020 The IBM Research Authors.
 //
@@ -17,12 +17,16 @@
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Target/TargetMachine.h"
 
-#include "CompilerUtils.hpp"
 #include "ExternalUtil.hpp"
+#include "src/Compiler/CompilerUtils.hpp"
 #include "src/Support/OMOptions.hpp"
 
 using namespace std;
@@ -88,14 +92,15 @@ llvm::cl::opt<string> shapeInformation("shapeInformation",
         "unknown dimensions)"),
     llvm::cl::value_desc("value"), llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<string> mtriple("mtriple", llvm::cl::desc("Target architecture"),
-    llvm::cl::value_desc("<llvm target triple>"),
-    llvm::cl::cat(OnnxMlirOptions), llvm::cl::ValueRequired);
-
-llvm::cl::opt<string> mcpu("mcpu", llvm::cl::desc("Target cpu"),
-    llvm::cl::value_desc("<llvm cpu value>"), llvm::cl::cat(OnnxMlirOptions),
+llvm::cl::opt<std::string> mtriple("mtriple",
+    llvm::cl::desc("Target architecture"),
+    llvm::cl::value_desc("LLVM target triple>"), llvm::cl::cat(OnnxMlirOptions),
     llvm::cl::ValueRequired);
 
+llvm::cl::opt<std::string> mcpu("mcpu", llvm::cl::desc("Target cpu"),
+    llvm::cl::value_desc("Target a specific CPU type"),
+    llvm::cl::cat(OnnxMlirOptions), llvm::cl::ValueRequired);
+
 // Make a function that forces preserving all files using the runtime arguments
 // and/or the overridePreserveFiles enum.
 enum class KeepFilesOfType { All, MLIR, Bitcode, Object, None };
@@ -291,7 +296,7 @@ void LoadMLIR(string inputFilename, mlir::MLIRContext &context,
   }
 }
 
-string getTargetCpuOption() {
+static std::string getTargetCpuOption() {
   string targetOptions = "";
   if (mcpu != "")
     targetOptions += "--mcpu=" + mcpu;
@@ -648,9 +653,6 @@ void outputCode(
 
   module->print(output->os(), flags);
   output->keep();
-
-  if (printIR)
-    module->print(llvm::outs(), flags);
 }
 
 void emitOutputFiles(string outputBaseName, EmissionTargetType emissionTarget,
@@ -713,8 +715,50 @@ void emitOutputFiles(string outputBaseName, EmissionTargetType emissionTarget,
   }
 }
 
+/// Return the module datalayout string. The datalayout string is determined by
+/// creating a target machine using the target triple and target cpu.
+static std::string getDataLayout(const Location &loc) {
+  const std::string targetTriple =
+      (mtriple != "") ? mtriple.getValue() : kDefaultTriple;
+  const std::string targetCpu = (mcpu != "") ? mcpu.getValue() : "";
+
+  std::string error;
+  const llvm::Target *LLVMTarget =
+      llvm::TargetRegistry::lookupTarget(targetTriple, error);
+  if (!LLVMTarget) {
+    emitError(loc, Twine("Target architecture is unknown: ") + error);
+    return nullptr;
+  }
+
+  llvm::TargetOptions ops;
+  llvm::TargetMachine *targetMachine = LLVMTarget->createTargetMachine(
+      targetTriple, targetCpu, "" /*features*/, ops, None);
+  if (!targetMachine) {
+    emitError(loc, "failed to create target machine");
+    return nullptr;
+  }
+
+  const llvm::DataLayout &dl = targetMachine->createDataLayout();
+  std::string dataLayoutString = dl.getStringRepresentation();
+  assert(dataLayoutString != "" && "Expecting a valid target datalayout");
+
+  return dataLayoutString;
+}
+
 int compileModule(mlir::OwningModuleRef &module, mlir::MLIRContext &context,
     std::string outputBaseName, EmissionTargetType emissionTarget) {
+  // Initialize the targets support for all targets LLVM was configured for.
+  llvm::InitializeAllTargets();
+  llvm::InitializeAllTargetMCs();
+  llvm::InitializeAllAsmPrinters();
+  llvm::InitializeAllAsmParsers();
+
+  // Set the module datalayout.
+  Operation &moduleOp = *(module->getOperation());
+  Location loc = moduleOp.getLoc();
+  moduleOp.setAttr(LLVM::LLVMDialect::getDataLayoutAttrName(),
+      StringAttr::get(&context, getDataLayout(loc)));
+
   mlir::PassManager pm(&context, mlir::OpPassManager::Nesting::Implicit);
 
   if (keepFiles(KeepFilesOfType::MLIR)) {
@@ -725,9 +769,8 @@ int compileModule(mlir::OwningModuleRef &module, mlir::MLIRContext &context,
 
   InputIRLevelType inputIRLevel = determineInputIRLevel(module);
 
-  if (inputIRLevel <= ONNXLevel && emissionTarget >= EmitONNXIR) {
+  if (inputIRLevel <= ONNXLevel && emissionTarget >= EmitONNXIR)
     addONNXToMLIRPasses(pm);
-  }
 
   if (emissionTarget >= EmitMLIR) {
     if (inputIRLevel <= ONNXLevel)
@@ -743,6 +786,13 @@ int compileModule(mlir::OwningModuleRef &module, mlir::MLIRContext &context,
   if (mlir::failed(pm.run(*module)))
     return 4;
 
-  emitOutputFiles(outputBaseName, emissionTarget, context, module);
+  if (printIR) {
+    mlir::OpPrintingFlags flags;
+    if (preserveLocations)
+      flags.enableDebugInfo();
+    module->print(llvm::outs(), flags);
+  } else
+    emitOutputFiles(outputBaseName, emissionTarget, context, module);
+
   return 0;
 }

src/Conversion/KrnlToLLVM/KrnlToLLVM.cpp

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
@@ -35,8 +36,13 @@
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/Sequence.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Endian.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/Target/TargetOptions.h"
 
 #include "onnx/onnx_pb.h"
 
@@ -1907,14 +1913,11 @@ struct ConvertKrnlToLLVMPass
 } // end anonymous namespace
 
 void ConvertKrnlToLLVMPass::runOnOperation() {
-  // Annotate ModuleOp with endian information so that LLVM global constants are
-  // handled correctly by the other LLVM tools such as 'opt'.
-
-  bool isLittleEndian = llvm::support::endian::system_endianness() ==
-                        llvm::support::endianness::little;
-  StringRef endian = isLittleEndian ? "e" : "E";
   ModuleOp module = getOperation();
-  module->setAttr("llvm.data_layout", StringAttr::get(&getContext(), endian));
+  const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
+  LowerToLLVMOptions options(
+      &getContext(), dataLayoutAnalysis.getAtOrAbove(module));
+  options.emitCWrappers = true;
 
   // Determine, for each output, whether it is a constant or not.
   SmallVector<bool, 4> constantOutputs;
@@ -1926,10 +1929,6 @@ void ConvertKrnlToLLVMPass::runOnOperation() {
   target.addLegalOp<ModuleOp>();
   target.addLegalOp<UnrealizedConversionCastOp>();
 
-  // Lower the MemRef types to a representation in LLVM.
-  LowerToLLVMOptions options(&getContext());
-  options.emitCWrappers = true;
-
   // Convert types to legal types for the LLVM dialect.
   LLVMTypeConverter typeConverter(&getContext(), options);
 

test/mlir/module_op_be/module_op.mlir

@@ -1,6 +1,6 @@
-// RUN: onnx-mlir-opt --convert-krnl-to-llvm %s -split-input-file | FileCheck %s
+// RUN: onnx-mlir --printIR %s | FileCheck %s
 
-// CHECK: module attributes {llvm.data_layout = "E"}
+// CHECK: module attributes {llvm.data_layout = "E-{{.*}}"}
 module {
 }
 

test/mlir/module_op_le/module_op.mlir

@@ -1,5 +1,5 @@
-// RUN: onnx-mlir-opt --convert-krnl-to-llvm %s -split-input-file | FileCheck %s
+// RUN: onnx-mlir --printIR %s | FileCheck %s
 
-// CHECK: module attributes {llvm.data_layout = "e"}
+// CHECK: module attributes {llvm.data_layout = "e-{{.*}}"}
 module {
 }

test/mlir/onnx/onnx_location.mlir

@@ -1,11 +1,11 @@
 
-// RUN: onnx-mlir --EmitMLIR  --preserveLocations --printIR %s |  FileCheck %s --check-prefix=PRESENT; rm %p/*.onnx.mlir ; rm %p/*.tmp
-// RUN: onnx-mlir --EmitMLIR  --printIR %s | FileCheck %s --check-prefix=ABSENT; rm %p/*.onnx.mlir ; rm %p/*.tmp
+// RUN: onnx-mlir --preserveLocations --printIR %s |  FileCheck %s --check-prefix=PRESENT
+// RUN: onnx-mlir --printIR %s | FileCheck %s --check-prefix=ABSENT
 
   func @main_graph(%arg0: tensor<1x16xf32>, %arg1: tensor<1x16xf32>) -> tensor<1x16xf32>  {
     %0 = "onnx.Add"(%arg0, %arg1) : (tensor<1x16xf32>, tensor<1x16xf32>) -> tensor<1x16xf32> loc("/build/workspace/addop.onnx":1:0)
      return %0 : tensor<1x16xf32>
   }
 
-// PRESENT: loc("onnx.Add"("{{(/[[:alnum:]]+)+}}.onnx":1:0))
-// ABSENT-NOT: loc("{{("onnx.Add"(/[[:alnum:]]+)+}}.onnx":1:0))
+// PRESENT:    loc("onnx.Add"("{{(/[[:alnum:]]+)+}}.onnx":1:0))
+// ABSENT-NOT: loc("onnx.Add"("{{(/[[:alnum:]]+)+}}.onnx":1:0))