diff --git a/qcom_arm64_backend/bin/arm64-mlir-opt.cpp b/qcom_arm64_backend/bin/arm64-mlir-opt.cpp index 22b40dc..391223d 100644 --- a/qcom_arm64_backend/bin/arm64-mlir-opt.cpp +++ b/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(); diff --git a/qcom_arm64_backend/include/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.h b/qcom_arm64_backend/include/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.h index 2f7a55a..9f88dcd 100644 --- a/qcom_arm64_backend/include/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.h +++ b/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> @@ -35,6 +36,10 @@ std::unique_ptr> createTritonARM64ProfPass( std::unique_ptr> createLowerConstantsSeparatelyPass( std::shared_ptr> sink = nullptr); +std::unique_ptr> createTritonARM64MathApproximationPass( + const TritonARM64MathApproximationOptions &options = + TritonARM64MathApproximationOptions()); + } // namespace triton } // namespace mlir diff --git a/qcom_arm64_backend/include/LowerTritonLinalgToLLVM/Passes.td b/qcom_arm64_backend/include/LowerTritonLinalgToLLVM/Passes.td index 15e3722..fbf8550 100644 --- a/qcom_arm64_backend/include/LowerTritonLinalgToLLVM/Passes.td +++ b/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 diff --git a/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/ARM64MathApproximation.cpp b/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/ARM64MathApproximation.cpp new file mode 100644 index 0000000..e023336 --- /dev/null +++ b/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 ®istry) const override { + registry.insert(); + } + + 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> +triton::createTritonARM64MathApproximationPass( + const TritonARM64MathApproximationOptions &options) { + return std::make_unique(options); +} diff --git a/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/CMakeLists.txt b/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/CMakeLists.txt index 1a08b8b..3310b77 100644 --- a/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/CMakeLists.txt +++ b/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/CMakeLists.txt @@ -1,4 +1,5 @@ add_mlir_conversion_library(LowerTritonLinalgToLLVM + ARM64MathApproximation.cpp ARM64ProfInstrumentation.cpp LowerTmTensor.cpp LowerTritonLinalgToLLVM.cpp diff --git a/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.cpp b/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.cpp index 69e46a4..8811234 100644 --- a/qcom_arm64_backend/lib/LowerTritonLinalgToLLVM/LowerTritonLinalgToLLVM.cpp +++ b/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)); diff --git a/qcom_arm64_backend/lib/OptimizeTritonShared/CMakeLists.txt b/qcom_arm64_backend/lib/OptimizeTritonShared/CMakeLists.txt index 329b42d..83c4611 100644 --- a/qcom_arm64_backend/lib/OptimizeTritonShared/CMakeLists.txt +++ b/qcom_arm64_backend/lib/OptimizeTritonShared/CMakeLists.txt @@ -16,6 +16,7 @@ add_mlir_conversion_library(OptimizeTritonShared LINK_LIBS PUBLIC MLIRBufferizationDialect + MLIRBufferizationPipelines MLIRBufferizationTransforms MLIRIR MLIRMemRefDialect diff --git a/qcom_arm64_backend/lib/OptimizeTritonShared/OptimizeTritonShared.cpp b/qcom_arm64_backend/lib/OptimizeTritonShared/OptimizeTritonShared.cpp index d8f0439..1bff92c 100644 --- a/qcom_arm64_backend/lib/OptimizeTritonShared/OptimizeTritonShared.cpp +++ b/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(createTritonARM64PuntBufferPass()); pm.addNestedPass(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); diff --git a/qcom_arm64_backend/python/examples/conftest.py b/qcom_arm64_backend/python/examples/conftest.py index 0df630c..7d63856 100644 --- a/qcom_arm64_backend/python/examples/conftest.py +++ b/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" } diff --git a/qcom_arm64_backend/python/examples/test_activations.py b/qcom_arm64_backend/python/examples/test_activations.py new file mode 100644 index 0000000..5162edf --- /dev/null +++ b/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) diff --git a/qcom_arm64_backend/python/torch-mlir/test_activations_torch.py b/qcom_arm64_backend/python/torch-mlir/test_activations_torch.py new file mode 100644 index 0000000..1a019cc --- /dev/null +++ b/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()) diff --git a/qcom_arm64_backend/python/torch-mlir/utils.py b/qcom_arm64_backend/python/torch-mlir/utils.py index fe83916..8e45e9c 100644 --- a/qcom_arm64_backend/python/torch-mlir/utils.py +++ b/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) diff --git a/qcom_arm64_backend/test/Transforms/arm64-math-approximation.mlir b/qcom_arm64_backend/test/Transforms/arm64-math-approximation.mlir new file mode 100644 index 0000000..85f59b1 --- /dev/null +++ b/qcom_arm64_backend/test/Transforms/arm64-math-approximation.mlir @@ -0,0 +1,78 @@ +// RUN: arm64-mlir-opt %s -triton-arm64-math-approximation | FileCheck %s +// RUN: arm64-mlir-opt %s -triton-arm64-math-approximation="ops=math.log" | FileCheck %s --check-prefix=OPS + +// Default set (erf, exp, tanh): expanded to arith ops, others untouched. + +// CHECK-LABEL: func.func @erf_scalable +// CHECK-NOT: math.erf +// CHECK: arith.select +// CHECK: return {{.*}} : vector<[4]xf32> +// OPS-LABEL: func.func @erf_scalable +// OPS: math.erf +func.func @erf_scalable(%x: vector<[4]xf32>) -> vector<[4]xf32> { + %0 = math.erf %x : vector<[4]xf32> + return %0 : vector<[4]xf32> +} + +// CHECK-LABEL: func.func @exp_scalable +// CHECK-NOT: math.exp +// CHECK: arith.bitcast {{.*}} : vector<[4]xi32> to vector<[4]xf32> +// CHECK: return {{.*}} : vector<[4]xf32> +func.func @exp_scalable(%x: vector<[4]xf32>) -> vector<[4]xf32> { + %0 = math.exp %x : vector<[4]xf32> + return %0 : vector<[4]xf32> +} + +// CHECK-LABEL: func.func @tanh_f32 +// CHECK-NOT: math.tanh +// CHECK: return {{.*}} : f32 +func.func @tanh_f32(%x: f32) -> f32 { + %0 = math.tanh %x : f32 + return %0 : f32 +} + +// f16 is computed in f32. +// CHECK-LABEL: func.func @exp_f16 +// CHECK: arith.extf {{.*}} : vector<8xf16> to vector<8xf32> +// CHECK-NOT: math.exp +// CHECK: arith.truncf {{.*}} : vector<8xf32> to vector<8xf16> +func.func @exp_f16(%x: vector<8xf16>) -> vector<8xf16> { + %0 = math.exp %x : vector<8xf16> + return %0 : vector<8xf16> +} + +// Not in the default set: keeps its libm lowering. +// CHECK-LABEL: func.func @log_f32 +// CHECK: math.log +// OPS-LABEL: func.func @log_f32 +// OPS-NOT: math.log +func.func @log_f32(%x: f32) -> f32 { + %0 = math.log %x : f32 + return %0 : f32 +} + +// fpowi is expanded whatever ops= says: llvm.intr.powi needs a scalar exponent, +// so a vector fpowi cannot be lowered to LLVM. x^3 (torch-mlir's GELU tanh) +// becomes multiplies. +// CHECK-LABEL: func.func @fpowi_const_scalable +// CHECK-NOT: math.fpowi +// CHECK: arith.mulf +// CHECK: return {{.*}} : vector<[4]xf32> +// OPS-LABEL: func.func @fpowi_const_scalable +// OPS-NOT: math.fpowi +// OPS: return +func.func @fpowi_const_scalable(%x: vector<[4]xf32>) -> vector<[4]xf32> { + %c3 = arith.constant dense<3> : vector<[4]xi64> + %0 = math.fpowi %x, %c3 : vector<[4]xf32>, vector<[4]xi64> + return %0 : vector<[4]xf32> +} + +// A non-constant power becomes powf, which has a vector lowering. +// CHECK-LABEL: func.func @fpowi_var +// CHECK-NOT: math.fpowi +// CHECK: arith.sitofp +// CHECK: math.powf +func.func @fpowi_var(%x: vector<[4]xf32>, %n: vector<[4]xi64>) -> vector<[4]xf32> { + %0 = math.fpowi %x, %n : vector<[4]xf32>, vector<[4]xi64> + return %0 : vector<[4]xf32> +} diff --git a/qcom_arm64_backend/test/Transforms/optimize-triton-shared-buffer-deallocation.mlir b/qcom_arm64_backend/test/Transforms/optimize-triton-shared-buffer-deallocation.mlir new file mode 100644 index 0000000..78d81af --- /dev/null +++ b/qcom_arm64_backend/test/Transforms/optimize-triton-shared-buffer-deallocation.mlir @@ -0,0 +1,67 @@ +// RUN: arm64-mlir-opt %s -split-input-file --vector-ext=sve -optimize-triton-shared | FileCheck %s +// RUN: arm64-mlir-opt %s -split-input-file --vector-ext=sme -optimize-triton-shared | FileCheck %s + +// A temporary used by two consumers: elementwise fusion keeps it, so it is +// bufferized into a buffer that must be freed before returning, along with the +// two results. With SME the loops are lowered to cf, so deallocation has to run +// before that. + +// CHECK-LABEL: func.func @two_uses +// CHECK-COUNT-3: memref.alloc +// CHECK-COUNT-3: memref.dealloc +// CHECK-NOT: memref.alloc +// CHECK: return + +#map = affine_map<(d0) -> (d0)> +func.func @two_uses(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>, %arg2: memref<1024xf32>) { + %x = bufferization.to_tensor %arg0 restrict : memref<1024xf32> to tensor<1024xf32> + %e0 = tensor.empty() : tensor<1024xf32> + %e = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%x : tensor<1024xf32>) outs(%e0 : tensor<1024xf32>) { + ^bb0(%in: f32, %out: f32): + %0 = math.exp %in : f32 + linalg.yield %0 : f32 + } -> tensor<1024xf32> + %a0 = tensor.empty() : tensor<1024xf32> + %a = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%e : tensor<1024xf32>) outs(%a0 : tensor<1024xf32>) { + ^bb0(%in: f32, %out: f32): + %0 = arith.negf %in : f32 + linalg.yield %0 : f32 + } -> tensor<1024xf32> + %b0 = tensor.empty() : tensor<1024xf32> + %b = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%e : tensor<1024xf32>) outs(%b0 : tensor<1024xf32>) { + ^bb0(%in: f32, %out: f32): + %0 = arith.mulf %in, %in : f32 + linalg.yield %0 : f32 + } -> tensor<1024xf32> + bufferization.materialize_in_destination %a in writable %arg1 : (tensor<1024xf32>, memref<1024xf32>) -> () + bufferization.materialize_in_destination %b in writable %arg2 : (tensor<1024xf32>, memref<1024xf32>) -> () + return +} + +// ----- + +// A buffer returned from a loop: whether to free it is decided at runtime by a +// generated helper, which must be inlined into the kernel. After the loop both +// %init and the loop result %r must be freed; the checks follow %r from its use +// in the copy to %out to its own dealloc. + +// CHECK-LABEL: func.func @loop_buffer +// CHECK-NOT: dealloc_helper +// CHECK: %[[INIT:.+]] = memref.alloc() : memref<16xf32> +// CHECK: memref.copy %[[R:[^,]+]], %arg1 : +// CHECK: %[[R_BASE:[^,]+]], {{.*}} = memref.extract_strided_metadata %[[R]] : +// CHECK-DAG: memref.dealloc %[[INIT]] : +// CHECK-DAG: memref.dealloc %[[R_BASE]] : +// CHECK-NOT: dealloc_helper +func.func @loop_buffer(%n: index, %out: memref<16xf32>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %init = memref.alloc() : memref<16xf32> + %r = scf.for %i = %c0 to %n step %c1 iter_args(%acc = %init) -> (memref<16xf32>) { + %new = memref.alloc() : memref<16xf32> + memref.copy %acc, %new : memref<16xf32> to memref<16xf32> + scf.yield %new : memref<16xf32> + } + memref.copy %r, %out : memref<16xf32> to memref<16xf32> + return +} diff --git a/qcom_arm64_backend/test/Transforms/optimize-triton-shared-elementwise-fusion.mlir b/qcom_arm64_backend/test/Transforms/optimize-triton-shared-elementwise-fusion.mlir new file mode 100644 index 0000000..ba83be9 --- /dev/null +++ b/qcom_arm64_backend/test/Transforms/optimize-triton-shared-elementwise-fusion.mlir @@ -0,0 +1,138 @@ +// RUN: arm64-mlir-opt %s --vector-ext=sve -optimize-triton-shared | FileCheck %s + +// GELU (tanh form) as emitted by triton-shared: one linalg.generic per op. +// Without elementwise fusion first, bubble-up-extract-slice hits the greedy +// rewrite limit on this chain and the pipeline fails without a diagnostic. +// With fusion: one vectorized loop, no temporary buffers. The checks follow +// the chain x -> 0.5x and x + c*x^3 -> exp(-2u) -> 1 / (1 + exp) -> 0.5x * (...) -> store +// inside the loop body. + +// CHECK-LABEL: func.func @gelu_tanh +// CHECK-NOT: memref.alloc +// CHECK: scf.for +// CHECK: %[[X:.+]] = vector.transfer_read {{.*}} vector<[4]xf32> +// CHECK: %[[HALF_X:.+]] = arith.mulf %[[X]], %{{.+}} : vector<[4]xf32> +// CHECK: arith.addf %[[X]], %{{.+}} : vector<[4]xf32> +// CHECK: %[[NEG:.+]] = arith.negf %{{.+}} : vector<[4]xf32> +// CHECK-NEXT: %[[EXP:.+]] = math.exp %[[NEG]] : vector<[4]xf32> +// CHECK-NEXT: %[[DEN:.+]] = arith.addf %[[EXP]], %{{.+}} : vector<[4]xf32> +// CHECK-NEXT: arith.divf %{{.+}}, %[[DEN]] : vector<[4]xf32> +// CHECK: %[[RES:.+]] = arith.mulf %[[HALF_X]], %{{.+}} : vector<[4]xf32> +// CHECK-NEXT: vector.transfer_write %[[RES]], {{.*}} vector<[4]xf32> +// CHECK-NEXT: } +// CHECK-NOT: scf.for +// CHECK-NOT: memref.alloc +// CHECK: return + +#map = affine_map<(d0) -> (d0)> +module { + func.func @gelu_tanh(%arg0: memref<*xf32> {tt.divisibility = 16 : i32}, %arg1: memref<*xf32> {tt.divisibility = 16 : i32}, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) { + %cst = arith.constant 4.471500e-02 : f32 + %c1024 = arith.constant 1024 : index + %cst_0 = arith.constant 0.797884583 : f32 + %cst_1 = arith.constant 5.000000e-01 : f32 + %cst_2 = arith.constant 2.000000e+00 : f32 + %cst_3 = arith.constant 1.000000e+00 : f32 + %c1024_i32 = arith.constant 1024 : i32 + %0 = tensor.empty() : tensor<1024xf32> + %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<1024xf32>) -> tensor<1024xf32> + %2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<1024xf32>) -> tensor<1024xf32> + %3 = linalg.fill ins(%cst_1 : f32) outs(%0 : tensor<1024xf32>) -> tensor<1024xf32> + %4 = linalg.fill ins(%cst_2 : f32) outs(%0 : tensor<1024xf32>) -> tensor<1024xf32> + %5 = linalg.fill ins(%cst_3 : f32) outs(%0 : tensor<1024xf32>) -> tensor<1024xf32> + %6 = arith.muli %arg6, %c1024_i32 : i32 + %7 = arith.index_cast %6 : i32 to index + %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [%7], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>> + %8 = arith.addi %7, %c1024 : index + %9 = arith.index_cast %arg2 : i32 to index + %10 = arith.minsi %8, %9 : index + %11 = arith.maxsi %10, %7 : index + %12 = arith.subi %11, %7 : index + %alloc = memref.alloc() : memref<1024xf32> + %subview = memref.subview %reinterpret_cast[0] [%12] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref> + %subview_4 = memref.subview %alloc[0] [%12] [1] : memref<1024xf32> to memref> + memref.copy %subview, %subview_4 : memref> to memref> + %13 = bufferization.to_tensor %alloc restrict writable : memref<1024xf32> to tensor<1024xf32> + %reinterpret_cast_5 = memref.reinterpret_cast %arg1 to offset: [%7], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>> + %14 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%13, %1 : tensor<1024xf32>, tensor<1024xf32>) outs(%13 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %15 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%14, %13 : tensor<1024xf32>, tensor<1024xf32>) outs(%14 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %16 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%15, %13 : tensor<1024xf32>, tensor<1024xf32>) outs(%15 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %17 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%13, %16 : tensor<1024xf32>, tensor<1024xf32>) outs(%13 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.addf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %18 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%17, %2 : tensor<1024xf32>, tensor<1024xf32>) outs(%17 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %19 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%13, %3 : tensor<1024xf32>, tensor<1024xf32>) outs(%13 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %20 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%18, %4 : tensor<1024xf32>, tensor<1024xf32>) outs(%18 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %21 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%20 : tensor<1024xf32>) outs(%20 : tensor<1024xf32>) { + ^bb0(%in: f32, %out: f32): + %29 = arith.negf %in : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %22 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%21 : tensor<1024xf32>) outs(%21 : tensor<1024xf32>) { + ^bb0(%in: f32, %out: f32): + %29 = math.exp %in : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %23 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%22, %5 : tensor<1024xf32>, tensor<1024xf32>) outs(%22 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.addf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %24 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%5, %23 : tensor<1024xf32>, tensor<1024xf32>) outs(%5 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.divf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %25 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%24, %4 : tensor<1024xf32>, tensor<1024xf32>) outs(%24 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %26 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%25, %5 : tensor<1024xf32>, tensor<1024xf32>) outs(%25 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.subf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %27 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%26, %5 : tensor<1024xf32>, tensor<1024xf32>) outs(%26 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.addf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %28 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%19, %27 : tensor<1024xf32>, tensor<1024xf32>) outs(%19 : tensor<1024xf32>) { + ^bb0(%in: f32, %in_7: f32, %out: f32): + %29 = arith.mulf %in, %in_7 : f32 + linalg.yield %29 : f32 + } -> tensor<1024xf32> + %extracted_slice = tensor.extract_slice %28[0] [%12] [1] : tensor<1024xf32> to tensor + %subview_6 = memref.subview %reinterpret_cast_5[0] [%12] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref> + bufferization.materialize_in_destination %extracted_slice in writable %subview_6 : (tensor, memref>) -> () + return + } +} + diff --git a/qcom_arm64_backend/test/Transforms/triton-linalg-to-llvm-math-approximation.mlir b/qcom_arm64_backend/test/Transforms/triton-linalg-to-llvm-math-approximation.mlir new file mode 100644 index 0000000..fa6f4ee --- /dev/null +++ b/qcom_arm64_backend/test/Transforms/triton-linalg-to-llvm-math-approximation.mlir @@ -0,0 +1,55 @@ +// RUN: arm64-mlir-opt %s --vector-ext=sve -triton-linalg-to-llvm | FileCheck %s + +// GELU (erf form) after -optimize-triton-shared. math.erf has no LLVM lowering, +// so it must be approximated inside the pipeline; the result stays an SVE +// vector computation with no libm calls. + +// CHECK-LABEL: llvm.func @gelu_erf +// CHECK-NOT: math. +// CHECK-NOT: llvm.call +// CHECK: llvm.intr.fma{{.*}} : (vector<[4]xf32>, vector<[4]xf32>, vector<[4]xf32>) -> vector<[4]xf32> +// CHECK-NOT: math. +// CHECK-NOT: llvm.call + +#map = affine_map<(d0)[s0, s1] -> (-d0 + s0, s1)> +module { + func.func @gelu_erf(%arg0: memref<*xf32> {tt.divisibility = 16 : i32}, %arg1: memref<*xf32> {tt.divisibility = 16 : i32}, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) { + %cst = arith.constant dense<1.000000e+00> : vector<[4]xf32> + %cst_0 = arith.constant dense<0.707106769> : vector<[4]xf32> + %cst_1 = arith.constant dense<5.000000e-01> : vector<[4]xf32> + %0 = ub.poison : f32 + %c0 = arith.constant 0 : index + %c4 = arith.constant 4 : index + %c1024 = arith.constant 1024 : index + %c1024_i32 = arith.constant 1024 : i32 + %1 = arith.muli %arg6, %c1024_i32 : i32 + %2 = arith.index_cast %1 : i32 to index + %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [%2], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>> + %3 = arith.addi %2, %c1024 : index + %4 = arith.index_cast %arg2 : i32 to index + %5 = arith.minsi %3, %4 : index + %6 = arith.maxsi %5, %2 : index + %7 = arith.subi %6, %2 : index + %reinterpret_cast_2 = memref.reinterpret_cast %arg1 to offset: [%2], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>> + %subview = memref.subview %reinterpret_cast_2[0] [%7] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref> + %subview_3 = memref.subview %reinterpret_cast[0] [%7] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref> + %vscale = vector.vscale + %c4_vscale = arith.muli %vscale, %c4 : index + scf.for %arg9 = %c0 to %7 step %c4_vscale { + %8 = affine.min #map(%arg9)[%7, %c4_vscale] + %subview_4 = memref.subview %subview_3[%arg9] [%8] [1] : memref> to memref> + %subview_5 = memref.subview %subview[%arg9] [%8] [1] : memref> to memref> + %9 = vector.create_mask %8 : vector<[4]xi1> + %10 = vector.transfer_read %subview_4[%c0], %0, %9 {in_bounds = [true]} : memref>, vector<[4]xf32> + %11 = arith.mulf %10, %cst_1 : vector<[4]xf32> + %12 = arith.mulf %10, %cst_0 : vector<[4]xf32> + %13 = math.erf %12 : vector<[4]xf32> + %14 = arith.addf %13, %cst : vector<[4]xf32> + %15 = arith.mulf %11, %14 : vector<[4]xf32> + vector.transfer_write %15, %subview_5[%c0], %9 {in_bounds = [true]} : vector<[4]xf32>, memref> + memref.copy %subview_5, %subview_5 : memref> to memref> + } + return + } +} +