Began adding allBlocks attribute to index tree #20

pflynn157 · pthomadakis · commit 39f1a013e91f · 2024-09-20T18:57:07.000-07:00
diff --git a/first.mlir b/first.mlir
diff --git a/first.ta b/first.ta
@@ -1,11 +1,19 @@
 def main() {
 	#IndexLabel Declarations
 	IndexLabel [a] = [?];
-	IndexLabel [b] = [?];           
-
-	#Tensor Declarations
-	Tensor<double> A([a, b], {BCSR});
+	IndexLabel [b] = [?];   
+    
+	Tensor<double> A([a, b], {CSR});
+	Tensor<double> B([b, a], {Dense});
+	Tensor<double> C([b, a], {Dense});
+	
     A[a, b] = comet_read(0);
+    B[b, a] = 1.0;
+    C[b, a] = A[a, b] * B[b, a];
+    #C[b, a] = 1.0;
+    
 	print(A);
+	print(B);
+	print(C);
 }
 
diff --git a/include/comet/Dialect/IndexTree/IR/IndexTree.h b/include/comet/Dialect/IndexTree/IR/IndexTree.h
@@ -214,7 +214,7 @@ class Index_Tree
 
 public:
   IndicesType getIndices(std::vector<mlir::Value> &lbls);
-  Tensor *getOrCreateTensor(mlir::Value v, std::vector<mlir::Value> &allIndexLabels, FormatsType &formats);
+  Tensor *getOrCreateTensor(mlir::Value v, std::vector<mlir::Value> &allIndexLabels, FormatsType &formats, BlocksType &blocks);
 
   vector<TreeNode *> getNodes();
 
diff --git a/include/comet/Dialect/IndexTree/IR/IndexTreeOps.td b/include/comet/Dialect/IndexTree/IR/IndexTreeOps.td
@@ -64,15 +64,15 @@ def IndexTreeComputeLHSOp : IndexTree_Op<"ComputeLHS", [Pure]>{
   let summary = "";
   let description = [{}];
 
-  let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats);
+  let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats, ArrayAttr:$allBlocks);
   let results = (outs AnyType:$output);
 }
 
 def IndexTreeComputeRHSOp : IndexTree_Op<"ComputeRHS", [Pure]>{
   let summary = "";
   let description = [{}];
 
-  let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats);
+  let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats, ArrayAttr:$allBlocks);
   let results = (outs AnyType:$output);
 }
 
@@ -118,4 +118,4 @@ def IndexTreeOp : IndexTree_Op<"itree", [Pure]>{
   //let hasVerifier = 1; 
 }
 
-#endif // INDEXTREE_OPS
+#endif // INDEXTREE_OPS
diff --git a/include/comet/Dialect/IndexTree/Transforms/Tensor.h b/include/comet/Dialect/IndexTree/Transforms/Tensor.h
@@ -35,6 +35,7 @@
 
 typedef std::vector<unsigned int> IndicesType;
 typedef std::vector<std::string> FormatsType;
+typedef std::vector<std::string> BlocksType;
 typedef std::vector<IterDomain *> DomainsType;
 
 using std::make_shared;
@@ -53,6 +54,7 @@ class Tensor
   string name;
   IndicesType indices;
   FormatsType format;
+  BlocksType block;
   IndicesType hidden;
   DomainsType domains;
   UnitExpression *definingExpr = nullptr;
@@ -61,12 +63,13 @@ class Tensor
 public:
   static int count;
 
-  Tensor(mlir::Value &value, IndicesType &indices, vector<string> &format)
+  Tensor(mlir::Value &value, IndicesType &indices, vector<string> &format, vector<string> &block)
   {
     // assert(format.size() == indices.size());
     this->value = value;
     this->indices = indices;
     this->format = format;
+    this->block = block;
     id = count++;
     domains = std::vector<IterDomain *>(indices.size());
   }
@@ -154,6 +157,21 @@ class Tensor
   {
     Tensor::format = format;
   }
