[Mlir-commits] [mlir] [mlir][ml_program] Add asm names for load (PR #70409)

[Jacques Pienaar]

https://github.com/jpienaar created https://github.com/llvm/llvm-project/pull/70409

None

>From 2b6eff4eef335bc2322ee6565b1948490106e610 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Thu, 26 Oct 2023 21:00:37 -0700
Subject: [PATCH] [mlir][ml_program] Add asm names for load

---
 .../mlir/Dialect/MLProgram/IR/MLProgramOps.td   | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td
index f1d4972f48cff64..fd213824b30ca45 100644
--- a/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td
+++ b/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td
@@ -16,6 +16,7 @@ include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/FunctionInterfaces.td"
+include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 
@@ -153,6 +154,7 @@ def MLProgram_GlobalOp : MLProgram_Op<"global", [
 //===----------------------------------------------------------------------===//
 
 def MLProgram_GlobalLoadOp : MLProgram_Op<"global_load", [
+    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     DeclareOpInterfaceMethods<SymbolUserOpInterface>
   ]> {
   let summary = "Direct load of a mutable value from a global";
@@ -190,6 +192,13 @@ def MLProgram_GlobalLoadOp : MLProgram_Op<"global_load", [
     /// Gets the corresponding GlobalOp (or nullptr).
     GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
   }];
+
+  let extraClassDefinition = [{
+    void $cppClass::getAsmResultNames(
+        function_ref<void(::mlir::Value, ::llvm::StringRef)> setNameFn) {
+      setNameFn(getResult(), getGlobal().getLeafReference());
+    }
+  }];
 }
 
 //===----------------------------------------------------------------------===//
@@ -198,6 +207,7 @@ def MLProgram_GlobalLoadOp : MLProgram_Op<"global_load", [
 
 def MLProgram_GlobalLoadConstOp : MLProgram_Op<"global_load_const", [
     Pure,
+    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     DeclareOpInterfaceMethods<SymbolUserOpInterface>
   ]> {
   let summary = "Direct load a constant value from a global";
@@ -229,6 +239,13 @@ def MLProgram_GlobalLoadConstOp : MLProgram_Op<"global_load_const", [
     /// Gets the corresponding GlobalOp (or nullptr).
     GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
   }];
+
+  let extraClassDefinition = [{
+    void $cppClass::getAsmResultNames(
+      function_ref<void(::mlir::Value, ::llvm::StringRef)> setNameFn) {
+        setNameFn(getResult(), getGlobal().getLeafReference());
+    }
+  }];
 }
 
 //===----------------------------------------------------------------------===//