Custom ops with dynamic tensor

[Olgacmt]

Is there a way to define a dynamic tensor and provide it as an input parameter into a custom operator?

For example, mult4 function from examples.

class MyModel(torch.nn.Module):
  def forward(x, y):
    return torch.my_ops.mult4(x + y)

def main():
    torch.ops.load_library("libcustom_ops_aot_lib.dylib")

    with torch.no_grad():
        tags = {
            "torch.cond",
            "torch.dynamic-shape",
            "torch.dynamic-value",
            "torch.escape-hatch",
        }
        model = CustomOpsConditionModel()
        # Define dynamic dimensions
        batch_x_dim = torch.export.Dim("batch_x", min=1, max=100)
        seq_len_x_dim = torch.export.Dim("seq_len_x", min=1, max=256)
        batch_y_dim = torch.export.Dim("batch_y", min=1, max=50)
        seq_len_y_dim = torch.export.Dim("seq_len_y", min=1, max=256)

        # Dynamic shapes configuration
        dynamic_shapes = {
            "x": {0: 1, 1: seq_len_x_dim, 2: 3},  # 'x' is dynamic
            "y": {0: 1, 1: seq_len_y_dim, 2: 3}   # 'y' is dynamic
        }

        # Create example inputs within the specified range
        batch_x = 1
        seq_len_x = 100
        feat_x = 3
        batch_y = 1
        seq_len_y = 100
        feat_y = 3

        x = torch.randn(batch_x, seq_len_x, feat_x, dtype=torch.float32, requires_grad=False)
        y = torch.randn(batch_y, seq_len_y, feat_y, dtype=torch.float32, requires_grad=False)

        example_inputs = (x, y)
        model = torch.export.export(model, example_inputs, dynamic_shapes=dynamic_shapes, strict=False)
        print(model.graph_module.print_readable())

    edge_config = get_xnnpack_edge_compile_config()
    edge_manager = to_edge(model, compile_config=edge_config)
    edge_manager = edge_manager.to_backend(XnnpackPartitioner())
    exec_prog = edge_manager.to_executorch()

    model_name = "custom_ops_example"
    with open(f"{model_name}.pte", "wb") as file:
        file.write(exec_prog.buffer)

Currently, it fails with an error:

torch._dynamo.exc.UserError: Constraints violated (seq_len_y)! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of seq_len_y = L['args'][0][1].size()[1] in the specified range seq_len_y <= 256 are valid because seq_len_y was inferred to be a constant (100).

[JacobSzwejbka]

Is it currently failing in torch.export.export? You would probably get better support in pytorch/pytorch issues. cc @tugsbayasgalan @angelayi

[tarun292]

@angelayi can confirm but i'm guessing the issue might be that there's no meta kernel here. Can you follow the instructions here to add a meta kernel.

https://pytorch.org/docs/stable/export.html#missing-fake-meta-abstract-kernels-for-operators

[angelayi]

Yeah I'd be curious to see your implementation of mult4!

[Olgacmt]

My mul4 impl:

Tensor mul4_impl(const Tensor &in) {
  // naive approach
  at::Tensor out = at::zeros_like(in);
  out.copy_(in);
  out.mul_(4);
  return out;
}

TORCH_LIBRARY_FRAGMENT(my_ops, m) {
  m.def("my_ops::mul4(Tensor input) -> Tensor");
}

TORCH_LIBRARY_IMPL(my_ops, CompositeExplicitAutograd, m) {
  m.impl("mul4", TORCH_FN(mul4_impl));
}

and

// mul4.out(Tensor input, *, Tensor(a!) output) -> Tensor(a!)
// ATen-compatible function signature, without a KernelRuntimeContext.
Tensor& mul4_out_impl(const Tensor& in, Tensor& out) {
  check_preconditions(in, out);
  float* out_data = out.mutable_data_ptr<float>();
  const float* in_data = in.const_data_ptr<float>();

  const ssize_t dim = in.dim();
  if (dim == 0) {
    return out;
  }
  if (dim == 1) {
    for (size_t out_idx = 0; out_idx < out.numel(); ++out_idx) {
      out_data[out_idx] = in_data[out_idx] * 4;
    }
    return out;
  }
  if (dim == 2) {
    for (size_t i = 0; i < in.size(0); ++i) {
      for (size_t j = 0; j < in.size(1); ++j) {
        out_data[i * out.size(1) + j] = in_data[i * in.size(1) + j] * 4;
      }
    }
    return out;
  }
  return out;
}

// mul4.out(Tensor input, *, Tensor(a!) output) -> Tensor(a!)
// ExecuTorch-compatible function signature with a KernelRuntimeContext.
Tensor& mul4_out_impl(
    ET_UNUSED KernelRuntimeContext& ctx,
    const Tensor& in,
    Tensor& out) {
  out = mul4_out_impl(in, out);
  print_impl("mul4_out", out);
  return out;
}

[angelayi]

Yeah I believe you're missing a fake kernel -- could you try adding one? Here are some docs on how you could do so: https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit?tab=t.0#heading=h.xvrg7clz290

[Olgacmt]

Good idea. Checking if it helps fix the issue.