[TorchToLinalg]Lower torch.gcd to linalg and scf #3732

bratislavSyrmia · 2024-09-25T14:10:26Z

Add verify() method to check if tensors are of
integer type. Also check if tensors are of same shape, or if the second tensor is a single element tensor.

Add e2e tests. Put them into onnx and stablehlo
xfailed sets.

bratislavSyrmia · 2024-09-30T14:25:26Z

force push: added math::cttz instead of counting trailing zeros manually

projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py

Add verify() method to check if tensors are of integer type. Also check if tensors are of same shape, or if the second tensor is a single element tensor. Add e2e tests. Put them into onnx and stablehlo xfailed sets.

vivekkhandelwal1

Hi @bratislavSyrmia, I would like you to explore a non-loop path for this lowering, since these kind of lowerings usually causes issues in the downstream pipeline especially the code-generation part.

IanWood1 · 2024-10-23T15:30:39Z

Hi @bratislavSyrmia, I would like you to explore a non-loop path for this lowering, since these kind of lowerings usually causes issues in the downstream pipeline especially the code-generation part.

@vivekkhandelwal1 out of curiosity, do you have an algorithm/solution in mind? The best I could think of was to use linalg.generic's library_call attr and define a GCD function but maybe that runs into the same problem.

vivekkhandelwal1 · 2024-10-25T11:58:50Z

Hi @bratislavSyrmia, I would like you to explore a non-loop path for this lowering, since these kind of lowerings usually causes issues in the downstream pipeline especially the code-generation part.

@vivekkhandelwal1 out of curiosity, do you have an algorithm/solution in mind? The best I could think of was to use linalg.generic's library_call attr and define a GCD function but maybe that runs into the same problem.

Hi @IanWood1, I did not spend time on thinking about it that's why I asked @bratislavSyrmia to explore the possibility of any such solution.

But if there exists a solution based on linalg.generic then it would still be a better approach then the current one.

bratislavSyrmia · 2024-10-31T11:05:35Z

Hi @bratislavSyrmia, I would like you to explore a non-loop path for this lowering, since these kind of lowerings usually causes issues in the downstream pipeline especially the code-generation part.

@vivekkhandelwal1 out of curiosity, do you have an algorithm/solution in mind? The best I could think of was to use linalg.generic's library_call attr and define a GCD function but maybe that runs into the same problem.

Hi @IanWood1, I did not spend time on thinking about it that's why I asked @bratislavSyrmia to explore the possibility of any such solution.

But if there exists a solution based on linalg.generic then it would still be a better approach then the current one.

I have thought about it but I have no idea how I would find the greatest common divisor between two numbers without using loops

bondhugula · 2024-11-25T23:52:16Z

Hi @bratislavSyrmia, I would like you to explore a non-loop path for this lowering, since these kind of lowerings usually causes issues in the downstream pipeline especially the code-generation part.

@vivekkhandelwal1 out of curiosity, do you have an algorithm/solution in mind? The best I could think of was to use linalg.generic's library_call attr and define a GCD function but maybe that runs into the same problem.

Hi @IanWood1, I did not spend time on thinking about it that's why I asked @bratislavSyrmia to explore the possibility of any such solution.
But if there exists a solution based on linalg.generic then it would still be a better approach then the current one.

I have thought about it but I have no idea how I would find the greatest common divisor between two numbers without using loops

This PR is acceptable and reviewable as is and I don't think it should be blocked because of a downstream's inability to deal with it. The lowering output is perfectly valid IR and it's contributing an otherwise missing lowering. Using loops itself, with a while loop (Euclid's GCD), it'd take log n steps while one could use an alternative approach with a countable for loop that would take O(n) steps. Downstream users could potentially add a lowering that uses a "GCD" intrinsic call to link with a library if really needed.

bondhugula · 2024-11-25T23:40:08Z

lib/Conversion/TorchToLinalg/Linear.cpp

+    auto other = adaptor.getOther(); // tensor B of the same size
+    auto loc = op.getLoc();
+
+    TensorType resultType =


bondhugula · 2024-11-25T23:40:46Z

lib/Conversion/TorchToLinalg/Linear.cpp

