Add TensorViewBuilder::shape(std::vector<Val*> shape)

torch/csrc/jit/codegen/cuda/executor_utils.cpp

@@ -794,7 +794,7 @@ kir::ExpressionEvaluator bindKernelInputs(
                 extent->toString(),
                 " to ",
                 value,
-                "but it's already set to ",
+                " but it's already set to ",
                 *prev_value);
             should_bind = false;
           }

torch/csrc/jit/codegen/cuda/ir_interface_nodes.h

@@ -572,7 +572,8 @@ class TORCH_CUDA_CU_API TensorViewBuilder {
   TensorViewBuilder& contiguity(std::vector<bool> contiguity);
 
   //! Set the shape (default 0 dimensional, ie. scalar)
-  TensorViewBuilder& shape(std::vector<int64_t> shape);
+  TensorViewBuilder& shape(std::vector<Val*> shape);
+  TensorViewBuilder& shape(const std::vector<int64_t>& shape);
 
   //! Creates a new TensorView with the specified options
   TensorView* build() const;
@@ -581,7 +582,7 @@ class TORCH_CUDA_CU_API TensorViewBuilder {
   size_t ndims_ = 0;
   DataType dtype_ = DataType::Float;
   std::vector<bool> contiguity_;
-  std::vector<int64_t> shape_;
+  std::vector<Val*> shape_;
 };
 
 } // namespace cuda

torch/csrc/jit/codegen/cuda/ops/normalization.cpp

@@ -529,8 +529,8 @@ ForwardNormResult batch_norm(
     auto invstd_bcast = broadcast(unbiased_invstd, broadcast_mask);
 
     // During inference, mean/invstd output are empty tensors
-    mean = TensorViewBuilder().shape({0}).build();
-    invstd = TensorViewBuilder().shape({0}).build();
+    mean = TensorViewBuilder().shape(std::vector<int64_t>{0}).build();
+    invstd = TensorViewBuilder().shape(std::vector<int64_t>{0}).build();
     y = mul(x_sub_mean, invstd_bcast);
   }
 
@@ -782,8 +782,8 @@ ForwardNormResult instance_norm(
         broadcast(unbiased_invstd, channels_only_broadcast_mask);
 
     // During inference, mean/invstd output are empty tensors
-    mean = TensorViewBuilder().shape({0}).build();
-    invstd = TensorViewBuilder().shape({0}).build();
+    mean = TensorViewBuilder().shape(std::vector<int64_t>{0}).build();
+    invstd = TensorViewBuilder().shape(std::vector<int64_t>{0}).build();
     y = mul(x_sub_mean, invstd_bcast);
   }
 

torch/csrc/jit/codegen/cuda/tensor_view.cpp

@@ -1193,7 +1193,29 @@ TensorViewBuilder& TensorViewBuilder::contiguity(std::vector<bool> contiguity) {
   return *this;
 }
 
