Add full op

Author: rdspring1

python/python_direct/ops.cpp

@@ -2281,6 +2281,15 @@ TensorView
 )");
 }
 
+template <class ShapeType>
+TensorView* full_op_fn(
+    ShapeType generic_output_shape,
+    Val* fill_value,
+    PrimDataType dtype) {
+  std::vector<Val*> output_shape = SequenceAsVector(generic_output_shape);
+  return full(output_shape, fill_value, dtype);
+}
+
 void bindTensorFactoryOps(py::module_& ops) {
   ops.def(
       "iota",
@@ -2308,6 +2317,37 @@ Returns
 -------
 TensorView
     The tensor with values from 0 to length-1.
+)",
+      py::return_value_policy::reference);
+  ops.def(
+      "full",
+      full_op_fn<py::list>,
+      py::arg("shape"),
+      py::arg("fill_value"),
+      py::arg("dtype"),
+      py::return_value_policy::reference);
+  ops.def(
+      "full",
+      full_op_fn<py::tuple>,
+      py::arg("shape"),
+      py::arg("fill_value"),
+      py::arg("dtype"),
+      R"(
+Create a tensor with all elements set to a specified value.
+
+Parameters
+----------
+shape : list or tuple
+    The shape of the tensor.
+fill_value : Val
+    The value to fill the tensor with.
+dtype : PrimDataType
+    The data type of the tensor.
+
+Returns
+-------
+TensorView
+    The tensor with all elements set to the specified value.
 )",
       py::return_value_policy::reference);
 }

python/python_direct/python_translate.cpp

@@ -981,6 +981,25 @@ class PythonTranslator : public OptInConstDispatch {
         {lsop->out()});
   }
 
+  // Map FullOp to python frontend
+  void handle(const FullOp* fop) final {
+    NVF_ERROR(fop != nullptr);
+    TensorView* out_tv = fop->output(0)->as<TensorView>();
+    visited_vals_.insert(out_tv);
+
+    // Fill value can be dynamic so create it
+    dispatch(fop->getFillValue());
+
+    static const std::vector<std::string> argument_names = {
+        "shape", "fill_value", "dtype"};
+    printer_.generateKwargsOperation(
+        "fd.ops.full",
+        std::make_tuple(),
+        argument_names,
+        std::make_tuple(getShape(out_tv), fop->getFillValue(), out_tv->dtype()),
+        {out_tv});
+  }
+
   // Map IotaOp to python frontend
   void handle(const IotaOp* iop) final {
     NVF_ERROR(iop != nullptr);

tests/python/direct/test_python_frontend.py

@@ -910,3 +910,23 @@ def fusion_func(fd: FusionDefinition):
     eager_out2 = torch.tensor([100, 101, 102], dtype=torch.int, device="cuda")
     nvfuser_direct_test.assertEqual(eager_out1, nvf_out[0])
     nvfuser_direct_test.assertEqual(eager_out2, nvf_out[1])
+
+
+def test_scalar_only_inputs(nvfuser_direct_test):
+    # We don't allow scalar outputs, currently,
+    # so a tensor has to be returned
+    def fusion_func(fd: FusionDefinition):
+        s0 = fd.define_scalar()
+        s1 = fd.define_scalar()
+        s2 = fd.ops.add(s0, s1)
+        c0 = fd.define_scalar(1.0, DataType.Float)
+        t3 = fd.ops.full(shape=[2, 2], fill_value=c0, dtype=DataType.Float)
+        t4 = fd.ops.mul(t3, s2)
+        fd.add_output(t4)
+
+    with FusionDefinition() as fd:
+        fusion_func(fd)
+
+    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [2.0, 3.0])
+    eager_out = torch.full([2, 2], 1.0) * 5.0
+    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

tests/python/opinfo/opinfos.py

@@ -1311,6 +1311,7 @@ def torch_reshape_sym_fn(input_tensor, output_shaped_tensor):
         ArgumentType.Symbolic,
         ArgumentType.Constant,
     ),
+    supports_direct_bindings=True,
 )
 tensor_creation_ops.append(full_opinfo)