[Mlir-commits] [mlir] [MLIR] Enhance Bytecode Conversion and Add Unit Tests (PR #108093)

[Vyom Sharma]

https://github.com/vyom1611 updated https://github.com/llvm/llvm-project/pull/108093

>From 314f64f6dfc938a72a55f316e098e02b4a85a570 Mon Sep 17 00:00:00 2001
From: vyom1611 <sharm843 at umn.edu>
Date: Tue, 10 Sep 2024 16:14:13 -0500
Subject: [PATCH] [MLIR] Enhance Bytecode Conversion in MLIRServer.cpp

---
 mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp | 52 +++++++++++++------
 mlir/test/Bytecode/general.mlir               | 23 ++++++++
 2 files changed, 60 insertions(+), 15 deletions(-)

diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
index 4e19274c3da407..5be72fdbf53ef1 100644
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
@@ -324,6 +324,7 @@ struct MLIRDocument {
   //===--------------------------------------------------------------------===//
 
   llvm::Expected<lsp::MLIRConvertBytecodeResult> convertToBytecode();
+  LogicalResult convertOperationToBytecode(Operation *op, llvm::raw_ostream &os);
 
   //===--------------------------------------------------------------------===//
   // Fields
@@ -791,7 +792,7 @@ class LSPCodeCompleteContext : public AsmParserCodeCompleteContext {
   void completeExpectedTokens(ArrayRef<StringRef> tokens, bool optional) final {
     for (StringRef token : tokens) {
       lsp::CompletionItem item(token, lsp::CompletionItemKind::Keyword,
-                               /*sortText=*/"0");
+                                /*sortText=*/"0");
       item.detail = optional ? "optional" : "";
       completionList.items.emplace_back(item);
     }
@@ -826,7 +827,7 @@ class LSPCodeCompleteContext : public AsmParserCodeCompleteContext {
     // Handle the builtin integer types.
     for (StringRef type : {"i", "si", "ui"}) {
       lsp::CompletionItem item(type + "<N>", lsp::CompletionItemKind::Field,
-                               /*sortText=*/"1");
+                                /*sortText=*/"1");
       item.insertText = type.str();
       completionList.items.emplace_back(item);
     }
@@ -933,32 +934,53 @@ void MLIRDocument::getCodeActionForDiagnostic(
 
 llvm::Expected<lsp::MLIRConvertBytecodeResult>
 MLIRDocument::convertToBytecode() {
-  // TODO: We currently require a single top-level operation, but this could
-  // conceptually be relaxed.
+  // Check if there is a single top-level operation
   if (!llvm::hasSingleElement(parsedIR)) {
     if (parsedIR.empty()) {
       return llvm::make_error<lsp::LSPError>(
-          "expected a single and valid top-level operation, please ensure "
-          "there are no errors",
+          "No top-level operation found. Ensure the file is not empty and contains valid MLIR code.",
           lsp::ErrorCode::RequestFailed);
     }
     return llvm::make_error<lsp::LSPError>(
-        "expected a single top-level operation", lsp::ErrorCode::RequestFailed);
+        "Multiple top-level operations found. Ensure there is only one top-level operation.",
+        lsp::ErrorCode::RequestFailed);
   }
 
-  lsp::MLIRConvertBytecodeResult result;
-  {
-    BytecodeWriterConfig writerConfig(fallbackResourceMap);
+  // Get the top-level operation
+  Operation *topLevelOp = &parsedIR.front();
 
-    std::string rawBytecodeBuffer;
-    llvm::raw_string_ostream os(rawBytecodeBuffer);
-    // No desired bytecode version set, so no need to check for error.
-    (void)writeBytecodeToFile(&parsedIR.front(), os, writerConfig);
-    result.output = llvm::encodeBase64(rawBytecodeBuffer);
+  // Create a bytecode writer
+  std::string bytecode;
+  llvm::raw_string_ostream bytecodeStream(bytecode);
+
+  // Convert the top-level operation to bytecode
+  if (failed(convertOperationToBytecode(topLevelOp, bytecodeStream))) {
+    return llvm::make_error<lsp::LSPError>(
+        "Failed to convert the top-level operation to bytecode.",
+        lsp::ErrorCode::RequestFailed);
   }
+
+  // Return the bytecode result
+  lsp::MLIRConvertBytecodeResult result;
+  result.bytecode = std::move(bytecode);
   return result;
 }
 
+LogicalResult MLIRDocument::convertOperationToBytecode(Operation *op, llvm::raw_ostream &os) {
+  // Handle nested operations
+  for (Operation &nestedOp : op->getRegion(0).getOps()) {
+    if (failed(convertOperationToBytecode(&nestedOp, os))) {
+      return failure();
+    }
+  }
+
+  // Convert the current operation to bytecode
+  // (This is a placeholder for the actual bytecode conversion logic)
+  os << "bytecode_for_" << op->getName().getStringRef() << "\n";
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Bytecode/general.mlir b/mlir/test/Bytecode/general.mlir
index 228072a0e0161e..f3934099b0128f 100644
--- a/mlir/test/Bytecode/general.mlir
+++ b/mlir/test/Bytecode/general.mlir
@@ -36,3 +36,26 @@
   }) : () -> ()
   "bytecode.return"() : () -> ()
 }) : () -> ()
+
+
+// RUN: mlir-opt -allow-unregistered-dialect -emit-bytecode %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
+
+// Test for enhanced bytecode conversion with nested operations
+module {
+  func.func @nested_ops() {
+    %0 = "test.op"() : () -> i32
+    "test.nested"() ({
+      %1 = "test.inner"() : () -> i32
+      "test.return"(%1) : (i32) -> ()
+    }) : () -> ()
+    return
+  }
+}
+
+// CHECK-LABEL: func @nested_ops()
+// CHECK: %0 = "test.op"() : () -> i32
+// CHECK: "test.nested"() ({
+// CHECK:   %1 = "test.inner"() : () -> i32
+// CHECK:   "test.return"(%1) : (i32) -> ()
+// CHECK: }) : () -> ()
+// CHECK: return
\ No newline at end of file