add HirAliasSelect

[samnordmann]

What

Add a SelectOp-like HIR to express indexing into ATen tensor.

Why

it is used in the context of stream lowering, see #4147 and
especially the discussion in #4147 (comment)

[github-actions]

Review updated until commit 82bc595

Description

• Added HirAliasSelect for tensor indexing in HIR.

• Implemented handling in HostIrEvaluator.

• Added tests for HirAliasSelect.

• Updated dispatch and header files.

---

Changes walkthrough 📝

	Relevant files
Enhancement	executor.cpp Implement HirAliasSelect handling                                                csrc/host_ir/executor.cpp Added handle method for HirAliasSelect. +9/-0      host_ir.cpp Define HirAliasSelect class                                                            csrc/host_ir/host_ir.cpp Added HirAliasSelect class definition. Implemented constructor and methods. +45/-0    executor.h Declare HirAliasSelect handling                                                    csrc/host_ir/executor.h Added handle method declaration for HirAliasSelect. +1/-0      host_ir.h Declare HirAliasSelect class                                                          csrc/host_ir/host_ir.h Added HirAliasSelect class declaration. +43/-0
Tests	test_host_irs.cpp Add HirAliasSelect test                                                                    tests/cpp/test_host_irs.cpp Added test case for HirAliasSelect. +38/-0
Configuration changes	dispatch.h Update dispatch list                                                                          csrc/dispatch.h Added HirAliasSelect to dispatch list. +2/-1

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Axis Validation The axis validation in the constructor of HirAliasSelect checks if the axis is within the bounds of the input tensor's dimensions. However, it uses in()->getLogicalDomain().at(axis) in the error message, which might be incorrect if the axis is out of bounds. It should use axis directly in the error message for clarity. static_cast<int64_t>(in->getLogicalDomain().size()) > axis, "Select axis ", axis, " is out of bounds for tensor ", in->toString(), " with ", in->getLogicalDomain().size(), Output Handling The comment mentions that "out" is not added as an output because the current op doesn't "define" it, but rather sets its allocation. This might lead to confusion. It would be helpful to clarify why "out" is treated differently from other outputs. // but rather sets its allocation. Since "out" will be used in another // producing expression, this avoids unnecessary cyclic dependencies. This // ressembles how kir::Allocate treats its allocated TensorView. Test Coverage The test SelectingTensor covers a basic case. It would be beneficial to add more test cases to cover edge cases, such as selecting the first or last dimension, and using different data types. using HirAliasSelectHostIrTest = NVFuserTest; TEST_F(HirAliasSelectHostIrTest, SelectingTensor) { constexpr int64_t ndims = 2; constexpr int64_t dim = 1; constexpr int64_t index = 3; const std::vector<int64_t> input_sizes = {32, 32}; ASSERT_LT(dim, ndims); ASSERT_EQ(input_sizes.size(), ndims); ASSERT_LT(index, input_sizes.at(dim)); auto hic = std::make_unique<HostIrContainer>(); FusionGuard fg(hic.get()); TensorView* in = makeContigTensor(ndims); TensorView* out = makeContigTensor(ndims - 1); auto* index_val = IrBuilder::create<Val>(index, DataType::Index); auto* select_op = IrBuilder::create<HirAliasSelect>(in, out, dim, index_val); hic->addInput(in); hic->addOutput(out); hic->pushBackTopLevelExprs(select_op); HostIrEvaluator hie(std::move(hic)); auto options = at::TensorOptions().device(at::kCUDA, 0).dtype(torch::kFloat); auto in_aten = at::randn(input_sizes, options); std::unordered_map<Val*, PolymorphicValue> concrete_input_buffers = { {in, in_aten}}; auto out_aten = hie.runWithInput(concrete_input_buffers)[0].as<at::Tensor>(); // validate auto ref_out = in_aten.select(dim, index); EXPECT_TRUE(ref_out.equal(out_aten)); }

[samnordmann]

!test

[samnordmann]

!test

[samnordmann]

!test

[samnordmann]

!test

[samnordmann]

!test