[MLIR][NVVM] Enable nvvm intrinsics import to LLVMIR (#68843)

Co-authored-by: Tobias Gysi <[email protected]>
Co-authored-by: Christian Ulmann <[email protected]>

Author: 3 people

mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt

@@ -56,6 +56,8 @@ add_mlir_dialect(NVVMOps nvvm)
 add_mlir_doc(NVVMOps NVVMDialect Dialects/ -gen-dialect-doc -dialect=nvvm)
 set(LLVM_TARGET_DEFINITIONS NVVMOps.td)
 mlir_tablegen(NVVMConversions.inc -gen-llvmir-conversions)
+mlir_tablegen(NVVMFromLLVMIRConversions.inc -gen-intr-from-llvmir-conversions)
+mlir_tablegen(NVVMConvertibleLLVMIRIntrinsics.inc -gen-convertible-llvmir-intrinsics)
 mlir_tablegen(NVVMOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(NVVMOpsEnums.cpp.inc -gen-enum-defs)
 mlir_tablegen(NVVMOpsAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=nvvm)

mlir/include/mlir/Target/LLVMIR/Dialect/All.h

@@ -22,6 +22,7 @@
 #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.h"
@@ -74,6 +75,7 @@ registerAllGPUToLLVMIRTranslations(DialectRegistry &registry) {
 static inline void
 registerAllFromLLVMIRTranslations(DialectRegistry &registry) {
   registerLLVMDialectImport(registry);
+  registerNVVMDialectImport(registry);
 }
 } // namespace mlir
 

mlir/include/mlir/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.h

@@ -0,0 +1,31 @@
+//===- LLVMIRToNVVMTranslation.h - LLVM IR to NVVM Dialect ------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This provides registration calls for LLVM IR to NVVM dialect translation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_LLVMIR_DIALECT_NVVM_LLVMIRTONVVMTRANSLATION_H
+#define MLIR_TARGET_LLVMIR_DIALECT_NVVM_LLVMIRTONVVMTRANSLATION_H
+
+namespace mlir {
+
+class DialectRegistry;
+class MLIRContext;
+
+/// Registers the NVVM dialect and its import from LLVM IR in the given
+/// registry.
+void registerNVVMDialectImport(DialectRegistry &registry);
+
+/// Registers the NVVM dialect and its import from LLVM IR with the given
+/// context.
+void registerNVVMDialectImport(MLIRContext &context);
+
+} // namespace mlir
+
+#endif // MLIR_TARGET_LLVMIR_DIALECT_NVVM_LLVMIRTONVVMTRANSLATION_H

mlir/lib/Target/LLVMIR/Dialect/NVVM/CMakeLists.txt

@@ -1,3 +1,21 @@
+set(LLVM_OPTIONAL_SOURCES
+  LLVMIRToNVVMTranslation.cpp
+  NVVMToLLVMIRTranslation.cpp
+  )
+
+add_mlir_translation_library(MLIRLLVMIRToNVVMTranslation
+  LLVMIRToNVVMTranslation.cpp
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRNVVMDialect
+  MLIRSupport
+  MLIRTargetLLVMIRImport
+  )
+
 add_mlir_translation_library(MLIRNVVMToLLVMIRTranslation
   NVVMToLLVMIRTranslation.cpp
 

mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp

@@ -0,0 +1,93 @@
+//===- LLVMIRToNVVMTranslation.cpp - Translate LLVM IR to NVVM dialect ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a translation between LLVM IR and the MLIR NVVM dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "mlir/Target/LLVMIR/ModuleImport.h"
+
+#include "llvm/IR/IntrinsicsNVPTX.h"
+
+using namespace mlir;
+using namespace mlir::NVVM;
+
+/// Returns true if the LLVM IR intrinsic is convertible to an MLIR NVVM dialect
+/// intrinsic. Returns false otherwise.
+static bool isConvertibleIntrinsic(llvm::Intrinsic::ID id) {
+  static const DenseSet<unsigned> convertibleIntrinsics = {
+#include "mlir/Dialect/LLVMIR/NVVMConvertibleLLVMIRIntrinsics.inc"
+  };
+  return convertibleIntrinsics.contains(id);
+}
+
+/// Returns the list of LLVM IR intrinsic identifiers that are convertible to
+/// MLIR NVVM dialect intrinsics.
+static ArrayRef<unsigned> getSupportedIntrinsicsImpl() {
+  static const SmallVector<unsigned> convertibleIntrinsics = {
+#include "mlir/Dialect/LLVMIR/NVVMConvertibleLLVMIRIntrinsics.inc"
+  };
+  return convertibleIntrinsics;
+}
+
+/// Converts the LLVM intrinsic to an MLIR NVVM dialect operation if a
+/// conversion exits. Returns failure otherwise.
+static LogicalResult convertIntrinsicImpl(OpBuilder &odsBuilder,
+                                          llvm::CallInst *inst,
+                                          LLVM::ModuleImport &moduleImport) {
+  llvm::Intrinsic::ID intrinsicID = inst->getIntrinsicID();
+
+  // Check if the intrinsic is convertible to an MLIR dialect counterpart and
+  // copy the arguments to an an LLVM operands array reference for conversion.
+  if (isConvertibleIntrinsic(intrinsicID)) {
+    SmallVector<llvm::Value *> args(inst->args());
+    ArrayRef<llvm::Value *> llvmOperands(args);
+#include "mlir/Dialect/LLVMIR/NVVMFromLLVMIRConversions.inc"
+  }
+
+  return failure();
+}
+
+namespace {
+
+/// Implementation of the dialect interface that converts operations belonging
+/// to the NVVM dialect.
+class NVVMDialectLLVMIRImportInterface : public LLVMImportDialectInterface {
+public:
+  using LLVMImportDialectInterface::LLVMImportDialectInterface;
+
+  /// Converts the LLVM intrinsic to an MLIR NVVM dialect operation if a
+  /// conversion exits. Returns failure otherwise.
+  LogicalResult convertIntrinsic(OpBuilder &builder, llvm::CallInst *inst,
+                                 LLVM::ModuleImport &moduleImport) const final {
+    return convertIntrinsicImpl(builder, inst, moduleImport);
+  }
+
+  /// Returns the list of LLVM IR intrinsic identifiers that are convertible to
+  /// MLIR NVVM dialect intrinsics.
+  ArrayRef<unsigned> getSupportedIntrinsics() const final {
+    return getSupportedIntrinsicsImpl();
+  }
+};
+
+} // namespace
+
+void mlir::registerNVVMDialectImport(DialectRegistry &registry) {
+  registry.insert<NVVM::NVVMDialect>();
+  registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+    dialect->addInterfaces<NVVMDialectLLVMIRImportInterface>();
+  });
+}
+
+void mlir::registerNVVMDialectImport(MLIRContext &context) {
+  DialectRegistry registry;
+  registerNVVMDialectImport(registry);
+  context.appendDialectRegistry(registry);
+}