+    auto gcdPayloadBody = [&](OpBuilder &b, Location loc,
+                              ValueRange genericInstructionArgs) {
+      auto A = genericInstructionArgs[0];
+      A = b.create<mlir::math::AbsIOp>(loc, A);
+      auto B = genericInstructionArgs[1];
+      B = b.create<mlir::math::AbsIOp>(loc, B);
+      auto zero = b.create<mlir::arith::ConstantIntOp>(loc, 0, A.getType());
+
+      Value AtrailingZerosCount =
+          b.create<mlir::math::CountTrailingZerosOp>(loc, A);
+      Value BtrailingZerosCount =
+          b.create<mlir::math::CountTrailingZerosOp>(loc, B);
+      auto smalerZerosCount = b.create<mlir::arith::MinSIOp>(
+          loc, AtrailingZerosCount, BtrailingZerosCount);
+      auto shiftedA = b.create<mlir::arith::ShRSIOp>(loc, A, smalerZerosCount);
+      auto shiftedB = b.create<mlir::arith::ShRSIOp>(loc, B, smalerZerosCount);
+
+      auto findGcdConditionBlock = [&](mlir::OpBuilder &b, mlir::Location loc,
+                                       mlir::ValueRange innerLoopArgs) {
+        Value min = b.create<mlir::arith::MinSIOp>(loc, innerLoopArgs[0],
+                                                   innerLoopArgs[1]);
+        Value max = b.create<mlir::arith::MaxSIOp>(loc, innerLoopArgs[0],
+                                                   innerLoopArgs[1]);
+
+        auto cmp = b.create<mlir::arith::CmpIOp>(
+            loc, mlir::arith::CmpIPredicate::ne, min, zero);
+        b.create<mlir::scf::ConditionOp>(loc, cmp, ValueRange{min, max});
+      };
+      auto findGcdBodyBlock = [&](mlir::OpBuilder &b, mlir::Location loc,
+                                  mlir::ValueRange innerLoopArgs) {
+        Value min = innerLoopArgs[0];
+        Value max = innerLoopArgs[1];
+        max = b.create<mlir::arith::SubIOp>(loc, max, min);
+
+        Value maxTrailingZerosCount =
+            b.create<mlir::math::CountTrailingZerosOp>(loc, max);
+        max = b.create<mlir::arith::ShRSIOp>(loc, max, maxTrailingZerosCount);
+        b.create<mlir::scf::YieldOp>(loc, ValueRange{min, max});
+      };
+
+      auto findGcdWhileOp = b.create<mlir::scf::WhileOp>(
+          loc, TypeRange{shiftedA.getType(), shiftedB.getType()},
+          ValueRange{shiftedA, shiftedB}, findGcdConditionBlock,
+          findGcdBodyBlock);
+
+      Value gcdResult = findGcdWhileOp.getResult(1);
+      gcdResult =
+          b.create<mlir::arith::ShLIOp>(loc, gcdResult, smalerZerosCount);
+
+      b.create<linalg::YieldOp>(loc, gcdResult);
+    };
+
+    other = torch_to_linalg::createElementwiseLinalgGeneric(
+        rewriter, loc, ValueRange{self, other},
+        cast<TensorType>(self.getType()).getElementType(), gcdPayloadBody);
+


Missing code comments for all the major blocks and a high-level description on the top for the lowering.

bondhugula · 2024-11-25T23:48:57Z

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

+"    }\n"
+"    return %0#1 : !torch.int\n"
+"  }\n"
+"  func.func @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.is_integer_dtype(%arg0: !torch.int) -> !torch.bool {\n"


Why is this part moved?

The file is auto-generated so I think it's a quirk of update_abstract_interp_lib.sh. Its also checked by ci here:

torch-mlir/.github/workflows/ci.yml

Line 77 in 99115dc

bash build_tools/ci/check_generated_sources.sh

bondhugula · 2024-11-26T02:06:31Z

lib/Conversion/TorchToLinalg/Linear.cpp