+  
+  const string &getBlock(int i) const
+  {
+    return block.at(i);
+  }
+
+  BlocksType &getBlocks()
+  {
+    return block;
+  }
+
+  void setBlock(const vector<string> &block)
+  {
+    Tensor::block = block;
+  }
 
   void setHiddenIndices(IndicesType &hiddenIndices)
   {
diff --git a/include/comet/Dialect/Utils/Utils.h b/include/comet/Dialect/Utils/Utils.h
@@ -126,6 +126,7 @@ namespace mlir
     std::vector<unsigned> getSumIndices(std::vector<unsigned> rhs_perm, std::vector<unsigned> rhs_perm_free);
     std::vector<unsigned> getIndexIterateOrder(std::vector<unsigned> rhs1_perm, std::vector<unsigned> rhs2_perm);
     std::vector<std::vector<std::string>> getAllFormats(ArrayAttr opFormatsArrayAttr, std::vector<std::vector<int64_t>> allPerms);
+    std::vector<std::vector<std::string>> getAllBlocks(ArrayAttr opFormatsArrayAttr, std::vector<std::vector<int64_t>> allPerms);
     bool checkIsElementwise(std::vector<std::vector<int>> allPerms);
     bool checkIsMixedMode(std::vector<std::vector<std::string>> formats);
     bool checkIsDense(std::vector<std::string> format);
diff --git a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp
@@ -264,24 +264,28 @@ void doTensorMultOp(TensorMultOp op, unique_ptr<Index_Tree> &tree, TargetDevice
 // comet_debug() << allPerms;
 #endif
 
+
+  auto allBlocks = getAllBlocks(op.getFormatsAttr(), allPerms);
   auto allFormats = getAllFormats(op.getFormatsAttr(), allPerms);
   auto SemiringOp = op.getSemiringAttr();
   auto MaskingTypeAttr = op.getMaskTypeAttr();
 
   /// If the operation is one of the chosen operations, then record output indices as parallel interators.
   bool is_chosen_operations = check_chosen_operations(allPerms, allFormats);
 
-  auto B = tree->getOrCreateTensor(rhs1_tensor, rhs1_labels, allFormats[0]);
-  auto C = tree->getOrCreateTensor(rhs2_tensor, rhs2_labels, allFormats[1]);
-  auto A = tree->getOrCreateTensor(lhs_tensor, lhs_labels, allFormats[2]);
+
+
+  auto B = tree->getOrCreateTensor(rhs1_tensor, rhs1_labels, allFormats[0], allBlocks[0]);
+  auto C = tree->getOrCreateTensor(rhs2_tensor, rhs2_labels, allFormats[1], allBlocks[1]);
+  auto A = tree->getOrCreateTensor(lhs_tensor, lhs_labels, allFormats[2], allBlocks[2]);
 
   Tensor *M;
   std::unique_ptr<UnitExpression> e;
   std::vector<mlir::Value> empty;
   if (mask_tensor != nullptr) /// mask is an optional input
   {
     comet_debug() << "mask input provided by user\n";
-    M = tree->getOrCreateTensor(mask_tensor, empty, allFormats[2]); /// We don't need indexlabel info for the mask
+    M = tree->getOrCreateTensor(mask_tensor, empty, allFormats[2], allBlocks[2]); /// We don't need indexlabel info for the mask
     e = make_unique<UnitExpression>(A, B, C, M, "*");
   }
   else
@@ -377,14 +381,15 @@ void doElementWiseOp(T op, unique_ptr<Index_Tree> &tree)
 
   auto allPerms = getAllPerms(op.getIndexingMaps());
   auto allFormats = getAllFormats(op.getFormatsAttr(), allPerms);
+  auto allBlocks = getAllBlocks(op.getFormatsAttr(), allPerms);
   auto SemiringOp = op.getSemiringAttr();
   auto maskAttr = "none";
 
   assert(allPerms.size() == 3);
 
-  auto B = tree->getOrCreateTensor(rhs1_tensor, rhs1_labels, allFormats[0]);
-  auto C = tree->getOrCreateTensor(rhs2_tensor, rhs2_labels, allFormats[1]);
-  auto A = tree->getOrCreateTensor(lhs_tensor, lhs_labels, allFormats[2]);
+  auto B = tree->getOrCreateTensor(rhs1_tensor, rhs1_labels, allFormats[0], allBlocks[0]);
+  auto C = tree->getOrCreateTensor(rhs2_tensor, rhs2_labels, allFormats[1], allBlocks[1]);
+  auto A = tree->getOrCreateTensor(lhs_tensor, lhs_labels, allFormats[2], allBlocks[2]);
 
   auto e = make_unique<UnitExpression>(A, B, C, "*");
 
@@ -498,6 +503,27 @@ IndexTreeComputeOp createComputeNodeOp(OpBuilder &builder, TreeNode *node, Locat
     }
     allFormats_lhs.push_back(builder.getStrArrayAttr(formats));
   }
+  
+  SmallVector<Attribute, 8> allBlocks_rhs;
+  for (auto t : expr->getOperands())
+  {
+    SmallVector<StringRef, 8> blocks;
+    for (auto &b : t->getBlocks())
+    {
+      blocks.push_back(b);
+    }
+    allBlocks_rhs.push_back(builder.getStrArrayAttr(blocks));
+  }
+  SmallVector<Attribute, 8> allBlocks_lhs;
+  for (auto t : expr->getResults())
+  {
+    SmallVector<StringRef, 8> blocks;
+    for (auto &b : t->getBlocks())
+    {
+      blocks.push_back(b);
+    }
+    allBlocks_lhs.push_back(builder.getStrArrayAttr(blocks));
+  }
 
   std::vector<Value> t_rhs;
   Value t_lhs = expr->getLHS()->getValue();
@@ -516,12 +542,14 @@ IndexTreeComputeOp createComputeNodeOp(OpBuilder &builder, TreeNode *node, Locat
   Value leafop_rhs = builder.create<indexTree::IndexTreeComputeRHSOp>(loc,
                                                                       mlir::UnrankedTensorType::get(builder.getF64Type()), t_rhs,
                                                                       builder.getArrayAttr(allIndices_rhs),
-                                                                      builder.getArrayAttr(allFormats_rhs));
+                                                                      builder.getArrayAttr(allFormats_rhs),
+                                                                      builder.getArrayAttr(allBlocks_rhs));
   comet_vdump(leafop_rhs);
   Value leafop_lhs = builder.create<indexTree::IndexTreeComputeLHSOp>(loc,
                                                                       mlir::UnrankedTensorType::get(builder.getF64Type()), t_lhs,
                                                                       builder.getArrayAttr(allIndices_lhs),
-                                                                      builder.getArrayAttr(allFormats_lhs));
+                                                                      builder.getArrayAttr(allFormats_lhs),
+                                                                      builder.getArrayAttr(allBlocks_lhs));
   comet_vdump(leafop_lhs);
 
   bool comp_worksp_opt = false; /// non-compressed workspace, this is a place-holder and it is updated in workspace transform pass.
