From 2def12ebc6cc904cddb4bc608df63014ec2bfe86 Mon Sep 17 00:00:00 2001
From: KareemErgawy <kareem.ergawy@tomtom.com>
Date: Mon, 7 Jun 2021 13:19:39 +0200
Subject: [PATCH] [MLIR][SPIRV] Use getAsmResultName(...) hook for AddressOfOp.

Implements better naming for results of spv.mlir.addressof ops by making it
inherit from OpAsmOpInterface and implementing the associated
getAsmResultName(...) hook.

Reviewed By: antiagainst

Differential Revision: https://reviews.llvm.org/D103594
---
 mlir/include/mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td |  3 ++-
 mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp                  |  8 ++++++++
 mlir/test/Dialect/SPIRV/IR/asm-op-interface.mlir        | 15 ++++++++++++++-
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td
index c185cf0..95daa98 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td
@@ -23,7 +23,8 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // -----
 
-def SPV_AddressOfOp : SPV_Op<"mlir.addressof", [InFunctionScope, NoSideEffect]> {
+def SPV_AddressOfOp : SPV_Op<"mlir.addressof",
+    [DeclareOpInterfaceMethods<OpAsmOpInterface>, InFunctionScope, NoSideEffect]> {
   let summary = "Get the address of a global variable.";
 
   let description = [{
diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
index 782d7dd..374fdad 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
@@ -1690,6 +1690,14 @@ void mlir::spirv::ConstantOp::getAsmResultNames(
   setNameFn(getResult(), specialName.str());
 }
 
+void mlir::spirv::AddressOfOp::getAsmResultNames(
+    llvm::function_ref<void(mlir::Value, llvm::StringRef)> setNameFn) {
+  SmallString<32> specialNameBuffer;
+  llvm::raw_svector_ostream specialName(specialNameBuffer);
+  specialName << variable() << "_addr";
+  setNameFn(getResult(), specialName.str());
+}
+
 //===----------------------------------------------------------------------===//
 // spv.EntryPoint
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/SPIRV/IR/asm-op-interface.mlir b/mlir/test/Dialect/SPIRV/IR/asm-op-interface.mlir
index a53f061..ebcf9f1 100644
--- a/mlir/test/Dialect/SPIRV/IR/asm-op-interface.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/asm-op-interface.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s | FileCheck %s
+// RUN: mlir-opt %s -split-input-file | FileCheck %s
 
 func @const() -> () {
   // CHECK: %true
@@ -26,3 +26,16 @@ func @const() -> () {
 
   return
 }
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.GlobalVariable @global_var : !spv.ptr<f32, Input>
+
+  spv.func @addressof() -> () "None" {
+    // CHECK: %global_var_addr = spv.mlir.addressof 
+    %0 = spv.mlir.addressof @global_var : !spv.ptr<f32, Input>
+    spv.Return
+  }
+}
+
-- 
2.7.4