From 6e002e3057287adde42b05531684659481fd032c Mon Sep 17 00:00:00 2001
From: rarbore2 <rarbore2@illinois.edu>
Date: Sat, 1 Feb 2025 17:57:02 -0600
Subject: [PATCH] Fix crash in einsum
---
.gitlab-ci.yml | 6 +++++-
hercules_ir/src/einsum.rs | 5 ++---
2 files changed, 7 insertions(+), 4 deletions(-)
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 4a3ca13d..4a2a883a 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -1,5 +1,9 @@
-test-job:
+test-cpu:
stage: test
script:
- cargo test --features=opencv
+
+test-gpu:
+ stage: test
+ script:
- cargo test --features=cuda,opencv
diff --git a/hercules_ir/src/einsum.rs b/hercules_ir/src/einsum.rs
index c2521c22..d1c276e5 100644
--- a/hercules_ir/src/einsum.rs
+++ b/hercules_ir/src/einsum.rs
@@ -235,12 +235,11 @@ impl<'a> EinsumContext<'a> {
control: _,
init: _,
reduct: _,
- } => {
- let reduce = self.so_far[&id];
+ } if let Some(reduce) = self.so_far.get(&id) => {
// Substitute opaque uses of thread ID nodes in inner expression
// with thread ID math expression, and increment inner-fork
// dimensions (alpha renaming).
- return self.substitute_new_dims(reduce);
+ return self.substitute_new_dims(*reduce);
}
_ => MathExpr::OpaqueNode(id),
};
--
GitLab