Slice python api

csrc/python_frontend/fusion_cache.cpp

@@ -247,7 +247,7 @@ bool TrieNode::isTerminal() const {
   return (record.get()->recordType() == serde::RecordType::End);
 }
 
-void TrieNode::markException(std::exception e) {
+void TrieNode::setException(const char* e) {
   std::lock_guard<std::mutex> guard(trie_node_lock);
   exception = e;
 }

csrc/python_frontend/fusion_cache.h

@@ -105,7 +105,7 @@ struct TrieNode {
   // a the end of Fusion entry in the cache.
   bool isTerminal() const;
   std::optional<std::exception> getException();
-  void markException(std::exception e);
+  void setException(const char* e);
   //! Serialize TrieNode using flatbuffers
   NVF_API flatbuffers::Offset<serde::TrieNode> serialize(
       flatbuffers::FlatBufferBuilder& builder,
@@ -127,7 +127,7 @@ struct TrieNode {
   TrieNode* parent;
   //! For thread-Safe locking of a node
   std::mutex trie_node_lock;
-  std::optional<std::exception> exception;
+  std::optional<std::string> exception = std::nullopt;
 };
 
 //! \class FusionCache

csrc/python_frontend/fusion_definition.cpp

@@ -96,9 +96,8 @@ void FusionDefinition::finalizeDefinition() {
 
       buildFusionIr(preschedFusion());
     } catch (const std::exception& e) {
-      trie_node_->markException(e);
+      trie_node_->setException(e.what());
       fusion_id_ = std::nullopt;
-      std::cout << "error: " << e.what() << std::endl;
       throw;
     }
 

csrc/python_frontend/fusion_record.h

@@ -395,19 +395,21 @@ struct SliceOpRecord : RecordFunctor {
       Val* start_idx = start.at(idx);
       Val* end_idx = end.at(idx);
       Val* stride_idx = stride.at(idx);
-          NVF_CHECK(
-              !start_idx->isConstInt() || start_idx->evaluate().as<int64_t>() >= 0,
-              "Slice operation start_indices must be greater-than-or-equal-to 0. Start Indices: ",
-              start_idx->evaluate().as<int64_t>());
-          NVF_CHECK(
-              !start_idx->isConstInt() || !end_idx->isConstInt() || end_idx->evaluate().as<int64_t>() >= start_idx->evaluate().as<int64_t>(),
-              "Slice operation end_indices must be greater-than-or-equal-to start_indices. Start Indices: ",
-              start_idx->evaluate().as<int64_t>(),
-              " End Indices: ",
-              end_idx->evaluate().as<int64_t>());
-          NVF_CHECK(
-              stride_idx->isConstInt() && stride_idx->evaluate().as<int64_t>() == 1,
-              "nvFuser Limitation: All slice operation strides must be of const size 1");
+      NVF_CHECK(
+          !start_idx->isConstInt() || start_idx->evaluate().as<int64_t>() >= 0,
+          "Slice operation start_indices must be greater-than-or-equal-to 0. Start Indices: ",
+          start_idx->evaluate().as<int64_t>());
+      NVF_CHECK(
+          !start_idx->isConstInt() || !end_idx->isConstInt() ||
+              end_idx->evaluate().as<int64_t>() >=
+                  start_idx->evaluate().as<int64_t>(),
+          "Slice operation end_indices must be greater-than-or-equal-to start_indices. Start Indices: ",
+          start_idx->evaluate().as<int64_t>(),
+          " End Indices: ",
+          end_idx->evaluate().as<int64_t>());
+      NVF_CHECK(
+          stride_idx->isConstInt() && stride_idx->evaluate().as<int64_t>() == 1,
+          "nvFuser Limitation: All slice operation strides must be of const size 1");
       vec_slice.push_back({start_idx, end_idx, stride_idx});
     }
     auto output = slice(arg, vec_slice);

csrc/python_frontend/python_bindings.cpp

@@ -100,11 +100,15 @@ Vector define_vector_explicit_fn(
     FusionDefinition& self,
     ITERABLE& values,
     PrimDataType dtype = DataType::Int) {
-  return define_vector_fn<ITERABLE>(self, values, /*inline_def=*/false, /*shape_check=*/true);
+  return define_vector_fn<ITERABLE>(
+      self, values, /*inline_def=*/false, /*shape_check=*/true);
 }
 
 template <class ShapeType>
-Vector SequenceAsVector(ShapeType shape, FusionDefinition& fd, bool shape_check=true) {
+Vector SequenceAsVector(
+    ShapeType shape,
+    FusionDefinition& fd,
+    bool shape_check = true) {
   static_assert(
       std::is_same_v<ShapeType, Vector> ||
       std::is_same_v<ShapeType, py::list> ||
@@ -122,7 +126,8 @@ Vector SequenceAsVector(ShapeType shape, FusionDefinition& fd, bool shape_check=
     // ```
     // would not work because the compiler would try to instantiate
     // define_vector_fn<Vector> and fail.
-    return define_vector_fn<ShapeType>(fd, shape, /*inline_def=*/true, /*shape_check=*/shape_check);
+    return define_vector_fn<ShapeType>(
+        fd, shape, /*inline_def=*/true, /*shape_check=*/shape_check);
   }
 }
 
@@ -251,7 +256,8 @@ Tensor slice_fn(
   size_t stride_index = 0;
 
   if (strides.has_value()) {
-    Vector new_stride = SequenceAsVector(strides.value(), *fd, /*shape_check=*/false);
+    Vector new_stride =
+        SequenceAsVector(strides.value(), *fd, /*shape_check=*/false);
     NVF_CHECK(
         new_start.size == new_stride.size,
         "Slice start_indices and strides don't match! Start Indices: ",