CSupportedOperations.cc

/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the following additional limitation. Functionality enabled by the
 * files subject to the Elastic License 2.0 may only be used in production when
 * invoked by an Elasticsearch process with a license key installed that permits
 * use of machine learning features. You may not use this file except in
 * compliance with the Elastic License 2.0 and the foregoing additional
 * limitation.
 */

#include "CSupportedOperations.h"

namespace ml {
namespace torch {

using namespace std::string_view_literals;

const CSupportedOperations::TStringViewSet CSupportedOperations::FORBIDDEN_OPERATIONS = {
    // Arbitrary memory access — enables heap scanning, address leaks, and
    // ROP chain construction.
    "aten::as_strided"sv,
    "aten::from_file"sv,
    "aten::save"sv,
    // After graph inlining, method and function calls should be resolved.
    // Their presence indicates an opaque call that cannot be validated.
    "prim::CallFunction"sv,
    "prim::CallMethod"sv,
};

// Generated by dev-tools/extract_model_ops/extract_model_ops.py against PyTorch 2.7.1.
// Reference models: bert-base-uncased, roberta-base, distilbert-base-uncased,
// google/electra-small-discriminator, microsoft/mpnet-base,
// microsoft/deberta-base, facebook/dpr-ctx_encoder-single-nq-base,
// google/mobilebert-uncased, xlm-roberta-base, elastic/bge-m3,
// elastic/distilbert-base-{cased,uncased}-finetuned-conll03-english,
// elastic/eis-elser-v2, elastic/elser-v2, elastic/hugging-face-elser,
// elastic/multilingual-e5-small-optimized, elastic/splade-v3,
// elastic/test-elser-v2.
// Additional ops from Elasticsearch integration test models
// (PyTorchModelIT, TextExpansionQueryIT, TextEmbeddingQueryIT).
const CSupportedOperations::TStringViewSet CSupportedOperations::ALLOWED_OPERATIONS = {
    // aten operations (core tensor computations)
    "aten::Int"sv,
    "aten::IntImplicit"sv,
    "aten::ScalarImplicit"sv,
    "aten::__and__"sv,
    "aten::abs"sv,
    "aten::add"sv,
    "aten::add_"sv,
    "aten::arange"sv,
    "aten::bitwise_not"sv,
    "aten::cat"sv,
    "aten::chunk"sv,
    "aten::clamp"sv,
    "aten::contiguous"sv,
    "aten::cumsum"sv,
    "aten::div"sv,
    "aten::div_"sv,
    "aten::dropout"sv,
    "aten::embedding"sv,
    "aten::expand"sv,
    "aten::full_like"sv,
    "aten::gather"sv,
    "aten::ge"sv,
    "aten::gelu"sv,
    "aten::hash"sv,
    "aten::index"sv,
    "aten::index_put_"sv,
    "aten::layer_norm"sv,
    "aten::len"sv,
    "aten::linear"sv,
    "aten::log"sv,
    "aten::lt"sv,
    "aten::manual_seed"sv,
    "aten::masked_fill"sv,
    "aten::matmul"sv,
    "aten::max"sv,
    "aten::mean"sv,
    "aten::min"sv,
    "aten::mul"sv,
    "aten::ne"sv,
    "aten::neg"sv,
    "aten::new_ones"sv,
    "aten::ones"sv,
    "aten::pad"sv,
    "aten::permute"sv,
    "aten::pow"sv,
    "aten::rand"sv,
    "aten::relu"sv,
    "aten::repeat"sv,
    "aten::reshape"sv,
    "aten::rsub"sv,
    "aten::scaled_dot_product_attention"sv,
    "aten::select"sv,
    "aten::size"sv,
    "aten::slice"sv,
    "aten::softmax"sv,
    "aten::sqrt"sv,
    "aten::squeeze"sv,
    "aten::str"sv,
    "aten::sub"sv,
    "aten::tanh"sv,
    "aten::tensor"sv,
    "aten::to"sv,
    "aten::transpose"sv,
    "aten::type_as"sv,
    "aten::unsqueeze"sv,
    "aten::view"sv,
    "aten::where"sv,
    "aten::zeros"sv,
    // prim operations (TorchScript graph infrastructure)
    "prim::Constant"sv,
    "prim::DictConstruct"sv,
    "prim::GetAttr"sv,
    "prim::If"sv,
    "prim::ListConstruct"sv,
    "prim::ListUnpack"sv,
    "prim::Loop"sv,
    "prim::NumToTensor"sv,
    "prim::TupleConstruct"sv,
    "prim::TupleUnpack"sv,
    "prim::device"sv,
    "prim::dtype"sv,
    "prim::max"sv,
    "prim::min"sv,
};
}
}