Merge pull request #491 from tensor-compiler/ell_fix

stephenchouca · web-flow · commit 4c24193f8da3 · 2022-05-12T13:26:35.000-04:00
Fixed support for ELL format
diff --git a/include/taco/format.h b/include/taco/format.h
@@ -97,7 +97,8 @@ class ModeFormat {
   /// Properties of a mode format
   enum Property {
     FULL, NOT_FULL, ORDERED, NOT_ORDERED, UNIQUE, NOT_UNIQUE, BRANCHLESS,
-    NOT_BRANCHLESS, COMPACT, NOT_COMPACT, ZEROLESS, NOT_ZEROLESS
+    NOT_BRANCHLESS, COMPACT, NOT_COMPACT, ZEROLESS, NOT_ZEROLESS, PADDED, 
+    NOT_PADDED
   };
 
   /// Instantiates an undefined mode format
@@ -129,6 +130,7 @@ class ModeFormat {
   bool isBranchless() const;
   bool isCompact() const;
   bool isZeroless() const;
+  bool isPadded() const;
 
   /// Returns true if a mode format has a specific capability, false otherwise
   bool hasCoordValIter() const;
diff --git a/include/taco/lower/mode_format_impl.h b/include/taco/lower/mode_format_impl.h
@@ -106,9 +106,10 @@ class ModeFormatImpl {
 public:
   ModeFormatImpl(std::string name, bool isFull, bool isOrdered, bool isUnique, 
                  bool isBranchless, bool isCompact, bool isZeroless, 
-                 bool hasCoordValIter, bool hasCoordPosIter, bool hasLocate, 
-                 bool hasInsert, bool hasAppend, bool hasSeqInsertEdge, 
-                 bool hasInsertCoord, bool isYieldPosPure);
+		 bool isPadded, bool hasCoordValIter, bool hasCoordPosIter, 
+		 bool hasLocate, bool hasInsert, bool hasAppend, 
+		 bool hasSeqInsertEdge, bool hasInsertCoord, 
+		 bool isYieldPosPure);
 
   virtual ~ModeFormatImpl();
 
@@ -246,6 +247,7 @@ class ModeFormatImpl {
   const bool isBranchless;
   const bool isCompact;
   const bool isZeroless;
+  const bool isPadded;
   
   const bool hasCoordValIter;
   const bool hasCoordPosIter;
diff --git a/include/taco/lower/mode_format_singleton.h b/include/taco/lower/mode_format_singleton.h
@@ -10,8 +10,9 @@ class SingletonModeFormat : public ModeFormatImpl {
   using ModeFormatImpl::getInsertCoord;
 
   SingletonModeFormat();
-  SingletonModeFormat(bool isFull, bool isOrdered,
-                      bool isUnique, bool isZeroless, long long allocSize = DEFAULT_ALLOC_SIZE);
+  SingletonModeFormat(bool isFull, bool isOrdered, bool isUnique, 
+                      bool isZeroless, bool isPadded, 
+		      long long allocSize = DEFAULT_ALLOC_SIZE);
 
   ~SingletonModeFormat() override {}
 
diff --git a/src/format.cpp b/src/format.cpp
@@ -187,6 +187,11 @@ bool ModeFormat::hasProperties(const std::vector<Property>& properties) const {
           return false;
         }
         break;	
+      case PADDED:
+        if (!isPadded()) {
+          return false;
+        }
+        break;	
       case NOT_FULL:
         if (isFull()) {
           return false;
@@ -217,6 +222,11 @@ bool ModeFormat::hasProperties(const std::vector<Property>& properties) const {
           return false;
         }
         break;
+      case NOT_PADDED:
+        if (isPadded()) {
+          return false;
+        }
+        break;
     }
   }
   return true;
@@ -252,6 +262,11 @@ bool ModeFormat::isZeroless() const {
   return impl->isZeroless;
 }
 
+bool ModeFormat::isPadded() const {
+  taco_iassert(defined());
+  return impl->isPadded;
+}
+
 bool ModeFormat::hasCoordValIter() const {
   taco_iassert(defined());
   return impl->hasCoordValIter;
diff --git a/src/index_notation/index_notation.cpp b/src/index_notation/index_notation.cpp
@@ -4155,6 +4155,22 @@ IndexStmt generatePackStmt(TensorVar tensor,
     packStmt = forall(indexVars[mode], packStmt);
   }
 
+  bool doAppend = true;
+  const Format lhsFormat = otherIsOnRight ? format : otherFormat;
+  for (int i = lhsFormat.getOrder() - 1; i >= 0; --i) {
+    const auto modeFormat = lhsFormat.getModeFormats()[i];
+    if (modeFormat.isBranchless() && i != 0) {
+      const auto parentModeFormat = lhsFormat.getModeFormats()[i - 1];
+      if (parentModeFormat.isUnique() || !parentModeFormat.hasAppend()) {
+        doAppend = false;
+        break;
+      }
+    }
+  }
+  if (!doAppend) {
+    packStmt = packStmt.assemble(otherIsOnRight ? tensor : other, AssembleStrategy::Insert);
+  }
+
   return packStmt; 
 }
 
diff --git a/src/ir/ir_rewriter.cpp b/src/ir/ir_rewriter.cpp
@@ -435,7 +435,7 @@ void IRRewriter::visit(const Allocate* op) {
     stmt = op;
   }
   else {
-    stmt = Allocate::make(var, num_elements, op->is_realloc, op->old_elements);
+    stmt = Allocate::make(var, num_elements, op->is_realloc, op->old_elements, op->clear);
   }
 }
 
diff --git a/src/lower/lowerer_impl_imperative.cpp b/src/lower/lowerer_impl_imperative.cpp
@@ -1391,24 +1391,28 @@ Stmt LowererImplImperative::lowerForallPosition(Forall forall, Iterator iterator
     endBound = endBounds[1];
   }
 
-  LoopKind kind = LoopKind::Serial;
-  if (forall.getParallelUnit() == ParallelUnit::CPUVector && !ignoreVectorize) {
-    kind = LoopKind::Vectorized;
-  }
-  else if (forall.getParallelUnit() != ParallelUnit::NotParallel
-           && forall.getOutputRaceStrategy() != OutputRaceStrategy::ParallelReduction && !ignoreVectorize) {
-    kind = LoopKind::Runtime;
+  Stmt loop = Block::make(strideGuard, declareCoordinate, boundsGuard, body);
+  if (iterator.isBranchless() && iterator.isCompact() && 
+      (iterator.getParent().isRoot() || iterator.getParent().isUnique())) {
+    loop = Block::make(VarDecl::make(iterator.getPosVar(), startBound), loop);
+  } else {
+    LoopKind kind = LoopKind::Serial;
+    if (forall.getParallelUnit() == ParallelUnit::CPUVector && !ignoreVectorize) {
+      kind = LoopKind::Vectorized;
+    }
+    else if (forall.getParallelUnit() != ParallelUnit::NotParallel && 
+	     forall.getOutputRaceStrategy() != OutputRaceStrategy::ParallelReduction && 
+	     !ignoreVectorize) {
+      kind = LoopKind::Runtime;
+    }
+
+    loop = For::make(iterator.getPosVar(), startBound, endBound, 1, loop, kind,
+                     ignoreVectorize ? ParallelUnit::NotParallel : forall.getParallelUnit(), 
+		     ignoreVectorize ? 0 : forall.getUnrollFactor());
   }
 
   // Loop with preamble and postamble
-  return Block::blanks(
-                       boundsCompute,
-                       For::make(iterator.getPosVar(), startBound, endBound, 1,
-                                 Block::make(strideGuard, declareCoordinate, boundsGuard, body),
-                                 kind,
-                                 ignoreVectorize ? ParallelUnit::NotParallel : forall.getParallelUnit(), ignoreVectorize ? 0 : forall.getUnrollFactor()),
-                       posAppend);
-
+  return Block::blanks(boundsCompute, loop, posAppend);
 }
 
 Stmt LowererImplImperative::lowerForallFusedPosition(Forall forall, Iterator iterator,
diff --git a/src/lower/mode_format_compressed.cpp b/src/lower/mode_format_compressed.cpp
@@ -17,8 +17,8 @@ CompressedModeFormat::CompressedModeFormat(bool isFull, bool isOrdered,
                                            bool isUnique, bool isZeroless, 
                                            long long allocSize) :
     ModeFormatImpl("compressed", isFull, isOrdered, isUnique, false, true,
-                   isZeroless, false, true, false, false, true, true, true, 
-                   false), 
+                   isZeroless, false, false, true, false, false, true, true,  
+                   true, false), 
     allocSize(allocSize) {
 }
 
diff --git a/src/lower/mode_format_dense.cpp b/src/lower/mode_format_dense.cpp
@@ -11,7 +11,7 @@ DenseModeFormat::DenseModeFormat() : DenseModeFormat(true, true, false) {
 DenseModeFormat::DenseModeFormat(const bool isOrdered, const bool isUnique, 
                                  const bool isZeroless) : 
     ModeFormatImpl("dense", true, isOrdered, isUnique, false, true, isZeroless, 
-                   false, false, true, true, false, false, false, true) {
+                   true, false, false, true, true, false, false, false, true) {
 }
 
 ModeFormat DenseModeFormat::copy(
diff --git a/src/lower/mode_format_impl.cpp b/src/lower/mode_format_impl.cpp
@@ -147,15 +147,16 @@ std::ostream& operator<<(std::ostream& os, const ModeFunction& modeFunction) {
 // class ModeTypeImpl
 ModeFormatImpl::ModeFormatImpl(const std::string name, bool isFull, 
                                bool isOrdered, bool isUnique, bool isBranchless, 
-                               bool isCompact, bool isZeroless, 
+                               bool isCompact, bool isZeroless, bool isPadded,
                                bool hasCoordValIter, bool hasCoordPosIter, 
                                bool hasLocate, bool hasInsert, bool hasAppend, 
                                bool hasSeqInsertEdge, bool hasInsertCoord,
                                bool isYieldPosPure) :
     name(name), isFull(isFull), isOrdered(isOrdered), isUnique(isUnique),
     isBranchless(isBranchless), isCompact(isCompact), isZeroless(isZeroless),
-    hasCoordValIter(hasCoordValIter), hasCoordPosIter(hasCoordPosIter),
-    hasLocate(hasLocate), hasInsert(hasInsert), hasAppend(hasAppend),
+    isPadded(isPadded), hasCoordValIter(hasCoordValIter),
+    hasCoordPosIter(hasCoordPosIter), hasLocate(hasLocate),
+    hasInsert(hasInsert), hasAppend(hasAppend), 
     hasSeqInsertEdge(hasSeqInsertEdge), hasInsertCoord(hasInsertCoord),
     isYieldPosPure(isYieldPosPure) {
 }
diff --git a/src/lower/mode_format_singleton.cpp b/src/lower/mode_format_singleton.cpp
@@ -10,15 +10,15 @@ using namespace taco::ir;
 namespace taco {
 
 SingletonModeFormat::SingletonModeFormat() : 
-    SingletonModeFormat(false, true, true, false) {
+    SingletonModeFormat(false, true, true, false, false) {
 }
 
 SingletonModeFormat::SingletonModeFormat(bool isFull, bool isOrdered,
                                          bool isUnique, bool isZeroless,
-                                         long long allocSize) :
+                                         bool isPadded, long long allocSize) :
     ModeFormatImpl("singleton", isFull, isOrdered, isUnique, true, true,
-                   isZeroless, false, true, false, false, true, false, true, 
-                   true), 
+                   isZeroless, isPadded, false, true, false, false, true, 
+		   false, true, true), 
     allocSize(allocSize) {
 }
 
@@ -28,6 +28,7 @@ ModeFormat SingletonModeFormat::copy(
   bool isOrdered = this->isOrdered;
   bool isUnique = this->isUnique;
   bool isZeroless = this->isZeroless;
+  bool isPadded = this->isPadded;
   for (const auto property : properties) {
     switch (property) {
       case ModeFormat::FULL:
@@ -54,13 +55,19 @@ ModeFormat SingletonModeFormat::copy(
       case ModeFormat::NOT_ZEROLESS:
         isZeroless = false;
         break;	
+      case ModeFormat::PADDED:
+        isPadded = true;
+        break;
+      case ModeFormat::NOT_PADDED:
+        isPadded = false;
+        break;	
       default:
         break;
     }
   }
   const auto singletonVariant = 
       std::make_shared<SingletonModeFormat>(isFull, isOrdered, isUnique, 
-                                            isZeroless);
+                                            isZeroless, isPadded);
   return ModeFormat(singletonVariant);
 }
 
@@ -128,7 +135,7 @@ Expr SingletonModeFormat::getAssembledSize(Expr prevSize, Mode mode) const {
 Stmt SingletonModeFormat::getInitCoords(Expr prevSize, 
     std::vector<AttrQueryResult> queries, Mode mode) const {
   Expr crdArray = getCoordArray(mode.getModePack());
-  return Allocate::make(crdArray, prevSize, false, Expr());
+  return Allocate::make(crdArray, prevSize, false, Expr(), isPadded);
 }
 
 ModeFunction SingletonModeFormat::getYieldPos(Expr parentPos, 
diff --git a/src/tensor.cpp b/src/tensor.cpp
@@ -965,6 +965,7 @@ TensorBase::getHelperFunctions(const Format& format, Datatype ctype,
     TensorVar packedTensor(Type(ctype, Shape(dims)), format);
 
     // Define packing and iterator routines in index notation.
+    // TODO: Use `generatePackCOOStmt` function to generate pack routine.
     std::vector<IndexVar> indexVars(format.getOrder());
     IndexStmt packStmt = (packedTensor(indexVars) = bufferTensor(indexVars));
     IndexStmt iterateStmt = Yield(indexVars, packedTensor(indexVars));
@@ -974,6 +975,21 @@ TensorBase::getHelperFunctions(const Format& format, Datatype ctype,
       iterateStmt = forall(indexVars[mode], iterateStmt);
     }
 
+    bool doAppend = true;
+    for (int i = format.getOrder() - 1; i >= 0; --i) {
+      const auto modeFormat = format.getModeFormats()[i];
+      if (modeFormat.isBranchless() && i != 0) {
+        const auto parentModeFormat = format.getModeFormats()[i - 1];
+        if (parentModeFormat.isUnique() || !parentModeFormat.hasAppend()) {
+          doAppend = false;
+          break;
+        }
+      }
+    }
+    if (!doAppend) {
+      packStmt = packStmt.assemble(packedTensor, AssembleStrategy::Insert);
+    }
+
     // Lower packing and iterator code.
     helperModule->addFunction(lower(packStmt, "pack", true, true));
     helperModule->addFunction(lower(iterateStmt, "iterate", false, true));
diff --git a/test/test_tensors.cpp b/test/test_tensors.cpp
@@ -137,6 +137,16 @@ TensorData<double> d5d_data() {
   });
 }
 
+TensorData<double> d5e_data() {
+  return TensorData<double>({5}, {
+    {{0}, 1},
+    {{1}, 2},
+    {{2}, 3},
+    {{3}, 4},
+    {{4}, 5}
+  });
+}
+
 TensorData<double> d8a_data() {
   return TensorData<double>({8}, {
     {{0}, 1},
@@ -328,6 +338,23 @@ TensorData<double> d333a_data() {
   });
 }
 
+TensorData<double> d355a_data() {
+  return TensorData<double>({3,5,5}, {
+    {{0,0,0}, 1},
+    {{0,1,1}, 2},
+    {{0,2,1}, 3},
+    {{0,3,1}, 4},
+    {{0,4,1}, 5},
+    {{1,0,1}, 6},
+    {{1,1,0}, 7},
+    {{1,2,0}, 8},
+    {{1,4,2}, 9},
+    {{2,1,2}, 10},
+    {{2,2,3}, 11},
+    {{2,4,4}, 12},
+  });
+}
+
 TensorData<double> d32b_data() {
   return TensorData<double>({3,2}, {
     {{0,0}, 10},
@@ -406,6 +433,10 @@ Tensor<double> d5d(std::string name, Format format) {
   return d5d_data().makeTensor(name, format);
 }
 
+Tensor<double> d5e(std::string name, Format format) {
+  return d5e_data().makeTensor(name, format);
+}
+
 Tensor<double> d8a(std::string name, Format format) {
   return d8a_data().makeTensor(name, format);
 }
@@ -486,6 +517,10 @@ Tensor<double> d333a(std::string name, Format format) {
   return d333a_data().makeTensor(name, format);
 }
 
+Tensor<double> d355a(std::string name, Format format) {
+  return d355a_data().makeTensor(name, format);
+}
+
 Tensor<double> d32b(std::string name, Format format) {
   return d32b_data().makeTensor(name, format);
 }
diff --git a/test/test_tensors.h b/test/test_tensors.h
@@ -101,6 +101,7 @@ TensorData<double> d5a_data();
 TensorData<double> d5b_data();
 TensorData<double> d5c_data();
 TensorData<double> d5d_data();
+TensorData<double> d5e_data();
 
 TensorData<double> d8a_data();
 TensorData<double> d8b_data();
@@ -127,6 +128,8 @@ TensorData<double> d233c_data();
 
 TensorData<double> d333a_data();
 
+TensorData<double> d355a_data();
+
 TensorData<double> d32b_data();
 TensorData<double> d3322a_data();
 
@@ -146,6 +149,7 @@ Tensor<double> d5a(std::string name, Format format);
 Tensor<double> d5b(std::string name, Format format);
 Tensor<double> d5c(std::string name, Format format);
 Tensor<double> d5d(std::string name, Format format);
+Tensor<double> d5e(std::string name, Format format);
 
 Tensor<double> d8a(std::string name, Format format);
 Tensor<double> d8b(std::string name, Format format);
@@ -175,6 +179,8 @@ Tensor<double> d233c(std::string name, Format format);
 
 Tensor<double> d333a(std::string name, Format format);
 
+Tensor<double> d355a(std::string name, Format format);
+
 Tensor<double> d32b(std::string name, Format format);
 Tensor<double> d3322a(std::string name, Format format);
 
diff --git a/test/tests-expr_storage.cpp b/test/tests-expr_storage.cpp
diff --git a/test/tests-merge_lattice.cpp b/test/tests-merge_lattice.cpp
diff --git a/tools/taco.cpp b/tools/taco.cpp

Original file line number	Diff line number	Diff line change
`@@ -435,7 +435,7 @@ void IRRewriter::visit(const Allocate* op) {`
`435`	`435`	`stmt = op;`
`436`	`436`	`}`
`437`	`437`	`else {`
`438`		`- stmt = Allocate::make(var, num_elements, op->is_realloc, op->old_elements);`
	`438`	`+ stmt = Allocate::make(var, num_elements, op->is_realloc, op->old_elements, op->clear);`
`439`	`439`	`}`
`440`	`440`	`}`
`441`	`441`
Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,7 @@ DenseModeFormat::DenseModeFormat() : DenseModeFormat(true, true, false) {`
`11`	`11`	`DenseModeFormat::DenseModeFormat(const bool isOrdered, const bool isUnique,`
`12`	`12`	`const bool isZeroless) :`
`13`	`13`	`ModeFormatImpl("dense", true, isOrdered, isUnique, false, true, isZeroless,`
`14`		`- false, false, true, true, false, false, false, true) {`
	`14`	`+ true, false, false, true, true, false, false, false, true) {`
`15`	`15`	`}`
`16`	`16`
`17`	`17`	`ModeFormat DenseModeFormat::copy(`