diff --git a/lib/Conversion/TensorAlgebraToSCF/LowerPCToLoops.cpp b/lib/Conversion/TensorAlgebraToSCF/LowerPCToLoops.cpp
@@ -140,23 +140,25 @@ void PCToLoopsLoweringPass::replicateOpsForLoopBody(Location loc, OpBuilder &bui
   {
     indexTree::IndexTreeComputeRHSOp it_compute_rhs_op = llvm::dyn_cast<indexTree::IndexTreeComputeRHSOp>(op);
     ArrayAttr op_formats_ArrayAttr = it_compute_rhs_op.getAllFormats();
+    ArrayAttr op_blocks_ArrayAttr = it_compute_rhs_op.getAllBlocks();
     ArrayAttr op_perms_ArrayAttr = it_compute_rhs_op.getAllPerms();
 
     rhs = builder.create<indexTree::IndexTreeComputeRHSOp>(loc, mlir::UnrankedTensorType::get(builder.getF64Type()),
                                                            it_compute_rhs_op->getOperands(), // tensors
-                                                           op_perms_ArrayAttr, op_formats_ArrayAttr);
+                                                           op_perms_ArrayAttr, op_formats_ArrayAttr, op_blocks_ArrayAttr);
   }
 
   // create IndexTreeComputeLHSOp, no dependency to earlier replications
   if (isa<indexTree::IndexTreeComputeLHSOp>(op))
   {
     indexTree::IndexTreeComputeLHSOp it_compute_lhs_op = llvm::dyn_cast<indexTree::IndexTreeComputeLHSOp>(op);
     ArrayAttr op_formats_ArrayAttr = it_compute_lhs_op.getAllFormats();
+    ArrayAttr op_blocks_ArrayAttr = it_compute_lhs_op.getAllBlocks();
     ArrayAttr op_perms_ArrayAttr = it_compute_lhs_op.getAllPerms();
 
     lhs = builder.create<indexTree::IndexTreeComputeLHSOp>(loc, mlir::UnrankedTensorType::get(builder.getF64Type()),
                                                            it_compute_lhs_op->getOperands(), /// tensors
-                                                           op_perms_ArrayAttr, op_formats_ArrayAttr);
+                                                           op_perms_ArrayAttr, op_formats_ArrayAttr, op_blocks_ArrayAttr);
   }
 
   /// create IndexTreeComputeOp only if rhs and lhs are ready
