Merge branch 'main' into keshavvinayak01/torch-aten-flex_attention

Author: keshavvinayak01

build_tools/ci/build_posix.sh

@@ -41,6 +41,7 @@ cmake -S "$repo_root/externals/llvm-project/llvm" -B "$build_dir" \
   -GNinja \
   -DCMAKE_BUILD_TYPE=Release \
   -DPython3_EXECUTABLE="$(which python3)" \
+  -DPython_EXECUTABLE="$(which python3)" \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DTORCH_MLIR_ENABLE_WERROR_FLAG=ON \
   -DCMAKE_INSTALL_PREFIX="$install_dir" \

build_tools/python_deploy/build_linux_packages.sh

@@ -244,6 +244,7 @@ function build_in_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir/externals/llvm-project/llvm
   cmake --build /main_checkout/torch-mlir/build --target tools/torch-mlir/all
   ccache -s
@@ -387,6 +388,7 @@ function build_out_of_tree() {
         -DLLVM_TARGETS_TO_BUILD=host \
         -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
         -DPython3_EXECUTABLE="$(which python3)" \
+        -DPython_EXECUTABLE="$(which python3)" \
         /main_checkout/torch-mlir/externals/llvm-project/llvm
     cmake --build /main_checkout/torch-mlir/llvm-build
   fi
@@ -409,6 +411,7 @@ function build_out_of_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir
   cmake --build /main_checkout/torch-mlir/build_oot
   ccache -s

build_tools/write_env_file.sh

@@ -13,7 +13,21 @@ portable_realpath() {
 
 td="$(portable_realpath "$(dirname "$0")"/..)"
 build_dir="$(portable_realpath "${TORCH_MLIR_BUILD_DIR:-$td/build}")"
-python_packages_dir="$build_dir/python_packages"
+
+in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
+out_of_tree_pkg_dir="${build_dir}/python_packages"
+
+if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
+  echo "Couldn't find in-tree or out-of-tree build, exiting."
+  exit 1
+fi
+
+# The `-nt` check works even if one of the two directories is missing.
+if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
+  python_packages_dir="${in_tree_pkg_dir}"
+else
+  python_packages_dir="${out_of_tree_pkg_dir}"
+fi
 
 write_env_file() {
   echo "Updating $build_dir/.env file"

docs/development.md

@@ -194,30 +194,48 @@ TIP: add multiple target options to stack build phases
 
 ### Setup Python Environment to export the built Python packages
 
+When CMake is configured with `-DMLIR_ENABLE_BINDINGS_PYTHON=ON`, the python packages will typically be located in either:
+
+1. `./build/tools/torch-mlir/python_packages/`  if doing an in-tree build.
+2. `./build/python_packages/` if doing an out-of-tree build.
+
+For the following sections, let `python_pkg_dir` represent whichever of the above is relevant for your build setup. On Linux and macOS, you can run `./build_tools/write_env_file.sh` to generate a file `./.env` in your root source directory with the correct `PYTHONPATH`.
+
 #### Linux and macOS
 
+To get the base `PYTHONPATH`, run:
+
 ```shell
-export PYTHONPATH=`pwd`/build/python_packages/torch_mlir:`pwd`/test/python/fx_importer
+./build_tools/write_env_file.sh
+source ./.env && export PYTHONPATH
+```
+
+To run fx_importer tests, you can append the following:
+
+```
+export PYTHONPATH="${PYTHONPATH}":/test/python/fx_importer"
 ```
 
 #### Windows PowerShell
 
+To get the base `PYTHONPATH`, identify your `python_pkg_dir` from above and set this variable in your environment:
+
+```shell
+$env:PYTHONPATH = "<python_pkg_dir>/torch-mlir"
+```
+
+To run fx_importer tests, you can append the following:
+
 ```shell
-$env:PYTHONPATH = "$PWD/build/tools/torch-mlir/python_packages/torch_mlir;$PWD/test/python/fx_importer"
+$env:PYTHONPATH += ";$PWD/test/python/fx_importer"
 ```
 
 ### Testing MLIR output in various dialects
 
 To test the MLIR output to torch dialect, you can use `test/python/fx_importer/basic_test.py`.
 
 Make sure you have activated the virtualenv and set the `PYTHONPATH` above
-(if running on Windows, modify the environment variable as shown above):
-
-```shell
-source mlir_venv/bin/activate
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/test/python/fx_importer
-python test/python/fx_importer/basic_test.py
-```
+(if running on Windows, modify the environment variable as shown above).
 
 This will display the basic example in TORCH dialect.
 
@@ -226,10 +244,10 @@ using torchscript with the example `projects/pt1/examples/torchscript_resnet18_a
 This path doesn't give access to the current generation work that is being driven via the fx_importer
 and may lead to errors.
 
-Same as above, but with different python path and example:
+The base `PYTHONPATH` should be set as above, then the example can be run with the following command (similar on Windows):
 
 ```shell
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/projects/pt1/examples
+export PYTHONPATH="${PYTHONPATH}:$PWD/projects/pt1/examples"
 python projects/pt1/examples/torchscript_resnet18_all_output_types.py
 ```
 

lib/Conversion/TorchToLinalg/Linear.cpp

@@ -1391,9 +1391,13 @@ class ConvertAtenConvolutionOp : public OpConversionPattern<AtenConvolutionOp> {
       return success();
     }
 
-    if (numSpatialDims != 2)
+    if (numSpatialDims != 2 && numSpatialDims != 3)
       return rewriter.notifyMatchFailure(
-          op, "unimplemented: only 1D and 2D grouped convolution supported");
+          op, "unimplemented: only 2D and 3D grouped convolution supported");
+    if (numSpatialDims == 3 && inputZp) {
+      return rewriter.notifyMatchFailure(
+          op, "unimplemented: quantized 3D grouped convolution not supported");
+    }
 
     // Grouped case, use the grouped conv linalg op
     auto expandGroups = [&](Value tensor, size_t dim) {
@@ -1435,21 +1439,101 @@ class ConvertAtenConvolutionOp : public OpConversionPattern<AtenConvolutionOp> {
     weight = transposed ? weight : expandWeight(weight);
     auto expandOutputTensor = expandGroups(outputTensor, 1);
 
-    // TODO: add 1D and 3D case
-    if (!inputZp) {
-      conv = rewriter
-                 .create<linalg::Conv2DNgchwGfchwOp>(
-                     loc, expandOutputTensor.getResultType(),
-                     ValueRange{paddedInputExpanded, weight},
-                     expandOutputTensor.getResult(), stridesAttr, dilationAttr)
-                 .getResult(0);
-    } else {
-      conv = rewriter
-                 .create<linalg::Conv2DNgchwGfchwQOp>(
-                     loc, expandOutputTensor.getResultType(),
-                     ValueRange{paddedInputExpanded, weight, inputZp, weightZp},
-                     expandOutputTensor.getResult(), stridesAttr, dilationAttr)
-                 .getResult(0);
+    if (numSpatialDims == 2) {
+      // 2D grouped convolution
+      if (!inputZp) {
+        conv =
+            rewriter
+                .create<linalg::Conv2DNgchwGfchwOp>(
+                    loc, expandOutputTensor.getResultType(),
+                    ValueRange{paddedInputExpanded, weight},
+                    expandOutputTensor.getResult(), stridesAttr, dilationAttr)
+                .getResult(0);
+      } else {
+        conv =
+            rewriter
+                .create<linalg::Conv2DNgchwGfchwQOp>(
+                    loc, expandOutputTensor.getResultType(),
+                    ValueRange{paddedInputExpanded, weight, inputZp, weightZp},
+                    expandOutputTensor.getResult(), stridesAttr, dilationAttr)
+                .getResult(0);
+      }
+    } else if (numSpatialDims == 3) {
+      // MLIR does not have a named 3D grouped convolution op, so we use
+      // linalg.generic instead.
+      AffineExpr d0, d1, d2, d3, d4, d5, d6, d7, d8, d9;
+      bindDims(context, d0, d1, d2, d3, d4, d5, d6, d7, d8, d9);
+
+      SmallVector<AffineExpr> inputExprs = {
+          d0,                                        // N
+          d1,                                        // G
+          d6,                                        // C/G
+          d3 * strideInts[0] + d7 * dilationInts[0], // D
+          d4 * strideInts[1] + d8 * dilationInts[1], // H
+          d5 * strideInts[2] + d9 * dilationInts[2]  // W
+      };
+
+      SmallVector<AffineExpr> weightExprs = {
+          d1, // G
+          d2, // F/G
+          d6, // C/G
+          d7, // KD
+          d8, // KH
+          d9  // KW
+      };
+
+      SmallVector<AffineExpr> outputExprs = {
+          d0, // N
+          d1, // G
+          d2, // F/G
+          d3, // OD
+          d4, // OH
+          d5, // OW
+      };
+
+      SmallVector<AffineMap> indexingMaps = {
+          AffineMap::get(10, 0, inputExprs, rewriter.getContext()),
+          AffineMap::get(10, 0, weightExprs, rewriter.getContext()),
+          AffineMap::get(10, 0, outputExprs, rewriter.getContext())};
+
+      SmallVector<utils::IteratorType> iteratorTypes = {
+          utils::IteratorType::parallel,  // N
+          utils::IteratorType::parallel,  // G
+          utils::IteratorType::parallel,  // F/G
+          utils::IteratorType::parallel,  // OD
+          utils::IteratorType::parallel,  // OH
+          utils::IteratorType::parallel,  // OW
+          utils::IteratorType::reduction, // C/G
+          utils::IteratorType::reduction, // KD
+          utils::IteratorType::reduction, // KH
+          utils::IteratorType::reduction  // KW
+      };
+
+      conv =
+          rewriter
+              .create<linalg::GenericOp>(
+                  loc, expandOutputTensor.getResultType(),
+                  ValueRange{paddedInputExpanded, weight},
+                  expandOutputTensor.getResult(), indexingMaps, iteratorTypes,
+                  [&](OpBuilder &b, Location loc, ValueRange args) {
+                    Value input = args[0];
+                    Value weight = args[1];
+                    Value output = args[2];
+
+                    // Convert input and weight to accumulator type if needed
+                    Type accType = output.getType();
+                    if (input.getType() != accType) {
+                      input = b.create<arith::ExtFOp>(loc, accType, input);
+                    }
+                    if (weight.getType() != accType) {
+                      weight = b.create<arith::ExtFOp>(loc, accType, weight);
+                    }
+
+                    Value mul = b.create<arith::MulFOp>(loc, input, weight);
+                    Value add = b.create<arith::AddFOp>(loc, mul, output);
+                    b.create<linalg::YieldOp>(loc, add);
+                  })
+              .getResult(0);
     }
     conv = rewriter.create<tensor::CollapseShapeOp>(
         loc, outputTensor.getType(), conv,

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -6435,6 +6435,11 @@ class ConvertAtenConstPatternOp : public OpConversionPattern<AtenOpT> {
     for (auto s : shape)
       size *= s;
 
+    if (size == 0) {
+      return rewriter.notifyMatchFailure(
+          op, "Shape must not have a dimension of size zero");
+    }
+
     SmallVector<int32_t> values(size, fillVal);
     auto constOp =
         tosa::getConstTensor<int32_t>(rewriter, op, values, shape).value();

lib/Dialect/Torch/Transforms/Passes.cpp

@@ -70,12 +70,16 @@ void mlir::torch::Torch::createTorchScriptModuleToTorchBackendPipeline(
 
 void mlir::torch::Torch::createTorchDynamoExportToTorchBackendPipeline(
     OpPassManager &pm, const TorchLoweringPipelineOptions &options) {
+  // Inline func.call operations created by higher-order ops like while_loop
+  // to conform to the linalg-on-tensors backend contract.
+  pm.addPass(createInlinerPass());
   pm.addNestedPass<func::FuncOp>(
       createReduceOpVariantsPass(options.extraLibrary));
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   if (options.decompose) {
     pm.addNestedPass<func::FuncOp>(
         Torch::createDecomposeComplexOpsPass(options.backendLegalOps));
+    pm.addNestedPass<func::FuncOp>(Torch::createRecomposeComplexOpsPass());
     pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   }
 }

lib/RefBackend/RefBackend.cpp

@@ -109,7 +109,7 @@ static std::string getConsumeReturnFunctionNameForReturnTypes(TypeRange types) {
     tokens.push_back(getTypeToken(type));
 
   return std::accumulate(tokens.begin(), tokens.end(), std::string(),
-                         [](std::string &a, std::string &b) {
+                         [](std::string a, std::string b) {
                            return a.empty() ? b : (a + "_" + b);
                          });
 }