Skip to content

[SWDEV-445018/SWDEV-445019] Retain regular mm functionality after f8 commit #1368

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 18, 2024
Merged

[SWDEV-445018/SWDEV-445019] Retain regular mm functionality after f8 commit #1368

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
edc4b8a
Retain regular matmul behaviour after f8 commit
jataylo Mar 13, 2024
f285ba1
Spacing
jataylo Mar 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 75 additions & 57 deletions torch/csrc/autograd/FunctionsManual.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1423,43 +1423,48 @@ Tensor mm_mat1_backward(
c10::Layout mat1_layout,
const Scalar& alpha) {

at::Tensor grad_copy = grad.clone();
at::Tensor mat2_copy = mat2.clone();
if (at::globalContext().allowF8ROCMGemm()) {
at::Tensor grad_copy = grad.clone();
at::Tensor mat2_copy = mat2.clone();

if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "FunctionalsManual.cpp: mm_mat1_backward(const Tensor& grad, const Tensor& mat2,...)" << std::endl;
std::cout << "\tgrad: " << &grad << std::endl;
std::cout << "\tmat2: " << &mat2 << std::endl;
std::cout << "\tgrad_copy: " << &grad_copy << std::endl;
std::cout << "\tmat2_copy: " << &mat2_copy << std::endl;
std::cout << "\tgrad_copy.is_grad(): before: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat2_copy.is_grad(): before: " << std::boolalpha << mat2_copy.is_grad() << std::endl;
}
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "FunctionalsManual.cpp: mm_mat1_backward(const Tensor& grad, const Tensor& mat2,...)" << std::endl;
std::cout << "\tgrad: " << &grad << std::endl;
std::cout << "\tmat2: " << &mat2 << std::endl;
std::cout << "\tgrad_copy: " << &grad_copy << std::endl;
std::cout << "\tmat2_copy: " << &mat2_copy << std::endl;
std::cout << "\tgrad_copy.is_grad(): before: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat2_copy.is_grad(): before: " << std::boolalpha << mat2_copy.is_grad() << std::endl;
}

grad_copy.set_is_grad(true);
mat2_copy.set_is_grad(false);
grad_copy.set_is_grad(true);
mat2_copy.set_is_grad(false);

if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tgrad_copy.is_grad(): after: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat2_copy.is_grad(): after: " << std::boolalpha << mat2_copy.is_grad() << std::endl;
}
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tgrad_copy.is_grad(): after: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat2_copy.is_grad(): after: " << std::boolalpha << mat2_copy.is_grad() << std::endl;
}

if (grad_copy.layout() == c10::kStrided && mat2_copy.layout() == c10::kStrided &&
mat1_layout == c10::kStrided) {
// if input was column-major, return grad as column-order for efficiency
if (mat1_strides[0] == 1 && mat1_strides[1] == mat1_sizes[0]) {
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(mat2_copy.conj().mm(grad_copy.t()).t(), alpha.conj())" << std::endl;
if (grad_copy.layout() == c10::kStrided && mat2_copy.layout() == c10::kStrided &&
mat1_layout == c10::kStrided) {
// if input was column-major, return grad as column-order for efficiency
if (mat1_strides[0] == 1 && mat1_strides[1] == mat1_sizes[0]) {
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(mat2_copy.conj().mm(grad_copy.t()).t(), alpha.conj())" << std::endl;
}
return maybe_multiply(mat2_copy.conj().mm(grad_copy.t()).t(), alpha.conj());
}
return maybe_multiply(mat2_copy.conj().mm(grad_copy.t()).t(), alpha.conj());
}

// General fallback, should work for any layout
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(grad_copy.mm(mat2_copy.t().conj()), alpha.conj())" << std::endl;
}
return maybe_multiply(grad_copy.mm(mat2_copy.t().conj()), alpha.conj());
}

// General fallback, should work for any layout
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(grad_copy.mm(mat2_copy.t().conj()), alpha.conj())" << std::endl;
}
return maybe_multiply(grad_copy.mm(mat2_copy.t().conj()), alpha.conj());
return maybe_multiply(grad.mm(mat2.t().conj()), alpha.conj());
}