@@ -338,4 +340,4 @@ void PCToLoopsLoweringPass::runOnOperation()
 std::unique_ptr<Pass> mlir::comet::createPCToLoopsLoweringPass()
 {
   return std::make_unique<PCToLoopsLoweringPass>();
-}
+}
diff --git a/lib/Dialect/IndexTree/IR/IndexTree.cpp b/lib/Dialect/IndexTree/IR/IndexTree.cpp
@@ -117,12 +117,12 @@ IndicesType Index_Tree::getIndices(std::vector<mlir::Value> &lbls)
   return indices;
 }
 
-Tensor *Index_Tree::getOrCreateTensor(mlir::Value v, std::vector<mlir::Value> &allIndexLabels, FormatsType &formats)
+Tensor *Index_Tree::getOrCreateTensor(mlir::Value v, std::vector<mlir::Value> &allIndexLabels, FormatsType &formats, BlocksType &blocks)
 {
   IndicesType indices = getIndices(allIndexLabels);
   comet_debug() << "Num Indices: " << indices.size() << ", Num formats " << formats.size() << "\n";
 
-  return new Tensor(v, indices, formats);
+  return new Tensor(v, indices, formats, blocks);
 }
 
 TreeNode *Index_Tree::addComputeNode(unique_ptr<UnitExpression> expr, TreeNode *parent)
@@ -261,4 +261,4 @@ void IteratorType::setType(std::string t) {
   } else {
     llvm::errs() << "Unsupported iterator type " + t + "\n";
   }
-}
+}
diff --git a/lib/Dialect/IndexTree/Transforms/Fusion.cpp b/lib/Dialect/IndexTree/Transforms/Fusion.cpp
@@ -440,8 +440,10 @@ mlir::Value IndexTreeKernelFusionPass::createReducedComputeLHS(
   mlir::indexTree::IndexTreeComputeLHSOp it_compute_lhs_op = llvm::dyn_cast<mlir::indexTree::IndexTreeComputeLHSOp>(
       lhs_op);
   ArrayAttr op_formats_ArrayAttr = it_compute_lhs_op.getAllFormats();
+  ArrayAttr op_blocks_ArrayAttr = it_compute_lhs_op.getAllBlocks();
   ArrayAttr op_perms_ArrayAttr = it_compute_lhs_op.getAllPerms();
   std::vector<std::vector<std::string>> old_formats_strs = convertArrayAttrStrTo2DVector(op_formats_ArrayAttr);
+  std::vector<std::vector<std::string>> old_blocks_strs = convertArrayAttrStrTo2DVector(op_blocks_ArrayAttr);
   std::vector<std::vector<int>> old_perms_ints = convertArrayAttrIntTo2DVector(op_perms_ArrayAttr);
 
   /// Create the new formats
@@ -450,6 +452,13 @@ mlir::Value IndexTreeKernelFusionPass::createReducedComputeLHS(
   SmallVector<StringRef, 8> formats;
   formats.insert(formats.end(), old_formats_strs[0].begin() + rank_base, old_formats_strs[0].end());
   new_formats.push_back(builder.getStrArrayAttr(formats));
+  
+  /// Create the new blocks
+  /// i.g., convert [["D", "D"]] to [["D"]]
+  SmallVector<Attribute, 8> new_blocks;
+  SmallVector<StringRef, 8> blocks;
+  blocks.insert(blocks.end(), old_blocks_strs[0].begin() + rank_base, old_blocks_strs[0].end());
+  new_blocks.push_back(builder.getStrArrayAttr(blocks));
 
   /// Create the new perms
   /// i.g., convert [[1, 0]] to [[0]]
@@ -466,7 +475,8 @@ mlir::Value IndexTreeKernelFusionPass::createReducedComputeLHS(
       mlir::UnrankedTensorType::get(builder.getF64Type()),
       tensors,
       builder.getArrayAttr(new_perms),
-      builder.getArrayAttr(new_formats));
+      builder.getArrayAttr(new_formats),
+      builder.getArrayAttr(new_blocks));
 
   return new_lhs_op;
 }
