[mlir][sparse] minor code layout edits

[bixia]

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

23–41

We can remove "static" keywords for these two function now.

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

 //===- SparseStorageSpecifierToLLVM.cpp - convert specifier to llvm -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 #include "SparseTensorStorageLayout.h"
 
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 
 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();
   unsigned rank = enc.getDimLevelType().size();
 
   SmallVector<Type, 2> result;
   auto indexType = tp.getSizesType();
   auto dimSizes = LLVM::LLVMArrayType::get(ctx, indexType, rank);
   auto memSizes = LLVM::LLVMArrayType::get(ctx, indexType,
                                            getNumDataFieldsFromEncoding(enc));
   result.push_back(dimSizes);
   result.push_back(memSizes);
   return result;
 }
 
 static Type convertSpecifier(StorageSpecifierType tp) {
   return LLVM::LLVMStructType::getLiteral(tp.getContext(),
                                           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;
 
 class SpecifierStructBuilder : public StructBuilder {
 public:
   explicit SpecifierStructBuilder(Value specifier) : StructBuilder(specifier) {
     assert(value);

 
 /// Builds IR inserting the pos-th memory size into the descriptor.
 void SpecifierStructBuilder::setMemSize(OpBuilder &builder, Location loc,
                                         unsigned pos, Value size) {
   value = builder.create<LLVM::InsertValueOp>(
       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:
   using OpAdaptor = typename SourceOp::Adaptor;
   using OpConversionPattern<SourceOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(SourceOp op, OpAdaptor adaptor,
matchAndRewrite(StorageSpecifierInitOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type llvmType = getTypeConverter()->convertType(op.getResult().getType());
     rewriter.replaceOp(op, SpecifierStructBuilder::getInitValue(
                                rewriter, op.getLoc(), llvmType));
     return success();
   }
 };
 
+//===----------------------------------------------------------------------===//
+// Public method for populating conversion rules.
+//===----------------------------------------------------------------------===//
+
 void mlir::populateStorageSpecifierToLLVMPatterns(TypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<StorageSpecifierGetOpConverter, StorageSpecifierSetOpConverter,
                StorageSpecifierInitOpConverter>(converter,
                                                 patterns.getContext());
 }

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

 //   ; sparse tensor metadata
 //   struct {
 //     array<rank x int> dimSizes    ; sizes for each dimension
 //     array<n x int> memSizes;      ; sizes for each data memref
 //   }
 // };
 //
 //===----------------------------------------------------------------------===//
+
 enum class SparseTensorFieldKind : uint32_t {
   StorageSpec = 0,
   PtrMemRef = 1,
   IdxMemRef = 2,
   ValMemRef = 3
 };
 
 static_assert(static_cast<uint32_t>(SparseTensorFieldKind::PtrMemRef) ==

 getMutDescriptorFromTensorTuple(Value tensor, SmallVectorImpl<Value> &fields) {
   auto tuple = getTuple(tensor);
   fields.assign(tuple.getInputs().begin(), tuple.getInputs().end());
   return MutSparseTensorDescriptor(tuple.getResultTypes()[0], fields);
 }
 
 } // namespace sparse_tensor
 } // namespace mlir
+
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_

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

 //===- SparseTensorStorageLayout.cpp --------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // 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)
     return builder.create<arith::IndexCastOp>(loc, to, value);
   return value;
 }
 
 static IntegerAttr fromOptionalInt(MLIRContext *ctx,
foreachFieldAndTypeInSparseTensor(
                 DimLevelType /*dlt*/) -> bool {
         assert(fieldIdx == fields.size());
         fields.push_back(fieldType);
         return true;
       });
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// The sparse tensor type converter (defined in Passes.h).
+//===----------------------------------------------------------------------===//
+
 SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });
   addConversion([&](RankedTensorType rtp, SmallVectorImpl<Type> &fields) {
     return convertSparseTensorType(rtp, fields);
   });
 
   // Required by scf.for 1:N type conversion.
   addSourceMaterialization([](OpBuilder &builder, RankedTensorType tp,
                               ValueRange inputs,
                               Location loc) -> std::optional<Value> {
     if (!getSparseTensorEncoding(tp))
       // Not a sparse tensor.
       return std::nullopt;
     // Sparse compiler knows how to cancel out these casts.
     return genTuple(builder, loc, tp, inputs);
   });
 }
 
+//===----------------------------------------------------------------------===//
+// StorageLayout methods.
+//===----------------------------------------------------------------------===//
+
 unsigned StorageLayout::getMemRefFieldIndex(SparseTensorFieldKind kind,
                                             std::optional<unsigned> dim) const {
   unsigned fieldIdx = -1u;
   foreachFieldInSparseTensor(
       enc,
       [dim, kind, &fieldIdx](unsigned fIdx, SparseTensorFieldKind fKind,
                              unsigned fDim, DimLevelType dlt) -> bool {
         if ((dim && fDim == dim.value() && kind == fKind) ||
             (kind == fKind && fKind == SparseTensorFieldKind::ValMemRef)) {
           fieldIdx = fIdx;
           // Returns false to break the iteration.
           return false;
         }
         return true;
       });
   assert(fieldIdx != -1u);
   return fieldIdx;
 }
 
 unsigned StorageLayout::getMemRefFieldIndex(StorageSpecifierKind kind,
                                             std::optional<unsigned> dim) const {
   return getMemRefFieldIndex(toFieldKind(kind), dim);
 }
 
+//===----------------------------------------------------------------------===//
+// StorageTensorSpecifier methods.
+//===----------------------------------------------------------------------===//
+
 Value SparseTensorSpecifier::getInitValue(OpBuilder &builder, Location loc,
                                           RankedTensorType rtp) {
   return builder.create<StorageSpecifierInitOp>(
       loc, StorageSpecifierType::get(getSparseTensorEncoding(rtp)));
 }
 
 Value SparseTensorSpecifier::getSpecifierField(OpBuilder &builder, Location loc,
                                                StorageSpecifierKind kind,
void SparseTensorSpecifier::setSpecifierField(OpBuilder &builder, Location loc,
                                               Value v,
                                               StorageSpecifierKind kind,
                                               std::optional<unsigned> dim) {
   specifier = builder.create<SetStorageSpecifierOp>(
       loc, specifier, kind, fromOptionalInt(specifier.getContext(), dim),
       createIndexCast(builder, loc, v, getFieldType(kind, dim)));
 }
 
+//===----------------------------------------------------------------------===//
+// Public methods.
+//===----------------------------------------------------------------------===//
+
 constexpr uint64_t kDataFieldStartingIdx = 0;
 
 void sparse_tensor::foreachFieldInSparseTensor(
     const SparseTensorEncodingAttr enc,
     llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
                             DimLevelType)>
         callback) {
   assert(enc);