Fixes cases where some lattice points doesn't have expressions for th…

…e next level This fixes #1
tensor-compiler · Jun 6, 2017 · ed500b4 · ed500b4
1 parent cd7b610
commit ed500b4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/src/lower/lower.cpp b/src/lower/lower.cpp
@@ -310,6 +310,11 @@ static vector<Stmt> lower(const Target&    target,
         IndexExpr childExpr = lqExpr;
         Target childTarget = target;
         if (indexVarCase == LAST_FREE || indexVarCase == BELOW_LAST_FREE) {
+          // Extract the expression to compute at the next level. If there's no
+          // computation on the next level for this lattice case then skip it
+          childExpr = getSubExpr(lqExpr, ctx.schedule.getDescendants(child));
+          if (!childExpr.defined()) continue;
+
           // Reduce child expression into temporary
           TensorBase t("t" + child.getName(), ComponentType::Double);
           Expr tensorVar = Var::make(t.getName(), Type(Type::Float,64));
@@ -320,14 +325,10 @@ static vector<Stmt> lower(const Target&    target,
             caseBody.push_back(VarAssign::make(tensorVar, 0.0, true));
           }
 
-          // Extract the expression to compute at the next level
-          childExpr = getSubExpr(lqExpr, ctx.schedule.getDescendants(child));
-
           // Rewrite lqExpr to substitute the expression computed at the next
           // level with the temporary
           lqExpr = replace(lqExpr, {{childExpr,taco::Access(t)}});
         }
-        taco_iassert(childExpr.defined());
         auto childCode = lower::lower(childTarget, childExpr, child, ctx);
         util::append(caseBody, childCode);
       }
@@ -353,17 +354,17 @@ static vector<Stmt> lower(const Target&    target,
         Expr rpos = resultIterator.getPtrVar();
         Stmt posInc = VarAssign::make(rpos, Add::make(rpos, 1));
 
-        // Conditionally resize idx and pos
+        // Conditionally resize result `idx` and `pos` arrays
         if (emitAssemble) {
           Expr resize =
               And::make(Eq::make(0, BitAnd::make(Add::make(rpos, 1), rpos)),
                         Lte::make(ctx.allocSize, Add::make(rpos, 1)));
           Expr newSize = ir::Mul::make(2, ir::Add::make(rpos, 1));
 
-          // Resize `idx`
+          // Resize result `idx` array
           Stmt resizeIndices = resultIterator.resizeIdxStorage(newSize);
 
-          // Resize `pos`
+          // Resize result `pos` array
           if (indexVarCase == ABOVE_LAST_FREE) {
             auto nextStep = resultPath.getStep(resultStep.getStep() + 1);
             Stmt resizePos = ctx.iterators[nextStep].resizePtrStorage(newSize);

diff --git a/test/expr_storage-tests.cpp b/test/expr_storage-tests.cpp
@@ -673,7 +673,6 @@ INSTANTIATE_TEST_CASE_P(axpy_3x3, expr,
                     },
                     {6, 0, 16}
                     ),
-/*
            TestData(Tensor<double>("a",{3},Format({Dense})),
                     {i},
                     d33a("B",Format({Sparse,Sparse}))(i,k) *
@@ -687,7 +686,6 @@ INSTANTIATE_TEST_CASE_P(axpy_3x3, expr,
                     },
                     {6, 0, 16}
                     ),
-*/
            TestData(Tensor<double>("a",{3},Format({Dense})),
                     {i},
                       da("alpha",Format())() *