@@ -485,8 +495,10 @@ mlir::Value IndexTreeKernelFusionPass::createReducedComputeRHS(
   mlir::indexTree::IndexTreeComputeRHSOp it_compute_rhs_op = llvm::dyn_cast<mlir::indexTree::IndexTreeComputeRHSOp>(
       rhs_op);
   ArrayAttr op_formats_ArrayAttr = it_compute_rhs_op.getAllFormats();
+  ArrayAttr op_blocks_ArrayAttr = it_compute_rhs_op.getAllBlocks();
   ArrayAttr op_perms_ArrayAttr = it_compute_rhs_op.getAllPerms();
   std::vector<std::vector<std::string>> old_formats_strs = convertArrayAttrStrTo2DVector(op_formats_ArrayAttr);
+  std::vector<std::vector<std::string>> old_blocks_strs = convertArrayAttrStrTo2DVector(op_blocks_ArrayAttr);
   std::vector<std::vector<int>> old_perms_ints = convertArrayAttrIntTo2DVector(op_perms_ArrayAttr);
 
   /// Locate the operand to be reduced
@@ -516,6 +528,23 @@ mlir::Value IndexTreeKernelFusionPass::createReducedComputeRHS(
     }
     new_formats.push_back(builder.getStrArrayAttr(formats));
   }
+  
+  /// Create the new blocks
+  /// Basically same algorithm as the formats
+  SmallVector<Attribute, 8> new_blocks;
+  for (uint32_t b_i = 0; b_i < old_blocks_strs.size(); ++b_i)
+  {
+    SmallVector<StringRef, 8> blocks;
+    if (b_i == tensor_id)
+    { /// for the new reduced tensor
+      blocks.insert(blocks.end(), old_formats_strs[b_i].begin() + rank_base, old_blocks_strs[b_i].end());
+    }
+    else
+    { /// for other remaining old operands
+      blocks.insert(blocks.end(), old_blocks_strs[b_i].begin(), old_blocks_strs[b_i].end());
+    }
+    new_blocks.push_back(builder.getStrArrayAttr(blocks));
+  }
 
   /// Create the new perms
   /// i.g., convert [[1, 0], [0, 2]] to [[0], [0, 2]]
@@ -554,7 +583,8 @@ mlir::Value IndexTreeKernelFusionPass::createReducedComputeRHS(
       mlir::UnrankedTensorType::get(builder.getF64Type()),
       tensors,
       builder.getArrayAttr(new_perms),
-      builder.getArrayAttr(new_formats));
+      builder.getArrayAttr(new_formats),
+      builder.getArrayAttr(new_blocks));
 
   return new_rhs_op;
 }
