[Mlir-commits] [mlir] [mlir] Add `convertInstruction` and `getSupportedInstructions` to `LLVMImportInterface` (PR #86799)

[Fabian Mora]

https://github.com/fabianmcg created https://github.com/llvm/llvm-project/pull/86799

This patch adds the `convertInstruction` and `getSupportedInstructions` to `LLVMImportInterface`, allowing any non-LLVM dialect to specify how to import LLVM IR instructions.

This patch is necessary for https://github.com/llvm/llvm-project/pull/73057

>From 5f71a95de1609a40592c3ca2a881bd2f7bb74711 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Wed, 27 Mar 2024 12:28:06 +0000
Subject: [PATCH] [mlir] Add `convertInstruction` and
 `getSupportedInstructions` to `LLVMImportInterface`

This patch adds the `convertInstruction` and `getSupportedInstructions` to
`LLVMImportInterface`, allowing any non-LLVM dialect to specify how to import
LLVM IR instructions.

This patch is necessary for https://github.com/llvm/llvm-project/pull/73057
---
 .../mlir/Target/LLVMIR/LLVMImportInterface.h  | 53 +++++++++++++++++++
 mlir/lib/Target/LLVMIR/ModuleImport.cpp       |  8 ++-
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Target/LLVMIR/LLVMImportInterface.h b/mlir/include/mlir/Target/LLVMIR/LLVMImportInterface.h
index 9f8da83ae9c205..1bd81fcd9400cb 100644
--- a/mlir/include/mlir/Target/LLVMIR/LLVMImportInterface.h
+++ b/mlir/include/mlir/Target/LLVMIR/LLVMImportInterface.h
@@ -52,6 +52,15 @@ class LLVMImportDialectInterface
     return failure();
   }
 
+  /// Hook for derived dialect interfaces to implement the import of
+  /// instructions into MLIR.
+  virtual LogicalResult
+  convertInstruction(OpBuilder &builder, llvm::Instruction *inst,
+                     ArrayRef<llvm::Value *> llvmOperands,
+                     LLVM::ModuleImport &moduleImport) const {
+    return failure();
+  }
+
   /// Hook for derived dialect interfaces to implement the import of metadata
   /// into MLIR. Attaches the converted metadata kind and node to the provided
   /// operation.
@@ -66,6 +75,11 @@ class LLVMImportDialectInterface
   /// returns the list of supported intrinsic identifiers.
   virtual ArrayRef<unsigned> getSupportedIntrinsics() const { return {}; }
 
+  /// Hook for derived dialect interfaces to publish the supported instructions.
+  /// As every LLVM IR instructions has a unique integer identifier, the
+  /// function returns the list of supported instructions identifiers.
+  virtual ArrayRef<unsigned> getSupportedInstructions() const { return {}; }
+
   /// Hook for derived dialect interfaces to publish the supported metadata
   /// kinds. As every metadata kind has a unique integer identifier, the
   /// function returns the list of supported metadata identifiers.
@@ -100,9 +114,27 @@ class LLVMImportInterface
                           *it, iface.getDialect()->getNamespace(),
                           intrinsicToDialect.lookup(*it)->getNamespace()));
       }
+      const auto *instIt =
+          llvm::find_if(iface.getSupportedInstructions(), [&](unsigned id) {
+            return instructionToDialect.count(id);
+          });
+      if (instIt != iface.getSupportedInstructions().end()) {
+        return emitError(
+            UnknownLoc::get(iface.getContext()),
+            llvm::formatv(
+                "expected unique conversion for instruction ({0}), but "
+                "got conflicting {1} and {2} conversions",
+                *it, iface.getDialect()->getNamespace(),
+                instructionToDialect.lookup(*it)
+                    ->getDialect()
+                    ->getNamespace()));
+      }
       // Add a mapping for all supported intrinsic identifiers.
       for (unsigned id : iface.getSupportedIntrinsics())
         intrinsicToDialect[id] = iface.getDialect();
+      // Add a mapping for all supported instruction identifiers.
+      for (unsigned id : iface.getSupportedInstructions())
+        instructionToDialect[id] = &iface;
       // Add a mapping for all supported metadata kinds.
       for (unsigned kind : iface.getSupportedMetadata())
         metadataToDialect[kind].push_back(iface.getDialect());
@@ -132,6 +164,26 @@ class LLVMImportInterface
     return intrinsicToDialect.count(id);
   }
 
+  /// Converts the LLVM instruction to an MLIR operation if a conversion exists.
+  /// Returns failure otherwise.
+  LogicalResult convertInstruction(OpBuilder &builder, llvm::Instruction *inst,
+                                   ArrayRef<llvm::Value *> llvmOperands,
+                                   LLVM::ModuleImport &moduleImport) const {
+    // Lookup the dialect interface for the given instruction.
+    const LLVMImportDialectInterface *iface =
+        instructionToDialect.lookup(inst->getOpcode());
+    if (!iface)
+      return failure();
+
+    return iface->convertInstruction(builder, inst, llvmOperands, moduleImport);
+  }
+
+  /// Returns true if the given LLVM IR instruction is convertible to an MLIR
+  /// operation.
+  bool isConvertibleInstruction(unsigned id) {
+    return instructionToDialect.count(id);
+  }
+
   /// Attaches the given LLVM metadata to the imported operation if a conversion
   /// to one or more MLIR dialect attributes exists and succeeds. Returns
   /// success if at least one of the conversions is successful and failure if
@@ -166,6 +218,7 @@ class LLVMImportInterface
 
 private:
   DenseMap<unsigned, Dialect *> intrinsicToDialect;
+  DenseMap<unsigned, const LLVMImportDialectInterface *> instructionToDialect;
   DenseMap<unsigned, SmallVector<Dialect *, 1>> metadataToDialect;
 };
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 6e70d52fa760b6..af419261c8919c 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -123,13 +123,17 @@ static SmallVector<int64_t> getPositionFromIndices(ArrayRef<unsigned> indices) {
 /// access to the private module import methods.
 static LogicalResult convertInstructionImpl(OpBuilder &odsBuilder,
                                             llvm::Instruction *inst,
-                                            ModuleImport &moduleImport) {
+                                            ModuleImport &moduleImport,
+                                            LLVMImportInterface &importIface) {
   // Copy the operands to an LLVM operands array reference for conversion.
   SmallVector<llvm::Value *> operands(inst->operands());
   ArrayRef<llvm::Value *> llvmOperands(operands);
 
   // Convert all instructions that provide an MLIR builder.
 #include "mlir/Dialect/LLVMIR/LLVMOpFromLLVMIRConversions.inc"
+  if (importIface.isConvertibleInstruction(inst->getOpcode()))
+    return importIface.convertInstruction(odsBuilder, inst, llvmOperands,
+                                          moduleImport);
   return failure();
 }
 
@@ -1596,7 +1600,7 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
   }
 
   // Convert all instructions that have an mlirBuilder.
-  if (succeeded(convertInstructionImpl(builder, inst, *this)))
+  if (succeeded(convertInstructionImpl(builder, inst, *this, iface)))
     return success();
 
   return emitError(loc) << "unhandled instruction: " << diag(*inst);