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

[Rob Suderman]

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

>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/4] [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 2bfca303b5fd48..c3d90260e2a278 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/4] 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 c3d90260e2a278..8fd1f37acdad6c 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 37999d6fc14ad1..9e00fd1bb5c73c 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<{{.*}}>)

>From 9c6eb4a6d162464f2ebd9ecf98eff42516d14a9d Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Wed, 24 Jan 2024 14:22:01 -0800
Subject: [PATCH 3/4] git-clang-format

---
 mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 8fd1f37acdad6c..0892e93fa9ac42 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -503,17 +503,18 @@ struct GlobalMemrefOpLowering
     Attribute initialValue = nullptr;
     if (!global.isExternal() && !global.isUninitialized()) {
       Attribute initialAttr = *global.getInitialValue();
-      if (auto resourceAttr = llvm::dyn_cast<DenseResourceElementsAttr>(initialAttr)) {
+      if (auto resourceAttr =
+              llvm::dyn_cast<DenseResourceElementsAttr>(initialAttr)) {
         auto blob = resourceAttr.getRawHandle().getBlob();
         initialAttr = DenseElementsAttr::getFromRawBuffer(
-          resourceAttr.getType(), blob->getData());
-      } 
+            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.
+        // 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>();
       }

>From 0d6a9b5b54e0acadfd5ebf35a49c9770d769bdc9 Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Wed, 24 Jan 2024 16:31:40 -0800
Subject: [PATCH 4/4] Added non-splat case

---
 mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
index 9e00fd1bb5c73c..a6e9b44a33e6e2 100644
--- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
@@ -331,14 +331,17 @@ 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>
+  // CHECK: llvm.mlir.global private constant @__constant_xf32_0(1.41421354 : f32) {addr_space = 0 : i32} : f32
+  memref.global "private" constant @__constant_xf32_0 : memref<f32> = dense_resource<NAME0>
+  // CHECK: llvm.mlir.global private constant @__constant_xf32_1(dense<[5.000000e-01, 2.500000e-01]> : tensor<2xf32>) {addr_space = 0 : i32} : !llvm.array<2 x f32>
+  memref.global "private" constant @__constant_xf32_1 : memref<2xf32> = dense_resource<NAME1>
 }
 
 {-#
   dialect_resources: {
     builtin: {
-      NAME: "0x08000000F304B53F"
+      NAME0: "0x08000000F304B53F",
+      NAME1: "0x080000000000003F0000803E"
     }
   }
 #-}