@@ -729,14 +759,19 @@ mlir::Value IndexTreeKernelFusionPass::createResetComputeRHS(
   SmallVector<Attribute, 1> formats_rhs;
   SmallVector<StringRef, 1> empty_format;
   formats_rhs.push_back(builder.getStrArrayAttr(empty_format));
+  
+  SmallVector<Attribute, 1> blocks_rhs;
+  SmallVector<StringRef, 1> empty_block;
+  blocks_rhs.push_back(builder.getStrArrayAttr(empty_block));
 
   /// TODO(zpeng): What if the type is not F64?
   mlir::Value compute_rhs = builder.create<indexTree::IndexTreeComputeRHSOp>(
       loc,
       mlir::UnrankedTensorType::get(builder.getF64Type()),
       tensors_rhs,
       builder.getArrayAttr(indices_rhs),
-      builder.getArrayAttr(formats_rhs));
+      builder.getArrayAttr(formats_rhs),
+      builder.getArrayAttr(blocks_rhs));
 
   return compute_rhs;
 }
@@ -770,13 +805,19 @@ mlir::Value IndexTreeKernelFusionPass::createResetComputeLHS(
   SmallVector<Attribute, 1> formats_lhs;
   SmallVector<StringRef, 1> one_format(rank, "D");
   formats_lhs.push_back(builder.getStrArrayAttr(one_format));
+  
+  /// Get blocks [["UNK"]]
+  SmallVector<Attribute, 1> blocks_lhs;
+  SmallVector<StringRef, 1> one_block(rank, "UNK");
+  blocks_lhs.push_back(builder.getStrArrayAttr(one_block));
 
   mlir::Value compute_lhs = builder.create<indexTree::IndexTreeComputeLHSOp>(
       loc,
       mlir::UnrankedTensorType::get(builder.getF64Type()),
       tensors_lhs,
       builder.getArrayAttr(indices_lhs),
-      builder.getArrayAttr(formats_lhs));
+      builder.getArrayAttr(formats_lhs),
+      builder.getArrayAttr(blocks_lhs));
 
   return compute_lhs;
 }
diff --git a/lib/Dialect/IndexTree/Transforms/WorkspaceTransforms.cpp b/lib/Dialect/IndexTree/Transforms/WorkspaceTransforms.cpp
diff --git a/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp b/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp
diff --git a/lib/Dialect/Utils/Utils.cpp b/lib/Dialect/Utils/Utils.cpp

Original file line number	Diff line number	Diff line change
`@@ -64,15 +64,15 @@ def IndexTreeComputeLHSOp : IndexTree_Op<"ComputeLHS", [Pure]>{`
`64`	`64`	`let summary = "";`
`65`	`65`	`let description = [{}];`
`66`	`66`
`67`		`- let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats);`
	`67`	`+ let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats, ArrayAttr:$allBlocks);`
`68`	`68`	`let results = (outs AnyType:$output);`
`69`	`69`	`}`
`70`	`70`
`71`	`71`	`def IndexTreeComputeRHSOp : IndexTree_Op<"ComputeRHS", [Pure]>{`
`72`	`72`	`let summary = "";`
`73`	`73`	`let description = [{}];`
`74`	`74`
`75`		`- let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats);`
	`75`	`+ let arguments = (ins Variadic<AnyType>:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats, ArrayAttr:$allBlocks);`
`76`	`76`	`let results = (outs AnyType:$output);`
`77`	`77`	`}`
`78`	`78`
`@@ -118,4 +118,4 @@ def IndexTreeOp : IndexTree_Op<"itree", [Pure]>{`
`118`	`118`	`//let hasVerifier = 1;`
`119`	`119`	`}`
`120`	`120`
`121`		`-#endif // INDEXTREE_OPS`
	`121`	`+#endif // INDEXTREE_OPS`
Original file line number	Diff line number	Diff line change
`@@ -140,23 +140,25 @@ void PCToLoopsLoweringPass::replicateOpsForLoopBody(Location loc, OpBuilder &bui`
`140`	`140`	`{`
`141`	`141`	`indexTree::IndexTreeComputeRHSOp it_compute_rhs_op = llvm::dyn_cast<indexTree::IndexTreeComputeRHSOp>(op);`
`142`	`142`	`ArrayAttr op_formats_ArrayAttr = it_compute_rhs_op.getAllFormats();`
	`143`	`+ ArrayAttr op_blocks_ArrayAttr = it_compute_rhs_op.getAllBlocks();`
