LoadStoreOp can have scalar inputs, so type check is required

NVIDIA · Nov 6, 2024 · c7e1345 · c7e1345
1 parent 967c282
commit c7e1345
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/csrc/scheduler/cache_policy_refiner.cpp b/csrc/scheduler/cache_policy_refiner.cpp
@@ -58,6 +58,12 @@ bool isLoadGlobalToLocal(const Expr* expr) {
   if (ldst->opType() != LoadStoreOpType::Set) {
     return false;
   }
+  // It should not be necessary to check the output since it should be
+  // always a TensorView as long as the input is a TensorView, but
+  // just in case.
+  if (!ldst->in()->isA<TensorView>() || !ldst->out()->isA<TensorView>()) {
+    return false;
+  }
   if (ldst->in()->as<TensorView>()->getMemoryType() != MemoryType::Global) {
     return false;
   }