rebase

Author: Peiming Liu

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

@@ -82,14 +82,14 @@ struct SparsificationPass
   SparsificationPass(const SparsificationPass &pass) = default;
   SparsificationPass(const SparsificationOptions &options) {
     parallelization = options.parallelizationStrategy;
-    debugSparseIteration = options.debugSparseIteration;
+    sparseEmitStrategy = options.sparseEmitStrategy;
     enableRuntimeLibrary = options.enableRuntimeLibrary;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     // Translate strategy flags to strategy options.
-    SparsificationOptions options(parallelization, debugSparseIteration,
+    SparsificationOptions options(parallelization, sparseEmitStrategy,
                                   enableRuntimeLibrary);
     // Apply sparsification and cleanup rewriting.
     RewritePatternSet patterns(ctx);

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

@@ -1369,7 +1369,7 @@ struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
       return failure();
 
     // Recursively generates code if admissible.
-    env.startEmit(options.debugSparseIteration);
+    env.startEmit(options.sparseEmitStrategy);
     genBuffers(env, rewriter);
     // TODO: Constant affine expression should be handled differently when using
     // slice-based codegen, it does not matter now because we already reject the

mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.cpp

@@ -59,7 +59,7 @@ LogicalResult CodegenEnv::initTensorExp() {
   return success();
 }
 
-void CodegenEnv::startEmit(DebugSparseIteration emitStrategy) {
+void CodegenEnv::startEmit(SparseEmitStrategy emitStrategy) {
   assert(insChain == nullptr && "must only start emitting once");
   if (sparseOut) {
     insChain = sparseOut->get();

mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.h

@@ -52,7 +52,7 @@ class CodegenEnv {
   Merger &merger() { return latticeMerger; }
   LoopEmitter &emitter() { return loopEmitter; }
 
-  void startEmit(DebugSparseIteration emitStrategy);
+  void startEmit(SparseEmitStrategy emitStrategy);
 
   /// Generates loop boundary statements (entering/exiting loops). The function
   /// passes and updates the passed-in parameters.

mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp

@@ -82,19 +82,19 @@ static Value genSliceStride(OpBuilder &builder, Location loc, Value tensor,
 LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
                          bool isSparseOut, unsigned numLoops,
                          DependentLvlGetter dimGetter,
-                         DebugSparseIteration emitStrategy) {
+                         SparseEmitStrategy emitStrategy) {
   initialize(tensors, loopTag, hasOutput, isSparseOut, numLoops, dimGetter);
 }
 
 void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
                              bool isSparseOut, unsigned numLoops,
                              DependentLvlGetter dimGetter,
-                             DebugSparseIteration emitStrategy) {
+                             SparseEmitStrategy emitStrategy) {
   // First initialize the top-level type of the fields.
   this->loopTag = loopTag;
   this->hasOutput = hasOutput;
   this->isSparseOut = isSparseOut;
-  SparseIterator::setDebugSparseIteration(emitStrategy);
+  SparseIterator::setSparseEmitStrategy(emitStrategy);
 
   const unsigned numManifestTensors = ts.size();
   const unsigned synTensorId = numManifestTensors;

mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h

@@ -89,13 +89,13 @@ class LoopEmitter {
   initialize(ValueRange tensors, StringAttr loopTag = nullptr,
              bool hasOutput = false, bool isSparseOut = false,
              unsigned numLoops = 0, DependentLvlGetter getter = nullptr,
-             DebugSparseIteration emitStrategy = DebugSparseIteration::kNone);
+             SparseEmitStrategy emitStrategy = SparseEmitStrategy::kFunctional);
 
   explicit LoopEmitter(
       ValueRange tensors, StringAttr loopTag = nullptr, bool hasOutput = false,
       bool isSparseOut = false, unsigned numLoops = 0,
       DependentLvlGetter getter = nullptr,
-      DebugSparseIteration emitStrategy = DebugSparseIteration::kNone);
+      SparseEmitStrategy emitStrategy = SparseEmitStrategy::kFunctional);
 
   /// Starts a loop emitting session by generating all the buffers needed
   /// for iterating over the tensors.

mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorLevel.cpp

@@ -777,11 +777,12 @@ class SubSectIterator : public SparseIterator {
 // SparseIterator derived classes implementation.
 //===----------------------------------------------------------------------===//
 
-DebugSparseIteration SparseIterator::emitStrategy = DebugSparseIteration::kNone;
+SparseEmitStrategy SparseIterator::emitStrategy =
+    SparseEmitStrategy::kFunctional;
 
 void SparseIterator::genInit(OpBuilder &b, Location l,
                              const SparseIterator *p) {
-  if (emitStrategy == DebugSparseIteration::kInterfaceOnly) {
+  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *begin = b.create(l, b.getStringAttr(prefix + ".begin"), {},
                                 getCursorValTypes(b));
@@ -793,7 +794,7 @@ void SparseIterator::genInit(OpBuilder &b, Location l,
 }
 
 Value SparseIterator::genNotEnd(OpBuilder &b, Location l) {
-  if (emitStrategy == DebugSparseIteration::kInterfaceOnly) {
+  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *notEnd = b.create(l, b.getStringAttr(prefix + ".not_end"),
                                  getCursor(), b.getI1Type());
@@ -804,7 +805,7 @@ Value SparseIterator::genNotEnd(OpBuilder &b, Location l) {
 }
 
 void SparseIterator::locate(OpBuilder &b, Location l, Value crd) {
-  if (emitStrategy == DebugSparseIteration::kInterfaceOnly) {
+  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     SmallVector<Value> args = getCursor();
     args.push_back(crd);
@@ -818,7 +819,7 @@ void SparseIterator::locate(OpBuilder &b, Location l, Value crd) {
 }
 
 Value SparseIterator::deref(OpBuilder &b, Location l) {
-  if (emitStrategy == DebugSparseIteration::kInterfaceOnly) {
+  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     SmallVector<Value> args = getCursor();
     Operation *deref = b.create(l, b.getStringAttr(prefix + ".deref"),
@@ -830,7 +831,7 @@ Value SparseIterator::deref(OpBuilder &b, Location l) {
 }
 
 ValueRange SparseIterator::forward(OpBuilder &b, Location l) {
-  if (emitStrategy == DebugSparseIteration::kInterfaceOnly) {
+  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *next = b.create(l, b.getStringAttr(prefix + ".next"),
                                getCursor(), getCursorValTypes(b));

mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorLevel.h

@@ -111,7 +111,7 @@ class SparseIterator {
 public:
   virtual ~SparseIterator() = default;
 
-  static void setDebugSparseIteration(DebugSparseIteration strategy) {
+  static void setSparseEmitStrategy(SparseEmitStrategy strategy) {
     SparseIterator::emitStrategy = strategy;
   }
 
@@ -247,7 +247,7 @@ class SparseIterator {
     return ref.take_front(cursorValsCnt);
   }
 
-  static DebugSparseIteration emitStrategy;
+  static SparseEmitStrategy emitStrategy;
 
 public:
   const IterKind kind;     // For LLVM-style RTTI.

mlir/test/Dialect/SparseTensor/sparse_conv_2d_slice_based.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --sparse-reinterpret-map --sparsification="debug-sparse-iteration=interface-only" --canonicalize --cse --allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s --sparse-reinterpret-map --sparsification="sparse-emit-strategy=debug-interface" --canonicalize --cse --allow-unregistered-dialect | FileCheck %s
 
 #map = affine_map<(d0, d1, d2, d3) -> (d0 + d2, d1 + d3)>
 #map1 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>