+      auto findGcdConditionBlock = [&](mlir::OpBuilder &b, mlir::Location loc,
+                                       mlir::ValueRange innerLoopArgs) {
+        Value min = b.create<mlir::arith::MinSIOp>(loc, innerLoopArgs[0],
+                                                   innerLoopArgs[1]);
+        Value max = b.create<mlir::arith::MaxSIOp>(loc, innerLoopArgs[0],
+                                                   innerLoopArgs[1]);
+
+        auto cmp = b.create<mlir::arith::CmpIOp>(
+            loc, mlir::arith::CmpIPredicate::ne, min, zero);
+        b.create<mlir::scf::ConditionOp>(loc, cmp, ValueRange{min, max});
+      };
+      auto findGcdBodyBlock = [&](mlir::OpBuilder &b, mlir::Location loc,
+                                  mlir::ValueRange innerLoopArgs) {
+        Value min = innerLoopArgs[0];
+        Value max = innerLoopArgs[1];
+        max = b.create<mlir::arith::SubIOp>(loc, max, min);
+
+        Value maxTrailingZerosCount =
+            b.create<mlir::math::CountTrailingZerosOp>(loc, max);
+        max = b.create<mlir::arith::ShRSIOp>(loc, max, maxTrailingZerosCount);
+        b.create<mlir::scf::YieldOp>(loc, ValueRange{min, max});
+      };
+
+      auto findGcdWhileOp = b.create<mlir::scf::WhileOp>(
+          loc, TypeRange{shiftedA.getType(), shiftedB.getType()},
+          ValueRange{shiftedA, shiftedB}, findGcdConditionBlock,
+          findGcdBodyBlock);
+
+      Value gcdResult = findGcdWhileOp.getResult(1);
+      gcdResult =
+          b.create<mlir::arith::ShLIOp>(loc, gcdResult, smalerZerosCount);
+
+      b.create<linalg::YieldOp>(loc, gcdResult);
+    };


You don't need mlir:: anywhere here since there is a namespace using for it at the top.

CoTinker requested review from rsuderman, vivekkhandelwal1 and qingyunqu September 30, 2024 08:42

bratislavSyrmia force-pushed the lower_torch_aten_gcd_to_linalg_and_scf branch from 7673a8f to 0815cd1 Compare September 30, 2024 14:24

IanWood1 reviewed Sep 30, 2024

View reviewed changes

projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py Outdated Show resolved Hide resolved

[TorchToLinalg]Lower torch.gcd to linalg and scf

53b1ec3

Add verify() method to check if tensors are of integer type. Also check if tensors are of same shape, or if the second tensor is a single element tensor. Add e2e tests. Put them into onnx and stablehlo xfailed sets.

bratislavSyrmia force-pushed the lower_torch_aten_gcd_to_linalg_and_scf branch from 0815cd1 to 53b1ec3 Compare October 1, 2024 09:18

Merge branch 'main' into lower_torch_aten_gcd_to_linalg_and_scf

ee7f6ee

vivekkhandelwal1 reviewed Oct 23, 2024

View reviewed changes

bondhugula reviewed Nov 25, 2024

View reviewed changes

bondhugula reviewed Nov 26, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TorchToLinalg]Lower torch.gcd to linalg and scf #3732

[TorchToLinalg]Lower torch.gcd to linalg and scf #3732

bratislavSyrmia commented Sep 25, 2024

bratislavSyrmia commented Sep 30, 2024

vivekkhandelwal1 left a comment

IanWood1 commented Oct 23, 2024

vivekkhandelwal1 commented Oct 25, 2024

bratislavSyrmia commented Oct 31, 2024

bondhugula commented Nov 25, 2024

bondhugula Nov 25, 2024

bondhugula Nov 25, 2024

bondhugula Nov 25, 2024

IanWood1 Nov 26, 2024

bondhugula Nov 26, 2024

[TorchToLinalg]Lower torch.gcd to linalg and scf #3732

Are you sure you want to change the base?

[TorchToLinalg]Lower torch.gcd to linalg and scf #3732

Conversation

bratislavSyrmia commented Sep 25, 2024

bratislavSyrmia commented Sep 30, 2024

vivekkhandelwal1 left a comment

Choose a reason for hiding this comment

IanWood1 commented Oct 23, 2024

vivekkhandelwal1 commented Oct 25, 2024

bratislavSyrmia commented Oct 31, 2024

bondhugula commented Nov 25, 2024

bondhugula Nov 25, 2024

Choose a reason for hiding this comment

bondhugula Nov 25, 2024

Choose a reason for hiding this comment

bondhugula Nov 25, 2024

Choose a reason for hiding this comment

IanWood1 Nov 26, 2024

Choose a reason for hiding this comment

bondhugula Nov 26, 2024

Choose a reason for hiding this comment