-TensorViewBuilder& TensorViewBuilder::shape(std::vector<int64_t> shape) {
+TensorViewBuilder& TensorViewBuilder::shape(const std::vector<int64_t>& shape) {
+  TORCH_CHECK(shape_.empty(), "Attempting to reset shape");
+  if (!shape.empty()) {
+    TORCH_CHECK(ndims_ == 0 || ndims_ == shape.size());
+    ndims_ = shape.size();
+  }
+  shape_.clear();
+  shape_.reserve(shape.size());
+  for (int64_t i : shape) {
+    if (i == -1) {
+      shape_.emplace_back(IrBuilder::create<Int>());
+    } else {
+      TORCH_CHECK(
+          i >= 0,
+          "Invalid extent value. ",
+          "For a tensor representing a single scalar use ndims = 0 with no sizes set.");
+      shape_.emplace_back(IrBuilder::create<Int>(i));
+    }
+  }
+  return *this;
+}
+
+TensorViewBuilder& TensorViewBuilder::shape(std::vector<Val*> shape) {
   TORCH_CHECK(shape_.empty(), "Attempting to reset shape");
   if (!shape.empty()) {
     TORCH_CHECK(ndims_ == 0 || ndims_ == shape.size());
@@ -1207,28 +1229,23 @@ TensorView* TensorViewBuilder::build() const {
   // Build the domain
   std::vector<IterDomain*> domain(ndims_, nullptr);
   for (const auto i : c10::irange(ndims_)) {
-    if (shape_.empty() || shape_[i] == -1) {
+    if (shape_.empty()) {
       domain[i] =
           IterDomainBuilder(
               FusionGuard::getCurFusion()->zeroVal(), IrBuilder::create<Int>())
               .build();
     } else {
-      TORCH_CHECK(
-          shape_[i] >= 0,
-          "Invalid extent value. ",
-          "For a tensor representing a single scalar use ndims = 0 with no sizes set.");
-      if (shape_[i] == 1) {
+      if (shape_[i]->isOneInt()) {
         // If size is known to be 1, assume it needs to be broadcasted.
         domain[i] = IterDomainBuilder(
                         FusionGuard::getCurFusion()->zeroVal(),
                         FusionGuard::getCurFusion()->oneVal())
                         .iter_type(IterType::Broadcast)
                         .build();
       } else {
-        domain[i] = IterDomainBuilder(
-                        FusionGuard::getCurFusion()->zeroVal(),
-                        IrBuilder::create<Int>(shape_[i]))
-                        .build();
+        domain[i] =
+            IterDomainBuilder(FusionGuard::getCurFusion()->zeroVal(), shape_[i])
+                .build();
       }
     }
   }

torch/csrc/jit/codegen/cuda/test/test_gpu.cpp

@@ -1,4 +1,5 @@
 #if defined(USE_CUDA)
+#include <gmock/gmock-matchers.h>
 #include <gtest/gtest.h>
 
 #include <torch/csrc/jit/codegen/cuda/arith.h>
@@ -25409,6 +25410,79 @@ TEST_F(NVFuserTest, FusionPrint_CUDA) {
   }
 }
 
+TEST_F(NVFuserTest, FusionCheckedSymbolicShape_CUDA) {
+  const auto options =
+      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+
+  at::Tensor a = at::randn({123, 456}, options);
+  at::Tensor b = at::randn({123, 456}, options);
+  at::Tensor c = at::randn({321, 654}, options);
+
+  using return_t =
+      std::pair<std::unique_ptr<FusionExecutorCache>, std::vector<at::Tensor>>;
+  auto matched_add = [](at::Tensor a, at::Tensor b) -> return_t {
+    auto fusion = std::make_unique<Fusion>();
+    FusionGuard fg(fusion.get());
+
+    Val* s1 = IrBuilder::create<Int>();
+    Val* s2 = IrBuilder::create<Int>();
+    auto builder = TensorViewBuilder().shape(std::vector<Val*>{s1, s2});
+    TensorView* tv0 = builder.build();
+    TensorView* tv1 = builder.build();
+
+    fusion->addInput(tv0);
+    fusion->addInput(tv1);
+
+    auto tv2 = add(tv0, tv1);
+
+    fusion->addOutput(tv2);
+
+    auto executor_cache =
+        std::make_unique<FusionExecutorCache>(std::move(fusion));
+    auto cg_outputs = executor_cache->runFusionWithInputs({a, b});
+    return {std::move(executor_cache), std::move(cg_outputs)};
+  };
+
+  {
+    auto ret1 = matched_add(a, b);
+    testValidate(
+        ret1.first->fusion(), ret1.second, {a, b}, {a + b}, __LINE__, __FILE__);
+  }
+
+  {
+    EXPECT_THAT(
+        [&]() { matched_add(a, c); },
+        ::testing::ThrowsMessage<c10::Error>(
+            ::testing::HasSubstr("Attempting to bind")));
+  }
+}
+
+TEST_F(NVFuserTest, FusionSizeDependentData_CUDA) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  Val* s1 = IrBuilder::create<Int>();
+  auto builder = TensorViewBuilder().shape(std::vector<Val*>{s1});
+  TensorView* tv0 = builder.build();
+
+  fusion->addInput(tv0);
+
+  auto tv1 = add(tv0, s1);
+
+  fusion->addOutput(tv1);
+
+  const auto options =
+      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+
+  at::Tensor a = at::zeros({123}, options);
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  auto cg_outputs = executor_cache.runFusionWithInputs({a});
+
+  testValidate(
+      executor_cache.fusion(), cg_outputs, {a}, {a + 123}, __LINE__, __FILE__);
+}
+
 } // namespace jit
 } // namespace torch
 #endif // #if defined(USE_CUDA)