Clip slice range expressions

csrc/ops/alias.cpp

@@ -691,12 +691,18 @@ TensorView* slice(TensorView* inp, const std::vector<Slice>& ranges) {
       out_root_id = inp_root_id->cloneWithoutRFactor();
       out_rf_id = out_root_id;
     } else {
+      // Clip the start and stop values to the extent of the input
+      auto clipped_start =
+          SimplifyingIrBuilder::minExpr(inp_root_id->extent(), range.start);
+      auto clipped_stop =
+          SimplifyingIrBuilder::minExpr(inp_root_id->extent(), range.stop);
+
       out_root_id =
           IterDomainBuilder(inp_root_id).is_rfactor_domain(true).build();
       out_rf_id = IterDomain::resize(
           out_root_id,
-          SimplifyingIrBuilder::negExpr(range.start),
-          sub(range.stop, inp_root_id->extent()),
+          SimplifyingIrBuilder::negExpr(clipped_start),
+          sub(clipped_stop, inp_root_id->extent()),
           true);
       needs_real_slicing = true;
     }

test/test_resize.cpp

@@ -1122,6 +1122,63 @@ TEST_F(NVFuserTest, FusionResizeSlice5_CUDA) {
   testValidate(&fusion, cg_outputs, aten_inputs, {t2, t4}, __LINE__, __FILE__);
 }
 
+// Slice with end beyond size of input. This should clip to input, not pad.
+TEST_F(NVFuserTest, FusionResizeSlice6_CUDA) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  std::vector<int64_t> shape({9});
+
+  // concrete shapes to avoid dynamic Fusion
+  auto tv0 = makeConcreteTensor(shape);
+  fusion.addInput(tv0);
+
+  auto tv1 = slice(tv0, {{fusion.zeroVal(), IrBuilder::create<Int>(11)}});
+  fusion.addOutput(tv1);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+
+  auto t0 = at::randn(shape, options);
+  std::vector<c10::IValue> aten_inputs({t0});
+
+  FusionExecutor fe;
+  fe.compileFusion(&fusion, aten_inputs);
+  auto cg_outputs = fe.runFusion(aten_inputs);
+
+  auto ref = t0.index({at::indexing::Slice(0, 11)});
+
+  testValidate(&fusion, cg_outputs, aten_inputs, {ref}, __LINE__, __FILE__);
+}
+
+// Slice with start beyond size of input. This should produce zero-size tensor.
+TEST_F(NVFuserTest, FusionResizeSlice7_CUDA) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  std::vector<int64_t> shape({9});
+
+  // concrete shapes to avoid dynamic Fusion
+  auto tv0 = makeConcreteTensor(shape);
+  fusion.addInput(tv0);
+
+  auto tv1 =
+      slice(tv0, {{IrBuilder::create<Int>(11), IrBuilder::create<Int>(13)}});
+  fusion.addOutput(tv1);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+
+  auto t0 = at::randn(shape, options);
+  std::vector<c10::IValue> aten_inputs({t0});
+
+  FusionExecutor fe;
+  fe.compileFusion(&fusion, aten_inputs);
+  auto cg_outputs = fe.runFusion(aten_inputs);
+
+  auto ref = t0.index({at::indexing::Slice(11, 13)});
+
+  testValidate(&fusion, cg_outputs, aten_inputs, {ref}, __LINE__, __FILE__);
+}
+
 // Auto scheduled version of Slice1
 TEST_F(NVFuserTest, FusionResizeSliceScheduler1_CUDA) {
   auto fusion_ptr = std::make_unique<Fusion>();