Skip to content

src/execution/ast_hash.cc

Namespaces

Name
tenzr

Source code

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "tenzr/execution/ast_hash.h"

#include <unordered_map>

namespace tenzr {

  namespace {

    constexpr uint64_t kFnvOffset = 14695981039346656037ULL;
    constexpr uint64_t kFnvPrime = 1099511628211ULL;

    inline uint64_t Fnv1aCombine(uint64_t hash, uint64_t val) {
      hash ^= val;
      hash *= kFnvPrime;
      return hash;
    }

    uint64_t HashNodeAttributes(const std::shared_ptr<Node>& node, uint64_t h) {
      auto combine = [&](uint64_t v) { return Fnv1aCombine(h, v); };
      auto combine_ints = [&](const std::vector<int64_t>& vals) {
        for (int64_t v : vals) {
          h = combine(static_cast<uint64_t>(v));
        }
      };
      auto combine_str = [&](const std::string& s) {
        for (char c : s) {
          h = combine(static_cast<uint64_t>(c));
        }
      };

      switch (node->Kind()) {
        case NodeKind::kTranspose:
          combine_ints(std::static_pointer_cast<TransposeNode>(node)->Permutation());
          break;
        case NodeKind::kSlice:
          for (const auto& sl : std::static_pointer_cast<SliceNode>(node)->Slices()) {
            h = combine(sl.start.has_value() ? static_cast<uint64_t>(*sl.start) : 0xDEADBEEFULL);
            h = combine(sl.stop.has_value() ? static_cast<uint64_t>(*sl.stop) : 0xBAADF00DULL);
            h = combine(static_cast<uint64_t>(sl.step));
          }
          break;
        case NodeKind::kReduce: {
          auto r = std::static_pointer_cast<ReduceNode>(node);
          h = combine(static_cast<uint64_t>(r->Reduction()));
          h = combine(r->Keepdims() ? 1ULL : 0ULL);
          combine_ints(r->CanonicalDimensions());
          break;
        }
        case NodeKind::kReduceWindow: {
          auto rw = std::static_pointer_cast<ReduceWindowNode>(node);
          combine_ints(rw->WindowDimensions());
          combine_ints(rw->WindowStrides());
          combine_ints(rw->PaddingLow());
          combine_ints(rw->PaddingHigh());
          combine_ints(rw->WindowDilations());
          combine_str(rw->Reduction());
          break;
        }
        case NodeKind::kConvolution: {
          auto c = std::static_pointer_cast<ConvolutionNode>(node);
          combine_ints(c->WindowStrides());
          combine_ints(c->LhsDilation());
          combine_ints(c->RhsDilation());
          combine_ints(c->PaddingLow());
          combine_ints(c->PaddingHigh());
          break;
        }
        case NodeKind::kDynamicSlice:
          combine_ints(std::static_pointer_cast<DynamicSliceNode>(node)->SliceSizes());
          break;
        case NodeKind::kPad: {
          auto p = std::static_pointer_cast<PadNode>(node);
          combine_ints(p->EdgePaddingLow());
          combine_ints(p->EdgePaddingHigh());
          combine_ints(p->InteriorPadding());
          break;
        }
        case NodeKind::kScatter:
          combine_str(std::static_pointer_cast<ScatterNode>(node)->UpdateOp());
          break;
        case NodeKind::kConcatenate:
          h = combine(static_cast<uint64_t>(std::static_pointer_cast<ConcatenateNode>(node)->Axis()));
          break;
        case NodeKind::kArgSort: {
          auto as = std::static_pointer_cast<ArgSortNode>(node);
          h = combine(static_cast<uint64_t>(as->Axis()));
          h = combine(as->Descending() ? 1ULL : 0ULL);
          h = combine(as->ReturnValues() ? 1ULL : 0ULL);
          break;
        }
        default:
          break;
      }
      return h;
    }

    uint64_t hash_node_internal(const std::shared_ptr<Node>& node, std::unordered_map<uint64_t, uint64_t>& memo,
                                std::unordered_map<uint64_t, size_t>& var_indices) {
      if (!node) {
        return 0;
      }

      auto it = memo.find(node->Id());
      if (it != memo.end()) {
        return it->second;
      }

      uint64_t h = kFnvOffset;
      h = Fnv1aCombine(h, static_cast<uint64_t>(node->Kind()));
      h = Fnv1aCombine(h, static_cast<uint64_t>(node->Dtype()));

      for (char c : node->OpName()) {
        h = Fnv1aCombine(h, static_cast<uint64_t>(c));
      }

      for (int64_t d : node->shape().Dims()) {
        h = Fnv1aCombine(h, static_cast<uint64_t>(d));
      }

      if (node->Kind() == NodeKind::kVariable) {
        auto [v_it, inserted] = var_indices.try_emplace(node->Id(), var_indices.size());
        h = Fnv1aCombine(h, static_cast<uint64_t>(v_it->second));
      }

      for (const auto& input : node->Inputs()) {
        uint64_t input_h = hash_node_internal(input, memo, var_indices);
        h = Fnv1aCombine(h, input_h);
      }

      h = HashNodeAttributes(node, h);

      memo[node->Id()] = h;
      return h;
    }

  }  // namespace

  uint64_t ComputeAstHash(const std::shared_ptr<Node>& node) {
    std::unordered_map<uint64_t, uint64_t> memo;
    std::unordered_map<uint64_t, size_t> var_indices;
    uint64_t h = hash_node_internal(node, memo, var_indices);
    h = Fnv1aCombine(h, var_indices.size());
    return h;
  }

}  // namespace tenzr