[mlir][sparse] minor code layout edits

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D140934

Author: aartbik

mlir/lib/Dialect/SparseTensor/Transforms/SparseStorageSpecifierToLLVM.cpp

@@ -14,6 +14,12 @@
 using namespace mlir;
 using namespace sparse_tensor;
 
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Helper methods.
+//===----------------------------------------------------------------------===//
+
 static SmallVector<Type, 2> getSpecifierFields(StorageSpecifierType tp) {
   MLIRContext *ctx = tp.getContext();
   auto enc = tp.getEncoding();
@@ -34,10 +40,9 @@ static Type convertSpecifier(StorageSpecifierType tp) {
                                           getSpecifierFields(tp));
 }
 
-StorageSpecifierToLLVMTypeConverter::StorageSpecifierToLLVMTypeConverter() {
-  addConversion([](Type type) { return type; });
-  addConversion([](StorageSpecifierType tp) { return convertSpecifier(tp); });
-}
+//===----------------------------------------------------------------------===//
+// Specifier struct builder.
+//===----------------------------------------------------------------------===//
 
 constexpr uint64_t kDimSizePosInSpecifier = 0;
 constexpr uint64_t kMemSizePosInSpecifier = 1;
@@ -102,6 +107,21 @@ void SpecifierStructBuilder::setMemSize(OpBuilder &builder, Location loc,
       loc, value, size, ArrayRef<int64_t>({kMemSizePosInSpecifier, pos}));
 }
 
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// The sparse storage specifier type converter (defined in Passes.h).
+//===----------------------------------------------------------------------===//
+
+StorageSpecifierToLLVMTypeConverter::StorageSpecifierToLLVMTypeConverter() {
+  addConversion([](Type type) { return type; });
+  addConversion([](StorageSpecifierType tp) { return convertSpecifier(tp); });
+}
+
+//===----------------------------------------------------------------------===//
+// Storage specifier conversion rules.
+//===----------------------------------------------------------------------===//
+
 template <typename Base, typename SourceOp>
 class SpecifierGetterSetterOpConverter : public OpConversionPattern<SourceOp> {
 public:
@@ -176,6 +196,10 @@ struct StorageSpecifierInitOpConverter
   }
 };
 
+//===----------------------------------------------------------------------===//
+// Public method for populating conversion rules.
+//===----------------------------------------------------------------------===//
+
 void mlir::populateStorageSpecifierToLLVMPatterns(TypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<StorageSpecifierGetOpConverter, StorageSpecifierSetOpConverter,

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.cpp

@@ -5,16 +5,22 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+
 #include "SparseTensorStorageLayout.h"
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace sparse_tensor;
 
+//===----------------------------------------------------------------------===//
+// Private helper methods.
+//===----------------------------------------------------------------------===//
+
 static Value createIndexCast(OpBuilder &builder, Location loc, Value value,
                              Type to) {
   if (value.getType() != to)
@@ -47,6 +53,10 @@ convertSparseTensorType(RankedTensorType rtp, SmallVectorImpl<Type> &fields) {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// The sparse tensor type converter (defined in Passes.h).
+//===----------------------------------------------------------------------===//
+
 SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });
   addConversion([&](RankedTensorType rtp, SmallVectorImpl<Type> &fields) {
@@ -65,6 +75,10 @@ SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   });
 }
 
+//===----------------------------------------------------------------------===//
+// StorageLayout methods.
+//===----------------------------------------------------------------------===//
+
 unsigned StorageLayout::getMemRefFieldIndex(SparseTensorFieldKind kind,
                                             std::optional<unsigned> dim) const {
   unsigned fieldIdx = -1u;
@@ -89,6 +103,10 @@ unsigned StorageLayout::getMemRefFieldIndex(StorageSpecifierKind kind,
   return getMemRefFieldIndex(toFieldKind(kind), dim);
 }
 
+//===----------------------------------------------------------------------===//
+// StorageTensorSpecifier methods.
+//===----------------------------------------------------------------------===//
+
 Value SparseTensorSpecifier::getInitValue(OpBuilder &builder, Location loc,
                                           RankedTensorType rtp) {
   return builder.create<StorageSpecifierInitOp>(
@@ -114,6 +132,10 @@ void SparseTensorSpecifier::setSpecifierField(OpBuilder &builder, Location loc,
       createIndexCast(builder, loc, v, getFieldType(kind, dim)));
 }
 
+//===----------------------------------------------------------------------===//
+// Public methods.
+//===----------------------------------------------------------------------===//
+
 constexpr uint64_t kDataFieldStartingIdx = 0;
 
 void sparse_tensor::foreachFieldInSparseTensor(

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

@@ -50,6 +50,7 @@ namespace sparse_tensor {
 // };
 //
 //===----------------------------------------------------------------------===//
+
 enum class SparseTensorFieldKind : uint32_t {
   StorageSpec = 0,
   PtrMemRef = 1,
@@ -355,4 +356,5 @@ getMutDescriptorFromTensorTuple(Value tensor, SmallVectorImpl<Value> &fields) {
 
 } // namespace sparse_tensor
 } // namespace mlir
+
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_