Fixed chain tensor multiplication and factorization optimization (#38)

Fixed chain tensor multiplication and factorization optimization (#32)

include/comet/Dialect/Utils/Utils.h

@@ -108,6 +108,7 @@ namespace mlir
 
     std::string dump2str(Value t);
     std::vector<std::string> stringSplit(std::string s, std::string delimiter);
+    std::vector<unsigned> getReverseIdentityPermutation(size_t size);
     std::vector<unsigned> getIdentityPermutation(size_t size);
 
     std::vector<std::vector<int64_t>> getAllPerms(ArrayAttr indexMaps);

integration_test/compound_exps/Dense_chain_mult_matrix.ta

@@ -0,0 +1,30 @@
+# RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> Dense_chain_mult_matrix.llvm
+# RUN: mlir-cpu-runner Dense_chain_mult_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%mlir_utility_library_dir/libmlir_runner_utils%shlibext,%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s
+
+
+def main() {
+	#IndexLabel Declarations
+	IndexLabel [i] = [2];
+	IndexLabel [j] = [2];  
+	IndexLabel [k] = [5];           
+	IndexLabel [l] = [2];           
+
+	#Tensor Declarations
+	Tensor<double> A([i, j], {Dense});	  
+	Tensor<double> B([j, k], {Dense});
+	Tensor<double> C([k, l], {Dense});
+	Tensor<double> D([i, l], {Dense});
+
+	#Tensor Fill Operation
+	A[i, j] = 2.2;
+	B[j, k] = 3.4;
+	C[k, l] = 1.0;
+	D[i, l] = 0.0;
+
+	D[i, l] = A[i, j] * B[j, k] * C[k,l];
+	print(D);
+}
+
+# Print the result for verification.
+# CHECK: data = 
+# CHECK-NEXT: 74.8,74.8,74.8,74.8,

integration_test/opts/chain_mult_factorize.ta

@@ -0,0 +1,30 @@
+# RUN: comet-opt --opt-multiop-factorize --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> chain_mult_factorize.llvm
+# RUN: mlir-cpu-runner chain_mult_factorize.llvm -O3 -e main -entry-point-result=void -shared-libs=%mlir_utility_library_dir/libmlir_runner_utils%shlibext,%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s
+
+
+def main() {
+	#IndexLabel Declarations
+	IndexLabel [i] = [2];
+	IndexLabel [j] = [2];  
+	IndexLabel [k] = [5];           
+	IndexLabel [l] = [2];           
+
+	#Tensor Declarations
+	Tensor<double> A([i, j], {Dense});	  
+	Tensor<double> B([j, k], {Dense});
+	Tensor<double> C([k, l], {Dense});
+	Tensor<double> D([i, l], {Dense});
+
+	#Tensor Fill Operation
+	A[i, j] = 2.2;
+	B[j, k] = 3.4;
+	C[k, l] = 1.0;
+	D[i, l] = 0.0;
+
+	D[i, l] = A[i, j] * B[j, k] * C[k,l];
+	print(D);
+}
+
+# Print the result for verification.
+# CHECK: data = 
+# CHECK-NEXT: 74.8,74.8,74.8,74.8,

lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp

@@ -178,6 +178,7 @@ void addTensorDecl(t op)
   if (ret_format.compare("Dense") == 0)
   {
     itensor = builder.create<DenseTensorDeclOp>(location, ret_value.getType(), lbls_value, ret_format);
+    builder.create<TensorFillOp>(location, itensor, builder.getF64FloatAttr(0));
   }
   else
   {

lib/Dialect/TensorAlgebra/Transforms/Passes.cpp

@@ -247,10 +247,9 @@ optimalOrder(ArrayRef<Operation *> inLTOps, Operation *outLTOp,
       // make getTotal optional to include only the operation count or
       // plus the cost  operation transpose
       totalCost += plan.getTotalTime();
-
       rhs1Labels = lhsLabels;
     }
-
+    
     // update global
     if (totalCost <= minCost)
     {
@@ -335,8 +334,10 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T
     lblMaps[op] = labelVec;
   }
 
-  auto outLabels = cast<tensorAlgebra::LabeledTensorOp>(rhsOp).getLabels();
-  LTOpsToRemove.push_back(cast<tensorAlgebra::LabeledTensorOp>(rhsOp).getOperation());
+  // auto outLabels = cast<tensorAlgebra::LabeledTensorOp>(rhsOp).getLabels();
+  // LTOpsToRemove.push_back(cast<tensorAlgebra::LabeledTensorOp>(rhsOp).getOperation());
+  auto outLabels = cast<tensorAlgebra::DenseTensorDeclOp>(rhsOp).getLabels();
+  // LTOpsToRemove.push_back(cast<tensorAlgebra::DenseTensorDeclOp>(rhsOp).getOperation());
   std::vector<Operation *> outLabelVec;
   for (auto lbl : outLabels)
   {
@@ -351,8 +352,9 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T
   std::vector<std::vector<int64_t>> lhsTensorShapes;
   std::tie(order, sumLabels, lhsTensorShapes) = optimalOrder(inLTOps, lhsOp, lblSizes, lblMaps);
 
-  bool same_order = hasSameOrder(getIdentityPermutation(order.size()), order);
+  bool same_order = hasSameOrder(getReverseIdentityPermutation(order.size()), order);
 
+  comet_debug() << "Same order " << same_order << "\n";
   // updated ta dialect generation.
   if (!same_order)
   {
@@ -517,11 +519,11 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T
                                                                MaskingAttr, nullptr);  //TODO: masking is an optional operand
       tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0));
       tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0));
-
+      comet_debug() << "New operation " << tcop << "\n";
       newRhs1 = tcop;
     }
 
-    mlir::tensorAlgebra::TensorSetOp newSetOp = builder.create<tensorAlgebra::TensorSetOp>(loc, newRhs1, operands[1].getDefiningOp()->getOperand(0));
+    mlir::tensorAlgebra::TensorSetOp newSetOp = builder.create<tensorAlgebra::TensorSetOp>(loc, newRhs1, operands[1]);
     newSetOp->setAttr("__beta__", builder.getF64FloatAttr(0.0));
 
     comet_debug() << "are they previous multop\n";

lib/Dialect/Utils/Utils.cpp

@@ -82,6 +82,16 @@ namespace mlir
       return MemRefType::get(type.getShape(), type.getElementType());
     }
 
+    std::vector<unsigned int> getReverseIdentityPermutation(size_t size)
+    {
+      std::vector<unsigned int> perm(size);
+      for (size_t i = 0; i < size; i++)
+      {
+        perm[i] = (size - 1) - i;
+      }
+      return perm;
+    }
+
     std::vector<unsigned int> getIdentityPermutation(size_t size)
     {
       std::vector<unsigned int> perm(size);