`143`	`144`	`ArrayAttr op_perms_ArrayAttr = it_compute_rhs_op.getAllPerms();`
`144`	`145`
`145`	`146`	`rhs = builder.create<indexTree::IndexTreeComputeRHSOp>(loc, mlir::UnrankedTensorType::get(builder.getF64Type()),`
`146`	`147`	`it_compute_rhs_op->getOperands(), // tensors`
`147`		`- op_perms_ArrayAttr, op_formats_ArrayAttr);`
	`148`	`+ op_perms_ArrayAttr, op_formats_ArrayAttr, op_blocks_ArrayAttr);`
`148`	`149`	`}`
`149`	`150`
`150`	`151`	`// create IndexTreeComputeLHSOp, no dependency to earlier replications`
`151`	`152`	`if (isa<indexTree::IndexTreeComputeLHSOp>(op))`
`152`	`153`	`{`
`153`	`154`	`indexTree::IndexTreeComputeLHSOp it_compute_lhs_op = llvm::dyn_cast<indexTree::IndexTreeComputeLHSOp>(op);`
`154`	`155`	`ArrayAttr op_formats_ArrayAttr = it_compute_lhs_op.getAllFormats();`
	`156`	`+ ArrayAttr op_blocks_ArrayAttr = it_compute_lhs_op.getAllBlocks();`
`155`	`157`	`ArrayAttr op_perms_ArrayAttr = it_compute_lhs_op.getAllPerms();`
`156`	`158`
`157`	`159`	`lhs = builder.create<indexTree::IndexTreeComputeLHSOp>(loc, mlir::UnrankedTensorType::get(builder.getF64Type()),`
`158`	`160`	`it_compute_lhs_op->getOperands(), /// tensors`
`159`		`- op_perms_ArrayAttr, op_formats_ArrayAttr);`
	`161`	`+ op_perms_ArrayAttr, op_formats_ArrayAttr, op_blocks_ArrayAttr);`
`160`	`162`	`}`
`161`	`163`
`162`	`164`	`/// create IndexTreeComputeOp only if rhs and lhs are ready`
`@@ -338,4 +340,4 @@ void PCToLoopsLoweringPass::runOnOperation()`
`338`	`340`	`std::unique_ptr<Pass> mlir::comet::createPCToLoopsLoweringPass()`
`339`	`341`	`{`
`340`	`342`	`return std::make_unique<PCToLoopsLoweringPass>();`
`341`		`-}`
	`343`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -117,12 +117,12 @@ IndicesType Index_Tree::getIndices(std::vector<mlir::Value> &lbls)`
`117`	`117`	`return indices;`
`118`	`118`	`}`
`119`	`119`
`120`		`-Tensor *Index_Tree::getOrCreateTensor(mlir::Value v, std::vector<mlir::Value> &allIndexLabels, FormatsType &formats)`
	`120`	`+Tensor *Index_Tree::getOrCreateTensor(mlir::Value v, std::vector<mlir::Value> &allIndexLabels, FormatsType &formats, BlocksType &blocks)`
`121`	`121`	`{`
`122`	`122`	`IndicesType indices = getIndices(allIndexLabels);`
`123`	`123`	`comet_debug() << "Num Indices: " << indices.size() << ", Num formats " << formats.size() << "\n";`
`124`	`124`
`125`		`- return new Tensor(v, indices, formats);`
	`125`	`+ return new Tensor(v, indices, formats, blocks);`
`126`	`126`	`}`
`127`	`127`
`128`	`128`	`TreeNode Index_Tree::addComputeNode(unique_ptr<UnitExpression> expr, TreeNode parent)`
`@@ -261,4 +261,4 @@ void IteratorType::setType(std::string t) {`
`261`	`261`	`} else {`
`262`	`262`	`llvm::errs() << "Unsupported iterator type " + t + "\n";`
`263`	`263`	`}`
`264`		`-}`
	`264`	`+}`