[Mlir-commits] [mlir] [mlir][sparse] rename map utility (PR #68611)

[Aart Bik]

https://github.com/aartbik updated https://github.com/llvm/llvm-project/pull/68611

>From 2da1792eeafcf5e0236c93465ab545d26c93babc Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 9 Oct 2023 09:53:27 -0700
Subject: [PATCH 1/2] [mlir][sparse] rename map utility

Rename util genReaderBuffers -> genMapBuffers since it is no longer
specific to the reader, but all MapRef data in general.
---
 .../Dialect/SparseTensor/Transforms/CodegenUtils.cpp | 12 ++++++------
 .../Dialect/SparseTensor/Transforms/CodegenUtils.h   | 10 +++++-----
 .../SparseTensor/Transforms/SparseTensorCodegen.cpp  |  4 ++--
 .../Transforms/SparseTensorConversion.cpp            |  6 +++---
 4 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
index c44c5985ee7bba1..dac6d6b64551c83 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
@@ -774,12 +774,12 @@ Value sparse_tensor::genReader(OpBuilder &builder, Location loc,
   return reader;
 }
 
-Value sparse_tensor::genReaderBuffers(OpBuilder &builder, Location loc,
-                                      SparseTensorType stt,
-                                      ArrayRef<Value> dimShapesValues,
-                                      Value dimSizesBuffer,
-                                      /*out*/ Value &dim2lvlBuffer,
-                                      /*out*/ Value &lvl2dimBuffer) {
+Value sparse_tensor::genMapBuffers(OpBuilder &builder, Location loc,
+                                   SparseTensorType stt,
+                                   ArrayRef<Value> dimShapesValues,
+                                   Value dimSizesBuffer,
+                                   /*out*/ Value &dim2lvlBuffer,
+                                   /*out*/ Value &lvl2dimBuffer) {
   const Dimension dimRank = stt.getDimRank();
   const Level lvlRank = stt.getLvlRank();
   // For an identity mapping, the dim2lvl and lvl2dim mappings are
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
index 698b6c491a9aef7..1562ea3f20f73df 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
@@ -353,11 +353,11 @@ Value genReader(OpBuilder &builder, Location loc, SparseTensorType stt,
                 /*out*/ SmallVectorImpl<Value> &dimShapeValues,
                 /*out*/ Value &dimSizesBuffer);
 
-/// Generates code to set up the buffer parameters for a reader.
-Value genReaderBuffers(OpBuilder &builder, Location loc, SparseTensorType stt,
-                       ArrayRef<Value> dimShapeValues, Value dimSizesBuffer,
-                       /*out*/ Value &dim2lvlBuffer,
-                       /*out*/ Value &lvl2dimBuffer);
+/// Generates code to set up the buffer parameters for a map.
+Value genMapBuffers(OpBuilder &builder, Location loc, SparseTensorType stt,
+                    ArrayRef<Value> dimShapeValues, Value dimSizesBuffer,
+                    /*out*/ Value &dim2lvlBuffer,
+                    /*out*/ Value &lvl2dimBuffer);
 
 //===----------------------------------------------------------------------===//
 // Inlined constant generators.
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 2c03f0a6020e6a8..e22789643c90af7 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1478,8 +1478,8 @@ struct SparseNewConverter : public OpConversionPattern<NewOp> {
     // Now construct the dim2lvl and lvl2dim buffers.
     Value dim2lvlBuffer;
     Value lvl2dimBuffer;
-    genReaderBuffers(rewriter, loc, dstTp, dimShapesValues, dimSizesBuffer,
-                     dim2lvlBuffer, lvl2dimBuffer);
+    genMapBuffers(rewriter, loc, dstTp, dimShapesValues, dimSizesBuffer,
+                  dim2lvlBuffer, lvl2dimBuffer);
 
     // Read the COO tensor data.
     Value xs = desc.getAOSMemRef();
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index e44d8565fc867de..68b337e2d2f3f04 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -205,7 +205,7 @@ class NewCallParams final {
     params[kParamLvlTypes] = genLvlTypesBuffer(builder, loc, stt);
     // Construct dimSizes, lvlSizes, dim2lvl, and lvl2dim buffers.
     params[kParamDimSizes] = allocaBuffer(builder, loc, dimSizesValues);
-    params[kParamLvlSizes] = genReaderBuffers(
+    params[kParamLvlSizes] = genMapBuffers(
         builder, loc, stt, dimSizesValues, params[kParamDimSizes],
         params[kParamDim2Lvl], params[kParamLvl2Dim]);
     // Secondary and primary types encoding.
@@ -446,8 +446,8 @@ class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
     Value dim2lvlBuffer;
     Value lvl2dimBuffer;
     Value lvlSizesBuffer =
-        genReaderBuffers(rewriter, loc, stt, dimShapesValues, dimSizesBuffer,
-                         dim2lvlBuffer, lvl2dimBuffer);
+        genMapBuffers(rewriter, loc, stt, dimShapesValues, dimSizesBuffer,
+                      dim2lvlBuffer, lvl2dimBuffer);
     // Use the `reader` to parse the file.
     Type opaqueTp = getOpaquePointerType(rewriter);
     Type eltTp = stt.getElementType();

>From 6bde50783c4b7c66795c09f0b3e2568ac9891488 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 9 Oct 2023 10:01:20 -0700
Subject: [PATCH 2/2] clang-format of course

---
 .../SparseTensor/Transforms/SparseTensorConversion.cpp      | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index 68b337e2d2f3f04..d2d7b46ab834e71 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -205,9 +205,9 @@ class NewCallParams final {
     params[kParamLvlTypes] = genLvlTypesBuffer(builder, loc, stt);
     // Construct dimSizes, lvlSizes, dim2lvl, and lvl2dim buffers.
     params[kParamDimSizes] = allocaBuffer(builder, loc, dimSizesValues);
-    params[kParamLvlSizes] = genMapBuffers(
-        builder, loc, stt, dimSizesValues, params[kParamDimSizes],
-        params[kParamDim2Lvl], params[kParamLvl2Dim]);
+    params[kParamLvlSizes] =
+        genMapBuffers(builder, loc, stt, dimSizesValues, params[kParamDimSizes],
+                      params[kParamDim2Lvl], params[kParamLvl2Dim]);
     // Secondary and primary types encoding.
     setTemplateTypes(stt);
     // Finally, make note that initialization is complete.