Tensor mm_mat2_backward(
Expand All @@ -1469,44 +1474,57 @@ Tensor mm_mat2_backward(
at::SymIntArrayRef mat2_strides,
c10::Layout mat2_layout,
const Scalar& alpha) {
at::Tensor grad_copy = grad.clone();
at::Tensor mat1_copy = mat1.clone();

if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "FunctionalsManual.cpp: mm_mat2_backward(const Tensor& grad, const Tensor& mat1,...)" << std::endl;
std::cout << "\tgrad: " << &grad << std::endl;
std::cout << "\tmat1: " << &mat1 << std::endl;
std::cout << "\tgrad_copy: " << &grad_copy << std::endl;
std::cout << "\tmat1_copy: " << &mat1_copy << std::endl;
std::cout << "\tgrad_copy.is_grad(): before: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat1_copy.is_grad(): before: " << std::boolalpha << mat1_copy.is_grad() << std::endl;
}
if (at::globalContext().allowF8ROCMGemm()) {
at::Tensor grad_copy = grad.clone();
at::Tensor mat1_copy = mat1.clone();

if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "FunctionalsManual.cpp: mm_mat2_backward(const Tensor& grad, const Tensor& mat1,...)" << std::endl;
std::cout << "\tgrad: " << &grad << std::endl;
std::cout << "\tmat1: " << &mat1 << std::endl;
std::cout << "\tgrad_copy: " << &grad_copy << std::endl;
std::cout << "\tmat1_copy: " << &mat1_copy << std::endl;
std::cout << "\tgrad_copy.is_grad(): before: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat1_copy.is_grad(): before: " << std::boolalpha << mat1_copy.is_grad() << std::endl;
}

grad_copy.set_is_grad(true);
mat1_copy.set_is_grad(false);
grad_copy.set_is_grad(true);
mat1_copy.set_is_grad(false);

if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tgrad_copy.is_grad(): after: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat1_copy.is_grad(): after: " << std::boolalpha << mat1_copy.is_grad() << std::endl;
}
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tgrad_copy.is_grad(): after: " << std::boolalpha << grad_copy.is_grad() << std::endl;
std::cout << "\tmat1_copy.is_grad(): after: " << std::boolalpha << mat1_copy.is_grad() << std::endl;
}


if (grad_copy.layout() == c10::kStrided && mat1_copy.layout() == c10::kStrided &&
mat2_layout == c10::kStrided) {
// if input was column-major, return grad as column-order for efficiency
if (mat2_strides[0] == 1 && mat2_strides[1] == mat2_sizes[0]) {
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(grad_copy.t().mm(mat1_copy.conj()).t(), alpha.conj())" << std::endl;
if (grad_copy.layout() == c10::kStrided && mat1_copy.layout() == c10::kStrided &&
mat2_layout == c10::kStrided) {
// if input was column-major, return grad as column-order for efficiency
if (mat2_strides[0] == 1 && mat2_strides[1] == mat2_sizes[0]) {
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(grad_copy.t().mm(mat1_copy.conj()).t(), alpha.conj())" << std::endl;
}
return maybe_multiply(grad_copy.t().mm(mat1_copy.conj()).t(), alpha.conj());
}
return maybe_multiply(grad_copy.t().mm(mat1_copy.conj()).t(), alpha.conj());
}
}

// General fallback, should work for any
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(mat1_copy.t().conj().mm(grad_copy), alpha.conj())" << std::endl;
// General fallback, should work for any
if (at::globalContext().allowF8ROCMLOG()) {
std::cout << "\tmaybe_multiply(mat1_copy.t().conj().mm(grad_copy), alpha.conj())" << std::endl;
}
return maybe_multiply(mat1_copy.t().conj().mm(grad_copy), alpha.conj());
} else {
if (grad.layout() == c10::kStrided && mat1.layout() == c10::kStrided &&
mat2_layout == c10::kStrided) {
// if input was column-major, return grad as column-order for efficiency
if (mat2_strides[0] == 1 && mat2_strides[1] == mat2_sizes[0]) {
return maybe_multiply(grad.t().mm(mat1.conj()).t(), alpha.conj());
}
}
// General fallback, should work for any layout
return maybe_multiply(mat1.t().conj().mm(grad), alpha.conj());
}
return maybe_multiply(mat1_copy.t().conj().mm(grad_copy), alpha.conj());
}

Tensor mm_mat1_sparse_backward(
Expand Down