[AdvancedCompiler] Fill(cpp wrapper) #759
Conversation
lib/fill.cpp
Outdated
| using namespace triton_jit; | ||
|
|
||
| // ----------- scalar fill：创建新张量并用标量填充 ----------- | ||
| at::Tensor fill_scalar(const at::Tensor& input, double value) { |
There was a problem hiding this comment.
aten's fill has schema
- fill.Scalar(Tensor self, Scalar value) -> Tensor
- fill.Tensor(Tensor self, Tensor value) -> Tensor
- fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!)
- fill_.Tensor(Tensor(a!) self, Tensor value) -> Tensor(a!)
c10::Scalar is torch's type for a numeric scalar whose type is not known at compile time. It is common for torch to use a macro to expand the cases with its underlaying dtype.
Tensor& scalar_fill(Tensor& self, const Scalar& value) {
AT_DISPATCH_V2(
self.scalar_type(), "fill_out", AT_WRAP([&]() {
fill_inplace<scalar_t>(self, value);
}), kComplexHalf, kHalf, kBool, kBFloat16, AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX), AT_EXPAND(AT_FLOAT8_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));
return self;
}We can also use these macros, but it may be a little hard to get it right.
There was a problem hiding this comment.
Maybe I should try to get it work.
There was a problem hiding this comment.
Scalar type is not same as basic type like float?
lib/fill.cpp
Outdated
| using namespace triton_jit; | ||
|
|
||
| // ----------- scalar fill：创建新张量并用标量填充 ----------- | ||
| at::Tensor fill_scalar(const at::Tensor& input, double value) { |
There was a problem hiding this comment.
Scalar type is not same as basic type like float?
include/flag_gems/operators.h
Outdated
| @@ -31,5 +30,16 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding( | |||
| const std::optional<at::Tensor> &position_ids = std::nullopt, | |||
| bool rotary_interleaved = false); | |||
|
|
|||
| at::Tensor fill_scalar(const at::Tensor &input, c10::Scalar value); | |||
There was a problem hiding this comment.
Use const c10::Scalar& for argument type Scalar in native_functions.yaml since it is the convention.
Check aten/src/ATen/native/Fill.cpp for details. We should follow the signature.
Tensor& fill_(Tensor& self, const Scalar& value);
Tensor& fill_(Tensor& self, const Tensor& value);
Tensor fill(const Tensor& self, const Scalar& value);
Tensor fill(const Tensor& self, const Tensor& value);
src/flag_gems/csrc/cstub.cpp
Outdated
| @@ -22,6 +22,11 @@ TORCH_LIBRARY(flag_gems, m) { | |||
| "rotary_embedding(Tensor q, Tensor k, Tensor cos, Tensor sin, Tensor? position_ids=None, " | |||
| "bool rotary_interleaved=False) -> (Tensor, Tensor)"); // q and k may be view to other size | |||
| m.def("bmm(Tensor self, Tensor mat2) -> Tensor"); | |||
| // fill operator declaration | |||
| m.def("fill_scalar(Tensor input, double value) -> Tensor"); | |||
There was a problem hiding this comment.
Be consistent with the function signature, use Scalar for fill_scalar.
Also, we suggest
- using the same signature as in Aten, since we may also register it for ATen in the future.
- use marker to mark aliasing, and ! to mark mutation.
fill.Scalar and fill.Tensor has the same op name with different overload.
fill.Scalar(Tensor self, Scalar value) -> Tensor
fill.Tensor(Tensor self, Tensor value) -> Tensor
fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!)
fill_.Tensor(Tensor(a!) self, Tensor value) -> Tensor(a!)
…:Tensor& for consistency with PyTorch conventions
|
Since pr #800 has been merged, I will close this pr. |
PR Category
Operator
Type of Change
New Feature
Description
CPP wrapper packaging for fill op.
Issue
Progress
Performance