[ARM64] Vectorized math approximations, elementwise fusion and buffer deallocation for activations

qcom_arm64_backend/bin/arm64-mlir-opt.cpp

@@ -93,6 +93,7 @@ int main(int argc, char **argv) {
   triton::registerTritonARM64LinalgVectorization();
   triton::registerTritonARM64VectorLowering();
   triton::registerTritonARM64ProfPass();
+  triton::registerTritonARM64MathApproximation();
   triton::registerTritonARM64RVO();
   triton::registerTritonARM64PuntBuffer();
   triton::registerTritonARM64BubbleUpExtractSlice();

qcom_arm64_backend/include/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.h

@@ -23,6 +23,7 @@ namespace mlir {
 namespace triton {
 #define GEN_PASS_DECL_TRITONARM64PROFPASS
 #define GEN_PASS_DECL_TRITONARM64LOWERCONSTANTSSEPARATELY
+#define GEN_PASS_DECL_TRITONARM64MATHAPPROXIMATION
 #include "LowerTritonLinalgToLLVM/Passes.h.inc"
 
 std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
@@ -35,6 +36,10 @@ std::unique_ptr<OperationPass<func::FuncOp>> createTritonARM64ProfPass(
 std::unique_ptr<OperationPass<ModuleOp>> createLowerConstantsSeparatelyPass(
     std::shared_ptr<std::vector<ModuleOp>> sink = nullptr);
 
+std::unique_ptr<OperationPass<ModuleOp>> createTritonARM64MathApproximationPass(
+    const TritonARM64MathApproximationOptions &options =
+        TritonARM64MathApproximationOptions());
+
 } // namespace triton
 } // namespace mlir
 

qcom_arm64_backend/include/LowerTritonLinalgToLLVM/Passes.td

@@ -40,4 +40,22 @@ def TritonARM64LowerConstantsSeparately :
   let constructor = "mlir::triton::createLowerConstantsSeparatelyPass()";
 }
 
+def TritonARM64MathApproximation :
+    Pass<"triton-arm64-math-approximation", "mlir::ModuleOp"> {
+  let summary = "Expand selected math ops into polynomial approximations";
+  let description = [{
+    Rewrites math ops into arith/vector ops (upstream MLIR polynomial
+    approximations), so they stay vectorized instead of becoming per-lane
+    libm calls. f16/bf16 ops are computed in f32. Must run before
+    convert-math-to-llvm; math.erf has no LLVM lowering otherwise.
+  }];
+
+  let options = [
+  ListOption<"ops", "ops", "std::string",
+             "Math op names to approximate (default: math.erf,math.exp,math.tanh)">,
+  ];
+
+  let constructor = "mlir::triton::createTritonARM64MathApproximationPass()";
+}
+
 #endif

qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/ARM64MathApproximation.cpp

@@ -0,0 +1,75 @@
+//===- ARM64MathApproximation.cpp -----------------------------------------===//
+//
+// Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+// SPDX-License-Identifier: BSD-3-Clause.
+// For more license information:
+//   https://github.com/qualcomm/cpu-mlir/LICENSE.txt
+//
+//===----------------------------------------------------------------------===//
+// Expands selected math ops into upstream MLIR polynomial approximations.
+// Without this, math.erf cannot be lowered to LLVM at all, and math.exp on
+// scalable vectors becomes llvm.exp.nxv*, which the AArch64 backend
+// scalarizes into one libm call per lane.
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/ADT/StringSet.h"
+
+#include "LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.h"
+
+using namespace mlir;
+using namespace triton;
+
+#define GEN_PASS_DEF_TRITONARM64MATHAPPROXIMATION
+#include "LowerTritonLinalgToLLVM/Passes.h.inc"
+
+namespace {
+
+// Ops used by activation functions. Other ops keep their libm lowering until
+// their approximation accuracy has been validated on this backend.
+constexpr StringLiteral kDefaultOps[] = {"math.erf", "math.exp", "math.tanh"};
+
+struct TritonARM64MathApproximationPass
+    : public ::impl::TritonARM64MathApproximationBase<
+          TritonARM64MathApproximationPass> {
+  using TritonARM64MathApproximationBase::TritonARM64MathApproximationBase;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<arith::ArithDialect, math::MathDialect,
+                    vector::VectorDialect>();
+  }
+
+  void runOnOperation() override {
+    llvm::StringSet<> selected;
+    if (ops.empty())
+      selected.insert(std::begin(kDefaultOps), std::end(kDefaultOps));
+    else
+      selected.insert(ops.begin(), ops.end());
+    auto predicate = [&](StringRef name) { return selected.contains(name); };
+
+    RewritePatternSet patterns(&getContext());
+    // f16/bf16 -> f32, then the f32 polynomial.
+    populateMathF32ExpansionPatterns(patterns, predicate);
+    populateMathPolynomialApproximationPatterns(patterns, predicate);
+    // Not an approximation, so always on: math.fpowi on vectors has no LLVM
+    // lowering (llvm.intr.powi takes a scalar exponent). A constant power
+    // becomes multiplies (torch-mlir emits x^3 for GELU tanh), any other powf.
+    math::populateExpansionPatterns(patterns, {"fpowi"});
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>>
+triton::createTritonARM64MathApproximationPass(
+    const TritonARM64MathApproximationOptions &options) {
+  return std::make_unique<TritonARM64MathApproximationPass>(options);
+}

qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_mlir_conversion_library(LowerTritonLinalgToLLVM
+        ARM64MathApproximation.cpp
         ARM64ProfInstrumentation.cpp
         LowerTmTensor.cpp
 	LowerTritonLinalgToLLVM.cpp

qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.cpp

@@ -76,6 +76,9 @@ void addTritonLinalgToLLVMPasses(
   pm.addPass(createSCFToControlFlowPass());
   pm.addPass(createConvertControlFlowToLLVMPass());
 
+  // Must run before ConvertMathToLLVM: math.erf has no LLVM lowering, and
+  // vector math.exp would otherwise be scalarized into libm calls.
+  pm.addPass(createTritonARM64MathApproximationPass());
   ConvertVectorToLLVMPassOptions options{};
   options.armSVE = (vec_ext == SVE || vec_ext == SME);
   pm.addPass(createConvertVectorToLLVMPass(options));

qcom_arm64_backend/lib/OptimizeTritonShared/CMakeLists.txt

@@ -16,6 +16,7 @@ add_mlir_conversion_library(OptimizeTritonShared
 
         LINK_LIBS PUBLIC
         MLIRBufferizationDialect
+        MLIRBufferizationPipelines
         MLIRBufferizationTransforms
         MLIRIR
         MLIRMemRefDialect

qcom_arm64_backend/lib/OptimizeTritonShared/OptimizeTritonShared.cpp

@@ -13,6 +13,7 @@
 #include "mlir/IR/Dialect.h"
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h"
 #include "mlir/Dialect/ArmSVE/Transforms/Passes.h"
+#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -93,6 +94,11 @@ ModuleOp arm64OptimizeTritonShared(ModuleOp mod, arm64VectorExtension vec_ext,
 
 void addOptimizeTritonSharedPasses(mlir::OpPassManager &pm,
     arm64VectorExtension vec_ext) {
+  // triton-shared emits one linalg.generic per elementwise op. Fusing them
+  // avoids temporaries and keeps long chains (e.g. GELU) within the greedy
+  // rewrite limit of bubble-up-extract-slice.
+  pm.addPass(createLinalgElementwiseOpFusionPass());
+  pm.addPass(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(createTritonARM64PuntBufferPass());
   pm.addNestedPass<func::FuncOp>(createTritonARM64RVOPass());
   pm.addPass(createCanonicalizerPass());
@@ -104,6 +110,14 @@ void addOptimizeTritonSharedPasses(mlir::OpPassManager &pm,
   pm.addPass(createTritonARM64LinalgVectorizationPass({vec_ext}));
   pm.addPass(bufferization::createEmptyTensorEliminationPass());
   pm.addPass(createTritonARM64VectorLoweringPass({vec_ext}));
+  // The vector lowering pass bufferizes, but bufferization never frees what it
+  // allocates, so every program instance leaked its temporaries. This must run
+  // while loops are still structured (SME lowers them to cf below).
+  bufferization::buildBufferDeallocationPipeline(pm);
+  // Deallocs whose aliasing is only known at runtime (e.g. buffers returned
+  // from loops) call a generated @dealloc_helper. Inline it: the driver
+  // expects a single function per kernel module.
+  pm.addPass(createInlinerPass());
 
   if (vec_ext == SME)
     addOptimizationPassesForSME(pm);

qcom_arm64_backend/python/examples/conftest.py

@@ -92,6 +92,7 @@ def set_target(config: pytest.Config):
     "test_layernorm.py",
 
     # QC-internal tests
+    "test_activations.py",
     "test_flash_attention.py"
 }
 

qcom_arm64_backend/python/examples/test_activations.py

@@ -0,0 +1,106 @@
+# ===- test_activations.py --------------------------------------------------===
+#
+#  Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+#  SPDX-License-Identifier: BSD-3-Clause.
+#  For more license information:
+#    https://github.com/qualcomm/cpu-mlir/LICENSE.txt
+#
+# ===------------------------------------------------------------------------===
+
+import pytest
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+BLOCK = 1024
+SQRT_2_OVER_PI = tl.constexpr(0.7978845608028654)
+INV_SQRT_2 = tl.constexpr(0.7071067811865476)
+
+
+@triton.jit
+def _relu(x):
+    return tl.maximum(x, 0.0)
+
+
+@triton.jit
+def _leaky_relu(x):
+    return tl.where(x >= 0.0, x, 0.01 * x)
+
+
+@triton.jit
+def _sigmoid(x):
+    return tl.sigmoid(x)
+
+
+@triton.jit
+def _silu(x):
+    return x * tl.sigmoid(x)
+
+
+@triton.jit
+def _tanh(x):
+    # No tl tanh on CPU; this form does not overflow for large |x|.
+    return 2.0 * tl.sigmoid(2.0 * x) - 1.0
+
+
+@triton.jit
+def _gelu_erf(x):
+    return 0.5 * x * (1.0 + tl.math.erf(x * INV_SQRT_2))
+
+
+@triton.jit
+def _gelu_tanh(x):
+    inner = SQRT_2_OVER_PI * (x + 0.044715 * x * x * x)
+    return 0.5 * x * (1.0 + (2.0 * tl.sigmoid(2.0 * inner) - 1.0))
+
+
+def _make_kernel(body):
+    @triton.jit
+    def kernel(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
+        offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+        mask = offs < n
+        x = tl.load(x_ptr + offs, mask=mask)
+        tl.store(out_ptr + offs, body(x), mask=mask)
+    return kernel
+
+
+# name: (triton body, torch reference, atol)
+ACTIVATIONS = {
+    "relu": (_relu, torch.relu, 0.0),
+    "leaky_relu": (_leaky_relu, lambda x: F.leaky_relu(x, 0.01), 1e-7),
+    "sigmoid": (_sigmoid, torch.sigmoid, 1e-6),
+    "silu": (_silu, F.silu, 1e-5),
+    "tanh": (_tanh, torch.tanh, 1e-6),
+    "gelu_erf": (_gelu_erf, F.gelu, 1e-5),
+    "gelu_tanh": (_gelu_tanh, lambda x: F.gelu(x, approximate="tanh"), 1e-5),
+}
+KERNELS = {name: _make_kernel(body) for name, (body, _, _) in ACTIVATIONS.items()}
+
+
+def run_activation(name, x):
+    out = torch.empty_like(x)
+    n = x.numel()
+    KERNELS[name][(triton.cdiv(n, BLOCK),)](x, out, n, BLOCK=BLOCK)
+    return out
+
+
+@pytest.mark.parametrize("name", ACTIVATIONS.keys())
+@pytest.mark.parametrize("n", [1, BLOCK, 3 * BLOCK + 7])
+def test_activation(name, n, device):
+    torch.manual_seed(0)
+    x = torch.randn(n, device=device) * 4
+    _, ref, atol = ACTIVATIONS[name]
+    torch.testing.assert_close(run_activation(name, x), ref(x), atol=atol, rtol=1e-5)
+
+
+@pytest.mark.parametrize("name", ACTIVATIONS.keys())
+def test_activation_edge_values(name, device):
+    # Saturation regions and signed zeros, where approximations usually break.
+    x = torch.tensor([0.0, -0.0, 1e-7, -1e-7, 5.0, -5.0, 10.0, -10.0, 30.0, -30.0,
+                      88.0, -88.0, 1e4, -1e4], device=device)
+    if name == "gelu_tanh":
+        x = x[x.abs() < 1e3]  # the cubic term overflows to inf in fp32 for huge |x|
+    _, ref, atol = ACTIVATIONS[name]
+    torch.testing.assert_close(run_activation(name, x), ref(x), atol=atol, rtol=1e-5)

qcom_arm64_backend/python/torch-mlir/test_activations_torch.py

@@ -0,0 +1,64 @@
+# ===- test_activations_torch.py --------------------------------------------===
+#
+#  Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+#  SPDX-License-Identifier: BSD-3-Clause.
+#  For more license information:
+#    https://github.com/qualcomm/cpu-mlir/LICENSE.txt
+#
+# ===------------------------------------------------------------------------===
+
+# The activations of examples/test_activations.py, compiled from PyTorch through
+# torch-mlir instead of Triton.
+
+import torch
+import torch.nn as nn
+import utils
+
+
+def gelu_tanh_as_imported(x):
+    # torch-mlir lowers GELU(approximate="tanh") with sqrt(2/pi) = 0.797724
+    # instead of 0.7978846, which alone moves outputs by up to 5e-5. Check the
+    # backend against the formula it is given, not against PyTorch.
+    return 0.5 * x * (1.0 + torch.tanh(0.797724 * (x + 0.044715 * x * x * x)))
+
+
+# name: (module, atol, reference; None = the module itself)
+ACTIVATIONS = {
+    "relu": (nn.ReLU, 0.0, None),
+    "leaky_relu": (lambda: nn.LeakyReLU(0.01), 1e-7, None),
+    "sigmoid": (nn.Sigmoid, 1e-6, None),
+    "silu": (nn.SiLU, 1e-5, None),
+    "tanh": (nn.Tanh, 1e-6, None),
+    "gelu_erf": (nn.GELU, 1e-5, None),
+    "gelu_tanh": (lambda: nn.GELU(approximate="tanh"), 1e-5, gelu_tanh_as_imported),
+}
+
+
+def make_input():
+    torch.manual_seed(0)
+    # A size that is not a multiple of any vector length, plus the saturation
+    # regions and signed zeros, where approximations usually break.
+    edges = torch.tensor([0.0, -0.0, 1e-7, -1e-7, 5.0, -5.0, 10.0, -10.0, 30.0, -30.0,
+                          88.0, -88.0, 1e4, -1e4])
+    return torch.cat([torch.randn(3 * 1024 + 7) * 4, edges])
+
+
+def test_activations_torch(options={}):
+    failed = []
+    for name, (make_module, atol, reference) in ACTIVATIONS.items():
+        model = make_module()
+        # Each activation needs its own function name (and so its own files).
+        model.__class__ = type(f"Activation_{name}", (model.__class__,), {})
+        print(f"== {name}")
+        x = make_input()
+        manager = utils.ModelManager(model, x)
+        manager.create_linalg_module()
+        manager.write_bytecode_to_file()
+        if not manager.execute_and_compare(rtol=1e-5, atol=atol, options=options,
+                                           reference=reference(x) if reference else None):
+            failed.append(name)
+    assert not failed, f"mismatch: {failed}"
+
+
+if __name__ == "__main__":
+    test_activations_torch(utils.get_vector_ext_options_from_args())

qcom_arm64_backend/python/torch-mlir/utils.py

@@ -55,9 +55,10 @@ def write_bytecode_to_file(self):
         with open(self.linalg_filename, "wb") as f:
             f.write(bytecode)
 
-    def execute_and_compare(self, rtol=1e-05, atol=1e-08, options={}, separate_constants=False):
-        # Generate reference output
-        reference = self.model(*self.inputs)
+    def execute_and_compare(self, rtol=1e-05, atol=1e-08, options={}, separate_constants=False, reference=None):
+        # Generate reference output, unless the caller gives one
+        if reference is None:
+            reference = self.model(*self.inputs)
 
         # ARM64 execution
         output = TorchMLIRLauncher().run_torch_mlir(
@@ -70,12 +71,18 @@ def execute_and_compare(self, rtol=1e-05, atol=1e-08, options={}, separate_const
                 print("\nReference output:\n", reference.shape, reference)
                 print("\narm64 output:\n", output[0].shape, output[0])
 
-            print("Match: ", torch.allclose(output[0], reference, rtol, atol))
+            match = torch.allclose(output[0], reference, rtol, atol)
+            print("Match: ", match)
+            return match
 
         elif len(output) > 1:
+            matches = []
             for index in range(len(output)):
                 if self.dump_outputs:
                     print("\nReference output:\n", reference[index])
                     print("\narm64 output:\n", output[index])
 
-                print("Match: ", torch.allclose(output[index], reference[index], rtol, atol))
+                match = torch.allclose(output[index], reference[index], rtol, atol)
+                print("Match: ", match)
+                matches.append(match)
+            return all(matches)