[Mlir-commits] [mlir] [memref] Support dense resources for memref-to-llvm (PR #79380)

[Rob Suderman]

https://github.com/rsuderman created https://github.com/llvm/llvm-project/pull/79380

Memref to llvm did not handle the DenseResourceElementsAttr. Reworked
the lowering so that it converts to a valid LLVM type.

>From 56a30c673b1b7c44e1b6c1ab1ff648f29f4fe39c Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Tue, 23 Jan 2024 21:48:45 -0800
Subject: [PATCH 1/2] [memref] Support dense resources for memref-to-llvm

Memref to llvm did not handle the DenseResourceElementsAttr. Reworked
the lowering so that it converts to a valid LLVM type.
---
 .../Conversion/MemRefToLLVM/MemRefToLLVM.cpp  | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 2bfca303b5fd489..c3d90260e2a2784 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -23,6 +23,7 @@
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/MathExtras.h"
@@ -501,13 +502,20 @@ struct GlobalMemrefOpLowering
 
     Attribute initialValue = nullptr;
     if (!global.isExternal() && !global.isUninitialized()) {
-      auto elementsAttr = llvm::cast<ElementsAttr>(*global.getInitialValue());
-      initialValue = elementsAttr;
-
-      // For scalar memrefs, the global variable created is of the element type,
-      // so unpack the elements attribute to extract the value.
-      if (type.getRank() == 0)
-        initialValue = elementsAttr.getSplatValue<Attribute>();
+      Attribute initialAttr = *global.getInitialValue();
+      if (auto resourceAttr = llvm::dyn_cast<DenseResourceElementsAttr>(initialAttr)) {
+        auto blob = resourceAttr.getRawHandle().getBlob();
+        initialAttr = DenseElementsAttr::get(resourceAttr.getType(), blob->getData());
+      } 
+
+      if (auto elementsAttr = llvm::cast<ElementsAttr>(initialAttr)) {
+        initialValue = elementsAttr;
+
+        // For scalar memrefs, the global variable created is of the element type,
+        // so unpack the elements attribute to extract the value.
+        if (type.getRank() == 0)
+          initialValue = elementsAttr.getSplatValue<Attribute>();
+      }
     }
 
     uint64_t alignment = global.getAlignment().value_or(0);

>From 479358627d2794cc929182d648fc8ac6b42e1d72 Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Wed, 24 Jan 2024 14:16:23 -0800
Subject: [PATCH 2/2] Added test, use different constructor

---
 mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp |  3 ++-
 .../Conversion/MemRefToLLVM/memref-to-llvm.mlir   | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index c3d90260e2a2784..8fd1f37acdad6c2 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -505,7 +505,8 @@ struct GlobalMemrefOpLowering
       Attribute initialAttr = *global.getInitialValue();
       if (auto resourceAttr = llvm::dyn_cast<DenseResourceElementsAttr>(initialAttr)) {
         auto blob = resourceAttr.getRawHandle().getBlob();
-        initialAttr = DenseElementsAttr::get(resourceAttr.getType(), blob->getData());
+        initialAttr = DenseElementsAttr::getFromRawBuffer(
+          resourceAttr.getType(), blob->getData());
       } 
 
       if (auto elementsAttr = llvm::cast<ElementsAttr>(initialAttr)) {
diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
index 37999d6fc14ad19..9e00fd1bb5c73cb 100644
--- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
@@ -330,6 +330,21 @@ memref.global "private" @gv4 : memref<f32> = dense<1.0> {alignment = 64}
 
 // -----
 
+module {
+  // CHECK: llvm.mlir.global private constant @__constant_xf32(1.41421354 : f32) {addr_space = 0 : i32} : f32
+  memref.global "private" constant @__constant_xf32 : memref<f32> = dense_resource<NAME>
+}
+
+{-#
+  dialect_resources: {
+    builtin: {
+      NAME: "0x08000000F304B53F"
+    }
+  }
+#-}
+
+// -----
+
 // Expand shapes need to be expanded outside of the memref-to-llvm pass.
 // CHECK-LABEL: func @expand_shape_static(
 // CHECK-SAME:         %[[ARG:.*]]: memref<{{.*}}>)