Fix breakage in unit tests

Author: Lunderberg

src/relax/transform/fuse_tir.cc

@@ -659,21 +659,25 @@ class FusedTIRConstructor : public ExprVisitor {
    */
   void MapInputBuffer(const tir::PrimFunc& func, const relax::Expr& args) {
     Array<Expr> arg_list;
-    Array<tir::Buffer> buffer_list;
     if (const auto* arg_tuple = args.as<TupleNode>()) {
       arg_list = arg_tuple->fields;
     } else {
       arg_list = {args};
     }
 
+    Array<Expr> relax_tensors;
+    Array<tir::Buffer> tir_buffers;
+
     ICHECK_GE(func->params.size(), arg_list.size());
     for (size_t i = 0; i < arg_list.size(); ++i) {
       const tir::Var& param = func->params[i];
-      const tir::Buffer& buffer = func->buffer_map.at(param);
-      buffer_list.push_back(buffer);
+      if (auto buffer = func->buffer_map.Get(param)) {
+        relax_tensors.push_back(arg_list[i]);
+        tir_buffers.push_back(buffer.value());
+      }
     }
 
-    MapArgsToBuffer(arg_list, buffer_list);
+    MapArgsToBuffer(relax_tensors, tir_buffers);
   }
 
   static Array<Integer> GetInplaceOutputIndices(const Array<Integer>& inplace_indices,

src/relax/transform/rewrite_dataflow_reshape.cc

@@ -34,12 +34,15 @@
 namespace tvm {
 namespace relax {
 
-std::vector<size_t> GetUsedArgsIndices(const tir::PrimFunc& fn, size_t num_args) {
+std::vector<size_t> GetUsedTensorArgIndices(const tir::PrimFunc& fn, size_t num_args) {
   std::vector<size_t> indices;
   for (size_t i = 0; i < num_args; ++i) {
-    auto buffer_var = fn->buffer_map[fn->params[i]]->data;
-    if (tir::UsesVar(fn->body, [=](const tir::VarNode* var) { return var == buffer_var.get(); })) {
-      indices.push_back(i);
+    if (auto buffer = fn->buffer_map.Get(fn->params[i])) {
+      auto buffer_var = buffer.value()->data;
+      if (tir::UsesVar(fn->body,
+                       [=](const tir::VarNode* var) { return var == buffer_var.get(); })) {
+        indices.push_back(i);
+      }
     }
   }
   return indices;
@@ -83,17 +86,17 @@ class DataflowReshapeRewriter : public ExprMutator {
 
     auto prim_fn = Downcast<tir::PrimFunc>(mod_->Lookup(Downcast<GlobalVar>(call->args[0])));
     auto arg_tuple = Downcast<Tuple>(call->args[1])->fields;
-    auto used_arg_indices = GetUsedArgsIndices(prim_fn, arg_tuple.size());
+    auto used_tensor_arg_indices = GetUsedTensorArgIndices(prim_fn, arg_tuple.size());
 
     // The number of inputs to call_tir(reshape, (...)) might not be one, since FuseOps
     // can generate a fused TupleGetItem + reshape function whose input is a tuple. FuseTIR
     // then flattens the tuple input so that the fused TIR reshape function ends up having
     // multiple input buffers. But only one of them should be accessed and reshaped.
-    if (used_arg_indices.size() != 1) {
+    if (used_tensor_arg_indices.size() != 1) {
       return GetRef<Call>(call);
     }
 
-    auto arg = arg_tuple[used_arg_indices[0]];
+    auto arg = arg_tuple[used_tensor_arg_indices[0]];
 
     if (!IsCallingTIRReshape(call, arg)) {
       return GetRef<Call>(call);