From 44d6c8c37f706a9d85c814340ee3ac0c7f7ed7b4 Mon Sep 17 00:00:00 2001
From: Johnson Zhang <johnson9009@163.com>
Date: Fri, 20 Jun 2025 12:09:52 +0800
Subject: [PATCH 1/3] [BACKEND] Update Arm China NPU "Zhouyi" Backend Support

---
 third_party/aipu/CMakeLists.txt               |   1 +
 third_party/aipu/backend/aipu_torch_dev.cpp   |  44 +++-
 third_party/aipu/backend/analysis/__init__.py |   6 +-
 .../backend/analysis/determine_vfactor.py     |  16 +-
 .../analysis/get_linalg_generic_size.py       |  25 ++
 third_party/aipu/backend/codegen.py           | 244 +++++++++++++-----
 third_party/aipu/backend/compiler.py          |  24 +-
 third_party/aipu/backend/driver.py            |  12 +-
 .../aipu/backend/transform/__init__.py        |   3 +-
 .../aipu/backend/transform/binding_tid.py     |  35 ++-
 .../transform/canonical_const_dtype.py        |  48 ++++
 third_party/aipu/include/Passes/Passes.td     |  20 ++
 third_party/aipu/language/aipu/libdevice.py   |  47 +---
 third_party/aipu/lib/Passes/CMakeLists.txt    |   1 +
 .../Passes/SCFLoopBufferizationPreprocess.cpp |  56 ++++
 third_party/aipu/lib/main.cpp                 |   5 +-
 .../aipu/python/test/test_codegen_dma.py      | 103 +++++---
 .../aipu/test/convert_bool_arg_to_i8.mlir     |   9 +
 third_party/aipu/test/lit.cfg                 |   4 +
 .../one-shot-bufferize-preprocessing.mlir     |  26 ++
 third_party/aipu/tools/CMakeLists.txt         |   1 +
 .../aipu/tools/aipu-opt/CMakeLists.txt        |  18 ++
 third_party/aipu/tools/aipu-opt/aipu-opt.cpp  |  15 ++
 23 files changed, 611 insertions(+), 152 deletions(-)
 create mode 100644 third_party/aipu/backend/analysis/get_linalg_generic_size.py
 create mode 100644 third_party/aipu/backend/transform/canonical_const_dtype.py
 create mode 100644 third_party/aipu/lib/Passes/SCFLoopBufferizationPreprocess.cpp
 create mode 100644 third_party/aipu/test/convert_bool_arg_to_i8.mlir
 create mode 100644 third_party/aipu/test/lit.cfg
 create mode 100644 third_party/aipu/test/one-shot-bufferize-preprocessing.mlir
 create mode 100644 third_party/aipu/tools/CMakeLists.txt
 create mode 100644 third_party/aipu/tools/aipu-opt/CMakeLists.txt
 create mode 100644 third_party/aipu/tools/aipu-opt/aipu-opt.cpp

diff --git a/third_party/aipu/CMakeLists.txt b/third_party/aipu/CMakeLists.txt
index 2ab8d44cb..691d3e393 100644
--- a/third_party/aipu/CMakeLists.txt
+++ b/third_party/aipu/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_subdirectory(include)
 add_subdirectory(lib)
+add_subdirectory(tools)
 
 add_triton_plugin(TritonAIPU ${CMAKE_CURRENT_SOURCE_DIR}/triton_aipu.cc)
 target_include_directories(TritonAIPU PRIVATE ${CMAKE_SOURCE_DIR}/third_party/flir/include)
diff --git a/third_party/aipu/backend/aipu_torch_dev.cpp b/third_party/aipu/backend/aipu_torch_dev.cpp
index dc5dfc6b8..31036a944 100644
--- a/third_party/aipu/backend/aipu_torch_dev.cpp
+++ b/third_party/aipu/backend/aipu_torch_dev.cpp
@@ -11,6 +11,8 @@
 
 #include <ATen/EmptyTensor.h>
 #include <ATen/InferSize.h>
+#include <ATen/core/GeneratorForPrivateuseone.h>
+#include <ATen/native/CPUFallback.h>
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/DistributionTemplates.h>
 #include <ATen/native/cpu/DistributionTemplates.h>
@@ -228,9 +230,29 @@ Tensor aipu_copy_from(const Tensor &self, const Tensor &dst,
   auto status = aipu_memcpy(aipu_ctx_, dst.data_ptr(), self.data_ptr(),
                             self.nbytes(), kind);
   AIPU_DRIVER_HANDLE_ERROR(status);
-  return self;
+  return dst;
 }
+Tensor aipu_copy_from_and_resize(const Tensor &self, const Tensor &dst) {
+  if (self.sizes() != dst.sizes()) {
+    auto new_dst =
+        custom_empty_symint(self.sizes(), self.scalar_type(), c10::nullopt,
+                            c10::nullopt, c10::nullopt, c10::nullopt);
+    auto kind = AIPU_MEMCPY_HOST_TO_DEVICE;
+    if (StrStartsWith(self.device().str(), "aipu")) {
+      kind = AIPU_MEMCPY_DEVICE_TO_HOST;
+      if (StrStartsWith(dst.device().str(), "aipu")) {
+        kind = AIPU_MEMCPY_DEVICE_TO_DEVICE;
+      }
+    }
+    auto aipu_ctx_ = AIPUAllocator::aipu_ctx_;
+    auto status = aipu_memcpy(aipu_ctx_, new_dst.data_ptr(), self.data_ptr(),
+                              self.nbytes(), kind);
+    AIPU_DRIVER_HANDLE_ERROR(status);
 
+    return new_dst;
+  }
+  return aipu_copy_from(self, dst, false);
+}
 template <template <typename> class RND>
 Tensor &random_kernel(Tensor &self, double cond1, double cond2,
                       c10::optional<Generator> gen) {
@@ -323,12 +345,30 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
   m.impl("aten::uniform_", &random_kernel<uniform_real_distribution>);
   m.impl("aten::normal_", &random_kernel<normal_distribution>);
   m.impl("aten::_copy_from", &aipu_copy_from);
+  m.impl("aten::_copy_from_and_resize", &aipu_copy_from_and_resize);
   m.impl("aten::random_.from",
          &random_from_to_kernel<uniform_int_from_to_distribution>);
   m.impl("aten::_local_scalar_dense", &_local_scalar_dense_aipu);
   m.impl("aten::fill_.Scalar", &fill_scalar_aipu);
 }
 
+void custom_cpu_fallback(const c10::OperatorHandle &op,
+                         torch::jit::Stack *stack) {
+  at::native::cpu_fallback(op, stack);
+}
+
+TORCH_LIBRARY_IMPL(_, PrivateUse1, m) {
+  m.fallback(torch::CppFunction::makeFromBoxedFunction<&custom_cpu_fallback>());
+}
+
+at::Generator make_custom_generator(c10::DeviceIndex device_index) {
+  return at::detail::getDefaultCPUGenerator();
+}
+
+REGISTER_GENERATOR_PRIVATEUSE1(make_custom_generator)
+
+C10_REGISTER_GUARD_IMPL(PrivateUse1, c10::impl::AIPUGuardImpl)
+
 // Register the autograd dispatch key for operators that have no dispatches
 TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {
   m.impl("isfinite", torch::autograd::autogradNotImplementedFallback());
@@ -352,6 +392,7 @@ struct _DeviceGuard {
 
 struct _Device {
   _Device(c10::Device device) { idx = device.index(); }
+  _Device(int index) { idx = index; }
 
   int idx = 0;
   int prev_idx = -1;
@@ -375,6 +416,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   py::class_<_Device>(m, "device", py::module_local())
       .def(py::init(
           [](c10::Device device) { return std::make_unique<_Device>(device); }))
+      .def(py::init([](int index) { return std::make_unique<_Device>(index); }))
       .def("__enter__", [](_Device &self) { ; })
       .def("__exit__",
            [](_Device &self, pybind11::object type, pybind11::object value,
diff --git a/third_party/aipu/backend/analysis/__init__.py b/third_party/aipu/backend/analysis/__init__.py
index 7b15c6ef2..fbaba7c30 100644
--- a/third_party/aipu/backend/analysis/__init__.py
+++ b/third_party/aipu/backend/analysis/__init__.py
@@ -1,3 +1,7 @@
 from .determine_vfactor import determine_vectorization_factor
+from .get_linalg_generic_size import get_linalg_generic_size
 
-__all__ = ["determine_vectorization_factor"]
+__all__ = [
+    "determine_vectorization_factor",
+    "get_linalg_generic_size",
+]
diff --git a/third_party/aipu/backend/analysis/determine_vfactor.py b/third_party/aipu/backend/analysis/determine_vfactor.py
index a2e1ef051..d29620f69 100644
--- a/third_party/aipu/backend/analysis/determine_vfactor.py
+++ b/third_party/aipu/backend/analysis/determine_vfactor.py
@@ -16,15 +16,23 @@ def determine_vectorization_factor(module, target_bitwidth=256, debug=False):
     """
     min_width = target_bitwidth
 
+    def _get_width(value):
+        elem_type = (value.type.element_type if hasattr(value.type, 'element_type') else value.type)
+        elem_width = 32 if isinstance(elem_type, ir.IndexType) else elem_type.width
+        # Except bool dtype
+        if elem_width == 1:
+            return None
+        return elem_width
+
     def walk_callback(op):
         nonlocal min_width
         if op.name == "affine.for":
             all_ops = (_op for region in op.regions for block in region.blocks for _op in block.operations)
             for _op in all_ops:
-                for result in _op.results:
-                    elem_type = (result.type.element_type if hasattr(result.type, 'element_type') else result.type)
-                    elem_width = 32 if isinstance(elem_type, ir.IndexType) else elem_type.width
-                    min_width = min(min_width, elem_width)
+                items = _op.results or _op.operands
+                for item in items:
+                    if width := _get_width(item):
+                        min_width = min(min_width, width)
 
         return ir.WalkResult.ADVANCE
 
diff --git a/third_party/aipu/backend/analysis/get_linalg_generic_size.py b/third_party/aipu/backend/analysis/get_linalg_generic_size.py
new file mode 100644
index 000000000..6c3a684c8
--- /dev/null
+++ b/third_party/aipu/backend/analysis/get_linalg_generic_size.py
@@ -0,0 +1,25 @@
+from mlir import ir
+
+
+def get_linalg_generic_size(module):
+    """
+    Get the size of linalg generic.
+
+    Args:
+        module: The Triton module to analyze.
+
+    Returns:
+        int: The size of linalg generic.
+    """
+    size = 0
+
+    def walk_callback(op):
+        nonlocal size
+        if op.name == "linalg.generic":
+            size += 1
+
+        return ir.WalkResult.ADVANCE
+
+    module.operation.walk(walk_callback, ir.WalkOrder.PRE_ORDER)
+
+    return size
diff --git a/third_party/aipu/backend/codegen.py b/third_party/aipu/backend/codegen.py
index 7436d2c01..627a856b9 100644
--- a/third_party/aipu/backend/codegen.py
+++ b/third_party/aipu/backend/codegen.py
@@ -32,6 +32,7 @@
     13: T.NE,
 }
 _MEMORY_SCOPE_MAPPING = {
+    -1: "local",
     4: "lsram",
     8: "shared",
     11: "alloc_event",
@@ -88,14 +89,34 @@ def _convert_vector_type(type):
     raise RuntimeError(f"not scalar type {type}")
 
 
+def _compute_strides(shape):
+    """compute tensor strides with given shape"""
+    strides = [1]
+    for dim_size in reversed(shape[1:]):
+        strides.append(strides[-1] * dim_size)
+    return list(reversed(strides))
+
+
+def _is_scalar_type(ty):
+    return isinstance(ty, (mlir_ir.IndexType, mlir_ir.IntegerType, mlir_ir.FloatType))
+
+
+def _is_vector_type(ty):
+    return isinstance(ty, mlir_ir.VectorType)
+
+
+def _is_pointer_type(ty):
+    return isinstance(ty, (mlir_ir.MemRefType, mlir_ir.UnrankedMemRefType))
+
+
 def _get_type(value):
     ty = value.type
 
-    if isinstance(ty, (mlir_ir.IndexType, mlir_ir.IntegerType, mlir_ir.FloatType)):
+    if _is_scalar_type(ty):
         return _convert_scalar_type(ty)
-    elif isinstance(ty, mlir_ir.VectorType):
+    elif _is_vector_type(ty):
         return _convert_vector_type(ty)
-    elif isinstance(ty, (mlir_ir.MemRefType, mlir_ir.UnrankedMemRefType)):
+    elif _is_pointer_type(ty):
         e_dtype = _convert_scalar_type(ty.element_type)
         return ir.PointerType(ir.PrimType(e_dtype))
 
@@ -112,6 +133,12 @@ def _get_shape(value):
     raise RuntimeError(f"Cannot parse shape {ty}")
 
 
+def _get_buffer(ptr):
+    if isinstance(ptr, tir.Buffer):
+        return ptr
+    return _get_buffer(ptr.buffer)
+
+
 class CodeGenerator():
 
     def __init__(self, mod) -> None:
@@ -128,6 +155,19 @@ def __init__(self, mod) -> None:
         self.after_args = None
         self.yeild_args = None
 
+    def _get_associated_dtype(self, op):
+        # Find the associated_dtype of the bool arg.
+        try:
+            owner = op.operands[0].owner
+            associated_dtype = self.get_operand(owner, 0).dtype
+            while associated_dtype.startswith("bool"):
+                owner = owner.operands[0].owner
+                associated_dtype = self.get_operand(owner, 0).dtype
+        except IndexError:
+            associated_dtype = "int32"
+        finally:
+            return associated_dtype
+
     def create_var_name(self):
         var_name = "var_" + str(self.name_idx)
         self.name_idx += 1
@@ -141,6 +181,9 @@ def emit_let(self, value, related_value):
     def get_operand(self, op, idx):
         return self.get_or_create_var(op.operands[idx])
 
+    def get_list_value(self, list_value):
+        return [self.get_or_create_var(i) for i in list_value]
+
     def get_or_create_var(self, value):
         if value in self.mlir_to_tir_mapping:
             return self.mlir_to_tir_mapping[value]
@@ -154,6 +197,18 @@ def get_or_create_var(self, value):
         self.mlir_to_tir_mapping[value] = var
         return var
 
+    def get_offsets_sizes_strides(self, op):
+
+        def _get_static_or_dynamic_value(static_v, dynamic_v):
+            dynamic_v = list(dynamic_v)
+            return [self.get_or_create_var(dynamic_v.pop(0)) if v < 0 else v for v in static_v]
+
+        offsets = _get_static_or_dynamic_value(op.static_offsets, op.offsets)
+        sizes = _get_static_or_dynamic_value(op.static_sizes, op.sizes)
+        strides = _get_static_or_dynamic_value(op.static_strides, op.strides)
+
+        return offsets, sizes, strides
+
     def for_range(self, begin, end, step, kind="serial"):
         self.ib._seq_stack.append([])
 
@@ -193,6 +248,17 @@ def exit_scope(self):
 
     def dispatch(self, op, stage):
         op_name = "func.func" if isinstance(op, func.FuncOp) else op.name
+        # Extern Call
+        if "extern_call" in op.attributes:
+            func_name = op.attributes["extern_call"].value
+            if func_name == "fmod":
+                fmod = lambda x, y: T.call_extern(_get_type(op.result), "fmod", x, y)
+                self.gen_binary(op, fmod)
+            if func_name == "div_rz":
+                self.ib.emit(tir.call_extern("void", "__vset_rounding_mode_rtz"))
+                self.gen_binary(op, T.Div)
+                self.ib.emit(tir.call_extern("void", "__vset_rounding_mode_rtn"))
+            return
         # Memref Dialect
         if op_name == "memref.reinterpret_cast":
             self.gen_memref_reinterpret_cast(op)
@@ -210,6 +276,12 @@ def dispatch(self, op, stage):
             self.gen_dma_start(op)
         elif op_name == "memref.dma_wait":
             self.gen_dma_wait(op)
+        elif op_name == "memref.cast":
+            self.gen_memref_cast(op)
+        elif op_name == "memref.expand_shape":
+            self.gen_memref_expand_shape(op)
+        elif op_name == "memref.collapse_shape":
+            self.gen_memref_collapse_shape(op)
         # Arith Dialect
         elif op_name == "arith.constant":
             self.gen_arith_constant(op)
@@ -235,6 +307,9 @@ def dispatch(self, op, stage):
             self.gen_binary(op, T.bitwise_xor)
         elif op_name in ("arith.remsi", "arith.remui"):
             self.gen_binary(op, T.Mod)
+        elif op_name in ("arith.remf", ):
+            remainder = lambda x, y: T.call_extern(_get_type(op.result), "remainder", x, y)
+            self.gen_binary(op, remainder)
         elif op_name == "arith.cmpi":
             self.gen_binary(op, _CMPI_MAPPING[op.predicate.value])
         elif op_name == "arith.cmpf":
@@ -244,14 +319,24 @@ def dispatch(self, op, stage):
             self.gen_arith_cast(op)
         elif op_name == "arith.select":
             self.gen_select(op)
+        elif op_name in ("arith.shrsi", "arith.shrui"):
+            self.gen_binary(op, T.shift_right)
+        elif op_name == "arith.shli":
+            self.gen_binary(op, T.shift_left)
+        elif op_name in ("arith.mulsi_extended", "arith.mului_extended"):
+            self.gen_arith_mul_extended(op)
         # Math Dialect
         elif op_name == "mathext.fmod":
             fmod = lambda x, y: T.call_extern(_get_type(op.result), "fmod", x, y)
             self.gen_binary(op, fmod)
         elif op_name == "math.powf":
             self.gen_binary(op, S.pow)
+        elif op_name == "math.rsqrt":
+            self.gen_unary(op, S.rsqrt)
         elif op_name == "math.tanh":
             self.gen_unary(op, S.tanh)
+        elif op_name == "math.exp2":
+            self.gen_unary(op, S.exp2)
         elif op_name == "math.exp":
             self.gen_unary(op, S.exp)
         elif op_name == "math.absf":
@@ -266,6 +351,11 @@ def dispatch(self, op, stage):
             self.gen_unary(op, S.erf)
         elif op_name == "math.log":
             self.gen_unary(op, S.log)
+        elif op_name == "math.floor":
+            self.gen_unary(op, S.floor)
+        elif op_name == "math.trunc":
+            trunc = lambda x: T.call_extern(_get_type(op.result), "trunc", x)
+            self.gen_unary(op, trunc)
         # Func Dialect
         elif op_name == "func.return":
             self.gen_func_return(op)
@@ -311,88 +401,120 @@ def gen_memref_reinterpret_cast(self, op):
         result = op.result
         arg = self.get_operand(op, 0)
         dtype = _get_type(result).element_type.dtype
-        offset = 0
-        if len(op.operands) == 2:
-            offset = self.get_operand(op, 1)
+        offsets, sizes, strides = self.get_offsets_sizes_strides(op)
 
-        buffer = T.Buffer((-1, ), elem_offset=offset, data=arg.base, dtype=dtype)
+        buffer = T.Buffer(sizes, elem_offset=offsets[0], data=arg.base, dtype=dtype, strides=strides)
         self.mlir_to_tir_mapping[result] = buffer
 
     def gen_memref_load(self, op):
         result = op.result
-        buffer = self.get_operand(op, 0)
-        index = [0]
-        if len(op.operands) >= 2:
-            index = [self.get_operand(op, i) for i in range(1, len(op.operands))]
-        self.emit_let(T.BufferLoad(buffer, index), result)
+        buffer = _get_buffer(self.get_operand(op, 0))
+        indices = self.get_list_value(op.indices) or [0]
+
+        self.emit_let(T.BufferLoad(buffer, indices), result)
 
     def gen_memref_store(self, op):
         value = self.get_operand(op, 0)
-        buffer = self.get_operand(op, 1)
+        buffer = _get_buffer(self.get_operand(op, 1))
+        indices = self.get_list_value(op.indices) or [0]
 
-        index = [0]
-        if len(op.operands) >= 3:
-            index = [self.get_operand(op, i) for i in range(2, len(op.operands))]
-        self.ib.emit(tir.BufferStore(buffer, value, index))
+        self.ib.emit(tir.BufferStore(buffer, value, indices))
 
     def gen_memref_alloc(self, op):
         result = op.result
         dtype = _get_type(result).element_type.dtype
         shape = _get_shape(result)
-        # set default memory space: lsram
-        scope_value = result.type.memory_space.value if result.type.memory_space else 4
+        scope_value = result.type.memory_space.value if result.type.memory_space else -1
+
         if scope_value == 11:
-            event = S.alloc_events(1)
-            self.mlir_to_tir_mapping[result] = event
+            self.emit_let(S.alloc_events(1), op.result)
         else:
             buf = self.ib.allocate(dtype, shape, scope=_MEMORY_SCOPE_MAPPING[scope_value])
+            buf._buffer.strides = _compute_strides(shape)
             self.mlir_to_tir_mapping[result] = buf._buffer
 
     def gen_dma_start(self, op):
         #  currently, we only support one event, skip stride
-        src = self.get_operand(op, 0)
-        src = src.buffer if isinstance(src, tir.Pointer) else src
-        dst = self.get_operand(op, 2)
-        dst = dst.buffer if isinstance(dst, tir.Pointer) else dst
+        src = _get_buffer(self.get_operand(op, 0))
+        dst = _get_buffer(self.get_operand(op, 2))
         src_index = self.get_operand(op, 1)
         dst_index = self.get_operand(op, 3)
         num_elements = self.get_operand(op, 4)
         event = self.get_operand(op, 5)
+
         self.ib.emit(S.async_dma_copy(dst.addr_of(dst_index), src.addr_of(src_index), num_elements, event=event))
 
     def gen_dma_wait(self, op):
         # currently, we only support one event
         event = self.get_operand(op, 0)
+
         self.ib.emit(S.wait_events(event))
 
     def gen_memref_copy(self, op):
         src = self.get_operand(op, 0)
         dst = self.get_operand(op, 1)
-        width = src.shape[0]
+        shape = src.shape
+
+        if len(shape) == 1:
+            dma_copy = S.dma_copy(dst, src, shape[0])
+        else:
+            if not len(shape) == 2:
+                raise RuntimeError(f"only suport 1d/2d DMA copy, but got shape={shape}")
+            dma_copy = S.dma_copy(dst, src, width=src.shape[1], src_stride=src.strides[0], times=src.shape[0],
+                                  dst_stride=dst.strides[0])
 
-        dma_copy = S.dma_copy(dst, src, width)
         self.ib.emit(dma_copy)
 
     def gen_memref_subview(self, op):
         result = op.result
-        arg0 = self.get_operand(op, 0)
-        buffer = arg0.buffer if isinstance(arg0, tir.Pointer) else arg0
-        size = self.get_operand(op, 1)
+        buffer = _get_buffer(self.get_operand(op, 0))
+        offsets, sizes, strides = self.get_offsets_sizes_strides(op)
+
+        buf_strides = buffer.strides
+        if len(strides) != len(buf_strides):
+            buf_strides = _compute_strides(buffer.shape)
+        update_strides = [buf_strides[i] * strides[i] for i in range(len(strides))]
+        update_offset = buffer.elem_offset + sum(off * stride for off, stride in zip(offsets, buf_strides))
 
-        subview = T.Buffer(size, elem_offset=buffer.elem_offset, data=buffer.data, dtype=buffer.dtype)
+        subview = T.Buffer(sizes, elem_offset=update_offset, data=buffer.data, dtype=buffer.dtype,
+                           strides=update_strides)
         self.mlir_to_tir_mapping[result] = subview
 
+    def gen_memref_expand_shape(self, op):
+        result = op.result
+        inp_buf = _get_buffer(self.get_operand(op, 0))
+        out_shape = list(op.static_output_shape)
+
+        expanded_buffer = T.Buffer(out_shape, elem_offset=inp_buf.elem_offset, data=inp_buf.data, dtype=inp_buf.dtype,
+                                   strides=_compute_strides(out_shape))
+        self.mlir_to_tir_mapping[result] = expanded_buffer
+
+    def gen_memref_collapse_shape(self, op):
+        result = op.result
+        inp_buf = _get_buffer(self.get_operand(op, 0))
+        out_shape = result.type.shape
+
+        out_buffer = T.Buffer(out_shape, elem_offset=inp_buf.elem_offset, data=inp_buf.data, dtype=inp_buf.dtype,
+                              strides=_compute_strides(out_shape))
+        self.mlir_to_tir_mapping[result] = out_buffer
+
+    def gen_memref_cast(self, op):
+        result = op.result
+        arg = self.get_operand(op, 0)
+        dtype = _get_type(result).element_type.dtype
+
+        self.mlir_to_tir_mapping[result] = arg.as_ptr(dtype)
+
     def gen_arith_constant(self, op):
 
         def _create_const_expr(op):
             ty = op.result.type
             dtype = _get_type(op.result)
-            # scalar
-            if isinstance(ty, (mlir_ir.IndexType, mlir_ir.IntegerType, mlir_ir.FloatType)):
+
+            if _is_scalar_type(ty):
                 value = bool(op.value) if dtype == "bool" else op.literal_value
                 return tir.const(value, dtype)
-            # vector
-            if isinstance(ty, mlir_ir.VectorType):
+            if _is_vector_type(ty):
                 const_value = op.value.maybe_downcast()
                 # For FP16, the C++ interface __get_item__ do not have a proper implementation.
                 # So here use np.array to directly using its raw data.
@@ -423,31 +545,32 @@ def gen_binary(self, op, method):
     def gen_select(self, op):
         #cond, true_value, false_value
         result = op.result
-
         arg0 = self.get_operand(op, 0)
         arg1 = self.get_operand(op, 1)
         arg2 = self.get_operand(op, 2)
-        if isinstance(result.type, mlir_ir.VectorType):
-            self.emit_let(S.vsel(arg1, arg2, mask=arg0), result)
-        else:
-            self.emit_let(tir.Select(arg0, arg1, arg2), result)
+
+        self.emit_let(tir.Select(arg0, arg1, arg2), result)
 
     def gen_arith_cast(self, op):
         result = op.result
         arg0 = self.get_operand(op, 0)
 
         if arg0.dtype.startswith("bool"):
-            # Find the associated_dtype of the bool arg.
-            owner = op.operands[0].owner
-            associated_dtype = self.get_operand(owner, 0).dtype
-            while associated_dtype.startswith("bool"):
-                owner = owner.operands[0].owner
-                associated_dtype = self.get_operand(owner, 0).dtype
+            associated_dtype = self._get_associated_dtype(op)
+            arg0 = T.Select(arg0, S.cast(1, associated_dtype), 0)
 
-            vsel = S.vsel(S.cast(1, associated_dtype), 0, arg0)
-            self.emit_let(S.cast(vsel, _get_type(result)), result)
-        else:
-            self.emit_let(S.cast(arg0, _get_type(result)), result)
+        self.emit_let(S.cast(arg0, _get_type(result)), result)
+
+    def gen_arith_mul_extended(self, op):
+        mull, mulh = op.results
+        arg0 = self.get_operand(op, 0)
+        arg1 = self.get_operand(op, 1)
+
+        self.emit_let(T.Mul(arg0, arg1), mull)
+        mulh_func = "mul_hi"
+        if isinstance(mulh.type, mlir_ir.VectorType):
+            mulh_func = "vmul_hi"
+        self.emit_let(T.call_extern(_get_type(mulh), mulh_func, arg0, arg1), mulh)
 
     def gen_unary(self, op, method):
         result = op.result
@@ -523,6 +646,8 @@ def gen_scf_if(self, op, stage):
         # Finish
         if stage.is_after_all_regions():
             self.exit_scope()
+            for i, value in enumerate(op.results):
+                self.mlir_to_tir_mapping[value] = self.yeild_args[i]
 
     def gen_scf_while(self, op, stage):
         if stage.is_before_all_regions():
@@ -555,28 +680,23 @@ def gen_scf_while(self, op, stage):
 
     def gen_vload(self, op):
         result = op.result
-        arg0 = self.get_operand(op, 0)
-        buffer = arg0.buffer if isinstance(arg0, tir.Pointer) else arg0
-        index = 0
-        if len(op.operands) >= 2:
-            index = self.get_operand(op, 1)
+        buffer = _get_buffer(self.get_operand(op, 0))
+        indices = self.get_list_value(op.indices) or 0
 
-        self.emit_let(S.vload(buffer.addr_of(index), lanes=result.type.shape[0]), result)
+        self.emit_let(S.vload(buffer.addr_of(indices), lanes=result.type.shape[0]), result)
 
     def gen_vstore(self, op):
         value = self.get_operand(op, 0)
-        arg1 = self.get_operand(op, 1)
-        buffer = arg1.buffer if isinstance(arg1, tir.Pointer) else arg1
-        index = 0
-        if len(op.operands) >= 3:
-            index = self.get_operand(op, 2)
+        buffer = _get_buffer(self.get_operand(op, 1))
+        indices = self.get_list_value(op.indices) or 0
 
-        self.ib.emit(S.vstore(value, buffer.addr_of(index)))
+        self.ib.emit(S.vstore(value, buffer.addr_of(indices)))
 
     def gen_vbcast(self, op):
         result = op.result
         value = self.get_operand(op, 0)
         dtype = result.type
+
         self.emit_let(S.vbcast(S.cast(value, value.dtype), lanes=dtype.shape[0]), result)
 
 
diff --git a/third_party/aipu/backend/compiler.py b/third_party/aipu/backend/compiler.py
index f3554128a..5a2c16b84 100644
--- a/third_party/aipu/backend/compiler.py
+++ b/third_party/aipu/backend/compiler.py
@@ -98,7 +98,7 @@ def make_ttir(mod, metadata, opt):
     def make_linalg(mod, metadata, opt):
         pm = ir.pass_manager(mod.context)
         pm.enable_debug()
-        # add pass here
+        # Add pass here.
         aipu.passes.convert.add_triton_to_linalg_pipeline(pm)
         pm.run(mod)
         return mod
@@ -109,8 +109,14 @@ def make_aipubin(mod, metadata, opt):
         ctx.allow_unregistered_dialects = True
         pm = PassManager("builtin.module", ctx)
         mod = Module.parse(aipu.common.generic_print(mod), ctx)
-        # add pass here
-        pm.add("func.func(linalg-fuse-elementwise-ops)")
+
+        # Add pass here.
+        linalg_generic_size = analysis.get_linalg_generic_size(mod)
+        # Affine loop fusion run failed when dealing with ir which owns more than 50 loops.
+        # So here fuse the loop in linalg.
+        if linalg_generic_size > 50:
+            pm.add("func.func(linalg-fuse-elementwise-ops)")
+        pm.add("scf-loop-bufferization-preprocessing")
         pm.add("one-shot-bufferize")
         pm.add("func.func(convert-bool-arg-to-i8)")
         pm.add("func.func(convert-linalg-to-affine-loops)")
@@ -119,10 +125,20 @@ def make_aipubin(mod, metadata, opt):
 
         pm = PassManager("builtin.module", ctx)
         vfactor = analysis.determine_vectorization_factor(mod, metadata["vector_register_bits"])
-        pm.add(f"func.func(affine-super-vectorize{{virtual-vector-size={vfactor}}})")
+        if vfactor > 1:
+            pm.add(f"func.func(affine-super-vectorize{{virtual-vector-size={vfactor}}})")
         pm.add("func.func(lower-affine)")
+
+        # Optimize pass.
+        pm.add("func.func(canonicalize)")
+        pm.add("func.func(cse)")
+        pm.add("func.func(reconcile-unrealized-casts)")
         pm.run(mod.operation)
+
+        # Post aipu pass.
         transform.binding_tid(mod, ctx)
+        transform.canonical_const_dtype(mod, ctx)
+
         ex = codegenAIPU(mod)
         metadata["name"] = ex._func_name
         metadata["shared"] = 1
diff --git a/third_party/aipu/backend/driver.py b/third_party/aipu/backend/driver.py
index b9c244750..9fa46b5a0 100644
--- a/third_party/aipu/backend/driver.py
+++ b/third_party/aipu/backend/driver.py
@@ -49,7 +49,7 @@ def __init__(self, src, metadata):
 
     # TODO(aipu-teams): This is just a temporary solution for now, because the real driver interface is not ready yet.
     # These code will be refactor later.
-    def __call__(self, gridX, gridY, gridZ, stream, function, *args):
+    def __call__(self, gridx_size, gridy_size, gridz_size, stream, function, *args):
         try:
             from flag_gems.utils.tensor_wrapper import StridedBuffer
         except ImportError:
@@ -74,10 +74,10 @@ def __call__(self, gridX, gridY, gridZ, stream, function, *args):
                 np_args[i] = arr.astype(np.int8)
                 bool_index.append(i)
 
-        tail_args = [gridX, gridY, gridZ, 0, 0, 0]
+        tail_args = [gridx_size, gridy_size, gridz_size, 0, 0, 0]
         tec_num = 4
 
-        for i in range((gridX + tec_num - 1) // tec_num):
+        for i in range((gridx_size * gridy_size * gridz_size + tec_num - 1) // tec_num):
             tail_args[3] = i
             ex(*(np_args + tail_args))
 
@@ -136,8 +136,12 @@ def is_active():
             torch.utils.generate_methods_for_privateuse1_backend(for_storage=True)
         return torch.aipu.is_available()
 
+    # TODO(aipu-teams): Support bechmarker later.
     def get_benchmarker(self):
-        from triton.testing import do_bench
+
+        def do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode="mean"):
+            return [float("inf"), float("inf"), float("inf")]
+
         return do_bench
 
     def get_empty_cache_for_benchmark(self):
diff --git a/third_party/aipu/backend/transform/__init__.py b/third_party/aipu/backend/transform/__init__.py
index 375d1376f..cd70545f0 100644
--- a/third_party/aipu/backend/transform/__init__.py
+++ b/third_party/aipu/backend/transform/__init__.py
@@ -1,3 +1,4 @@
 from .binding_tid import binding_tid
+from .canonical_const_dtype import canonical_const_dtype
 
-__all__ = ["binding_tid"]
+__all__ = ["binding_tid", "canonical_const_dtype"]
diff --git a/third_party/aipu/backend/transform/binding_tid.py b/third_party/aipu/backend/transform/binding_tid.py
index 18e2b58c0..85701212d 100644
--- a/third_party/aipu/backend/transform/binding_tid.py
+++ b/third_party/aipu/backend/transform/binding_tid.py
@@ -1,10 +1,19 @@
 from mlir import ir
-from mlir.dialects import arith, func
+from mlir.dialects import arith, func, scf
 
 
 def binding_tid(module, ctx):
     """
     Binding tid to the third-to-last parameter.
+    Triton convention: [gridx_size, gridy_size, gridz_size, gridx_idx, gridy_idx, gridz_idx].
+    Aipu_driver set args: [gridx_size, gridy_size, gridz_size, loop_tec_idx, 0, 0].
+
+    Recover gridx_idx, gridy_idx in func_body:
+
+    grid_flat_idx = loop_tec_idx * tec_num + tec_id
+    gridx_idx = grid_flat_idx % gridx_size
+    gridy_idx = (grid_flat_idx // gridx_size) % gridy_size
+    gridz_idx = grid_flat_idx // (gridx_size * gridy_size)
 
     Args:
         module: The mlir module to analyze
@@ -17,15 +26,29 @@ def binding_tid(module, ctx):
     def walk_callback(op):
         if op.name == "func.func":
             block = op.regions[0].blocks[0]
-            gridx = block.arguments[-3]
+            gridx, gridy, gridz = block.arguments[-3:]
+            gridx_size, gridy_size, gridz_size = block.arguments[-6:-3]
             with ctx, op.location, ir.InsertionPoint.at_block_begin(block):
                 i32 = ir.IntegerType.get_signless(32)
                 local_size = func.call([i32], "local_size", [])
                 local_id = func.call([i32], "local_id", [])
-                var_a = arith.muli(gridx, local_size)
-                var_b = arith.addi(var_a, local_id)
-
-                gridx.replace_all_uses_except(var_b, var_a.owner)
+                mul_tid = arith.muli(gridx, local_size)
+                flatten_tid = arith.addi(mul_tid, local_id)
+
+                xy_size = arith.muli(gridx_size, gridy_size)
+                total_size = arith.muli(xy_size, gridz_size)
+                cond = arith.cmpi(arith.CmpIPredicate.sge, flatten_tid, total_size)
+                if_op = scf.IfOp(cond)
+                with ir.InsertionPoint(if_op.then_block):
+                    func.ReturnOp([])
+
+                new_gridx = arith.remsi(flatten_tid, gridx_size)
+                new_gridy = arith.remsi(arith.divsi(flatten_tid, gridx_size), gridy_size)
+                new_gridz = arith.divsi(flatten_tid, xy_size)
+
+                gridx.replace_all_uses_except(new_gridx, mul_tid.owner)
+                gridy.replace_all_uses_with(new_gridy)
+                gridz.replace_all_uses_with(new_gridz)
 
         return ir.WalkResult.ADVANCE
 
diff --git a/third_party/aipu/backend/transform/canonical_const_dtype.py b/third_party/aipu/backend/transform/canonical_const_dtype.py
new file mode 100644
index 000000000..f2cc845ce
--- /dev/null
+++ b/third_party/aipu/backend/transform/canonical_const_dtype.py
@@ -0,0 +1,48 @@
+from mlir import ir
+from mlir.dialects import arith
+import numpy as np
+
+
+def _is_in_range(val, dtype):
+    iinfo = np.iinfo(dtype)
+    return iinfo.min <= val <= iinfo.max
+
+
+def _get_cononical_type(val):
+    if _is_in_range(val, np.int32):
+        return ir.IntegerType.get_signed(32)
+    elif _is_in_range(val, np.uint32):
+        return ir.IntegerType.get_unsigned(32)
+    else:
+        return None
+
+
+def canonical_const_dtype(module, ctx):
+    """
+    Canonical i64 value in arith.constant to 32-bit value.
+
+    Args:
+        module: The mlir module to analyze
+        ctx: The mlir ctx
+
+    Returns:
+        None
+    """
+
+    def walk_callback(op):
+        if op.name == "arith.constant":
+            op_type = op.result.type
+            if isinstance(op_type, ir.IntegerType) and op_type.width == 64:
+                with ctx, op.location, ir.InsertionPoint.at_block_begin(module.body):
+                    const_value = op.attributes["value"].value
+                    type = _get_cononical_type(const_value)
+                    if type is None:
+                        raise RuntimeError(f"Cannot support i64 value {const_value} in arith.constant.")
+
+                    new_const = arith.constant(type, const_value)
+                    op.result.replace_all_uses_with(new_const)
+                    op.erase()
+
+        return ir.WalkResult.ADVANCE
+
+    module.operation.walk(walk_callback)
diff --git a/third_party/aipu/include/Passes/Passes.td b/third_party/aipu/include/Passes/Passes.td
index 8d9b3437a..52fd9b625 100644
--- a/third_party/aipu/include/Passes/Passes.td
+++ b/third_party/aipu/include/Passes/Passes.td
@@ -9,4 +9,24 @@ def AIPUConvertBoolArg2I8 : Pass<"convert-bool-arg-to-i8", "func::FuncOp"> {
     let dependentDialects = ["func::FuncDialect", "arith::ArithDialect"];
 }
 
+// The following code snippet is from:
+// https://chromium.googlesource.com/external/github.com/llvm/llvm-project/+/refs/heads/upstream/users/matthias-springer/scf_bufferization_preprocessing
+// License: Apache 2.0
+// Modifications: None, used as-is
+def SCFLoopBufferizationPreprocessing
+    : Pass<"scf-loop-bufferization-preprocessing"> {
+  let summary = "Preprocess loops before One-Shot Bufferize";
+  let description = [{
+    Preprocess `scf.for` loops before running One-Shot Bufferize to support
+    loops where a yielded tensor is not equivalent to the respective iter_arg.
+    Such IR is currently not supported by One-Shot Bufferize.
+    This pass inserts a `bufferization.materialize_in_destination` op for every
+    yielded tensor, such that the yielded value is guaranteed to materialize in
+    the future buffer of the iter_arg; this is done by copying the tensor
+    contents into the iter_arg buffer. Such memcpys are a no-op in case the
+    tensor contents already materialize in the iter_arg buffer.
+  }];
+  let dependentDialects = ["bufferization::BufferizationDialect",
+                           "scf::SCFDialect"];
+}
 #endif // MLIR_DIALECT_AIPU_PASSES
diff --git a/third_party/aipu/language/aipu/libdevice.py b/third_party/aipu/language/aipu/libdevice.py
index 4f71cba5d..74d1b7e0d 100644
--- a/third_party/aipu/language/aipu/libdevice.py
+++ b/third_party/aipu/language/aipu/libdevice.py
@@ -41,25 +41,10 @@ def rcp64h(arg0):
     ...
 
 
-def rsqrt(arg0):
-    ...
-
-
 def ceil(arg0):
     ...
 
 
-@core.extern
-def trunc(arg0, _builder=None):
-    return core.extern_elementwise("trunc", "", [arg0], {
-        (core.dtype("fp32"), ): ("__truncf", core.dtype("fp32")),
-    }, is_pure=True, _builder=_builder)
-
-
-def exp2(arg0):
-    ...
-
-
 def saturatef(arg0):
     ...
 
@@ -84,17 +69,6 @@ def fast_dividef(arg0, arg1):
     ...
 
 
-def div_rn(arg0, arg1):
-    ...
-
-
-@core.extern
-def div_rz(arg0, arg1, _builder=None):
-    return core.extern_elementwise("div_rz", "", [arg0, arg1], {
-        (core.dtype("fp32"), core.dtype("fp32")): ("__div_rz", core.dtype("fp32")),
-    }, is_pure=True, _builder=_builder)
-
-
 def div_rd(arg0, arg1):
     ...
 
@@ -547,14 +521,6 @@ def copysign(arg0, arg1):
     ...
 
 
-def finitef(arg0):
-    ...
-
-
-def isinf(arg0):
-    ...
-
-
 def nextafter(arg0, arg1):
     ...
 
@@ -743,10 +709,6 @@ def scalbn(arg0, arg1):
     ...
 
 
-def fmod(arg0, arg1):
-    ...
-
-
 def remainder(arg0, arg1):
     ...
 
@@ -809,5 +771,14 @@ def binary_op(arg0, arg1, _builder=None):
 
 fmod = create_binary_op_wrapper("fmod", ["fp32", "fp16"])
 pow = create_binary_op_wrapper("powf", ["fp32", "fp16"])
+div_rz = create_binary_op_wrapper("div_rz", ["fp32", "fp16"])
+div_rn = create_binary_op_wrapper("div_rn", ["fp32", "fp16"])
+fmod = create_binary_op_wrapper("fmod", ["fp32", "fp16"])
+
 tanh = create_unary_op_wrapper("tanh", ["fp32", "fp16"])
 erf = create_unary_op_wrapper("erf", ["fp32", "fp16"])
+trunc = create_unary_op_wrapper("trunc", ["fp32", "fp16"])
+exp2 = create_unary_op_wrapper("exp2", ["fp32", "fp16"])
+finitef = create_unary_op_wrapper("isfinite", ["fp32", "fp16"])
+isinf = create_unary_op_wrapper("isinf", ["fp32", "fp16"])
+rsqrt = create_unary_op_wrapper("rsqrt", ["fp32", "fp16"])
diff --git a/third_party/aipu/lib/Passes/CMakeLists.txt b/third_party/aipu/lib/Passes/CMakeLists.txt
index 75308e004..f305a1b14 100644
--- a/third_party/aipu/lib/Passes/CMakeLists.txt
+++ b/third_party/aipu/lib/Passes/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_mlir_dialect_library(MLIRAIPUPasses
     ConvertBoolArg2I8.cpp
+    SCFLoopBufferizationPreprocess.cpp
 
     DEPENDS
     MLIRAIPUPassesIncGen
diff --git a/third_party/aipu/lib/Passes/SCFLoopBufferizationPreprocess.cpp b/third_party/aipu/lib/Passes/SCFLoopBufferizationPreprocess.cpp
new file mode 100644
index 000000000..0be9a411b
--- /dev/null
+++ b/third_party/aipu/lib/Passes/SCFLoopBufferizationPreprocess.cpp
@@ -0,0 +1,56 @@
+// The following code snippet is from:
+// https://chromium.googlesource.com/external/github.com/llvm/llvm-project/+/refs/heads/upstream/users/matthias-springer/scf_bufferization_preprocessing
+// License: Apache 2.0
+// Modifications:aipu namespace, add condition `if(operand.get() != bbArg)`
+
+#include "Passes/Passes.h"
+
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/Patterns.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+
+namespace aipu {
+
+#define GEN_PASS_DEF_SCFLOOPBUFFERIZATIONPREPROCESSING
+#include "Passes/Passes.h.inc"
+
+using namespace mlir::scf;
+
+struct SCFLoopBufferizationPreprocessingPass
+    : public impl::SCFLoopBufferizationPreprocessingBase<
+          SCFLoopBufferizationPreprocessingPass> {
+  void runOnOperation() override {
+    OpBuilder builder(getOperation()->getContext());
+    getOperation()->walk([&](scf::YieldOp yieldOp) {
+      builder.setInsertionPoint(yieldOp);
+      // TODO: Support scf.while.
+      auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
+      if (!forOp)
+        return WalkResult::skip();
+      for (OpOperand &operand : yieldOp->getOpOperands()) {
+        auto tensorType = dyn_cast<TensorType>(operand.get().getType());
+        if (!tensorType)
+          continue;
+        auto bbArg = forOp.getRegionIterArgs()[operand.getOperandNumber()];
+        if (operand.get() != bbArg) {
+          Value materialized =
+              builder
+                  .create<bufferization::MaterializeInDestinationOp>(
+                      yieldOp.getLoc(), tensorType, operand.get(), bbArg)
+                  .getResult();
+          operand.set(materialized);
+        }
+      }
+      return WalkResult::advance();
+    });
+  }
+};
+
+} // namespace aipu
+
+} // namespace mlir
diff --git a/third_party/aipu/lib/main.cpp b/third_party/aipu/lib/main.cpp
index 14fb353be..fbc31bae3 100644
--- a/third_party/aipu/lib/main.cpp
+++ b/third_party/aipu/lib/main.cpp
@@ -6,7 +6,10 @@
 using namespace mlir::aipu;
 
 void init_aipu_passes(pybind11::module &&m) {
-  m.def("register_all_passes", []() { registerAIPUConvertBoolArg2I8(); });
+  m.def("register_all_passes", []() {
+    registerAIPUConvertBoolArg2I8();
+    registerSCFLoopBufferizationPreprocessing();
+  });
 }
 
 PYBIND11_MODULE(libaipu_interface, m) {
diff --git a/third_party/aipu/python/test/test_codegen_dma.py b/third_party/aipu/python/test/test_codegen_dma.py
index b651dc8eb..8f26eb2cc 100644
--- a/third_party/aipu/python/test/test_codegen_dma.py
+++ b/third_party/aipu/python/test/test_codegen_dma.py
@@ -1,49 +1,92 @@
+import triton
 import numpy as np
 from mlir import ir
 from tvm import aipu
 from tvm.aipu import testing
+from tvm.aipu.utils import rand
 from triton.backends.aipu.codegen import AIPUModule, CodeGenerator
+from triton.backends.aipu import transform
 
 
-def test_dma():
+def get_vadd_mlir():
     mod_str = """
-  func.func @dma_ops(%inp: memref<16xf32>, %out: memref<16xf32>) {
+module {
+  func.func @add_kernel(%arg0: memref<*xf32> {tt.divisibility = 16 : i32}, %arg1: memref<*xf32> {tt.divisibility = 16 : i32}, %arg2: memref<*xf32> {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32) {
     %c0 = arith.constant 0 : index
+    %c1024_i32 = arith.constant 1024 : i32
+    %c1024 = arith.constant 1024 : index
+    %0 = arith.muli %arg7, %c1024_i32 : i32
+    %1 = arith.index_cast %0 : i32 to index
+    %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [%1], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>>
+    %2 = arith.addi %1, %c1024 : index
+    %3 = arith.index_cast %arg3 : i32 to index
+    %4 = arith.minsi %2, %3 : index
+    %5 = arith.maxsi %4, %1 : index
+    %6 = arith.subi %5, %1 : index
+    %alloc = memref.alloc() : memref<1024xf32>
+    %subview = memref.subview %reinterpret_cast[0] [%6] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
+    %subview_0 = memref.subview %alloc[0] [%6] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
+    %alloc_1 = memref.alloc() : memref<1xi32, 11 : i32>
+    memref.dma_start %subview[%c0], %subview_0[%c0], %6, %alloc_1[%c0] : memref<?xf32, strided<[1], offset: ?>>, memref<?xf32, strided<[1]>>, memref<1xi32, 11 : i32>
+    memref.dma_wait %alloc_1[%c0], %6 : memref<1xi32, 11 : i32>
+    %reinterpret_cast_2 = memref.reinterpret_cast %arg1 to offset: [%1], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>>
+    %alloc_3 = memref.alloc() : memref<1024xf32>
+    %subview_4 = memref.subview %reinterpret_cast_2[0] [%6] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
+    %subview_5 = memref.subview %alloc_3[0] [%6] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
+    %alloc_6 = memref.alloc() : memref<1xi32, 11 : i32>
+    memref.dma_start %subview_4[%c0], %subview_5[%c0], %6, %alloc_6[%c0] : memref<?xf32, strided<[1], offset: ?>>, memref<?xf32, strided<[1]>>, memref<1xi32, 11 : i32>
+    memref.dma_wait %alloc_6[%c0], %6 : memref<1xi32, 11 : i32>
+    %alloc_7 = memref.alloc() {alignment = 64 : i64} : memref<1024xf32>
+    %c0_8 = arith.constant 0 : index
+    %c1024_9 = arith.constant 1024 : index
+    %c8 = arith.constant 8 : index
+    scf.for %arg10 = %c0_8 to %c1024_9 step %c8 {
+      %cst = arith.constant 0.000000e+00 : f32
+      %7 = vector.transfer_read %alloc[%arg10], %cst : memref<1024xf32>, vector<8xf32>
+      %cst_13 = arith.constant 0.000000e+00 : f32
+      %8 = vector.transfer_read %alloc_3[%arg10], %cst_13 : memref<1024xf32>, vector<8xf32>
+      %9 = arith.addf %7, %8 : vector<8xf32>
+      vector.transfer_write %9, %alloc_7[%arg10] : vector<8xf32>, memref<1024xf32>
+    }
+    %reinterpret_cast_10 = memref.reinterpret_cast %arg2 to offset: [%1], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>>
+    %subview_11 = memref.subview %alloc_7[0] [%6] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
+    %subview_12 = memref.subview %reinterpret_cast_10[0] [%6] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
+    memref.copy %subview_11, %subview_12 : memref<?xf32, strided<[1]>> to memref<?xf32, strided<[1], offset: ?>>
+    return
+  }
+}
+"""
+    return mod_str
 
-    %lsram = memref.alloc() : memref<16 x f32, affine_map<(d0) -> (d0)>, 4>
-    %tag = memref.alloc() : memref<1 x i32, affine_map<(d0) -> (d0)>, 11>
-
-    %num_elements = arith.constant 16 : index
 
-    memref.dma_start %inp[%c0], %lsram[%c0], %num_elements, %tag[%c0] : memref<16 x f32>, memref<16 x f32, 4>, memref<1 x i32, 11>
-    memref.dma_start %lsram[%c0], %out[%c0], %num_elements, %tag[%c0] : memref<16 x f32, 4>, memref<16 x f32>, memref<1 x i32, 11>
+def test_dma():
+    mod_str = get_vadd_mlir()
+    ctx = ir.Context()
+    mod = ir.Module.parse(mod_str, ctx)
+    transform.binding_tid(mod, ctx)
 
-    memref.dma_wait %tag[%c0], %num_elements : memref<1 x i32, 11>
-    return
-  }"""
-    mod = AIPUModule(ir.Module.parse(mod_str, ir.Context()))
-    cg = CodeGenerator(mod)
+    cg = CodeGenerator(AIPUModule(mod))
     cg.mod.walk_mod(cg.dispatch)
 
     bm = aipu.tir.BuildManager()
     ex = bm.build(cg.prim_func)
-    print(ex.c_code)
-    """
-    __kernel void dma_ops(__global float* var_3, __global float* var_5) {
-      *addr_of_event_state() = 0;
-      __lsram float buf[16];
-      int cse_var_1 = (16 * 4);
-      AsyncDmaDirect_kGlobal_to_kLsram((int)buf, (int)var_3, cse_var_1, cse_var_1, cse_var_1, cse_var_1, alloc_event());
-      AsyncDmaDirect_kLsram_to_kGlobal((int)var_5, (int)buf, cse_var_1, cse_var_1, cse_var_1, cse_var_1, alloc_event());
-      wait_events((1 << alloc_event()));
-      barrier(CLK_LOCAL_MEM_FENCE);return;
-      barrier(CLK_LOCAL_MEM_FENCE);
-    }
-    """
-    a = np.array(list(range(16)), dtype=np.float32)
-    aipu_out = np.empty((16, ), dtype=np.float32)
-    ex(a, aipu_out)
-    testing.assert_allclose(a, aipu_out)
+
+    size = 4097
+    BLOCKSIZE = 1024
+    dtype = "float32"
+    a = rand(size, dtype)
+    b = rand(size, dtype)
+    aipu_out = np.empty(size, dtype=dtype)
+
+    gridX = triton.cdiv(size, BLOCKSIZE)
+    np_args = [a, b, aipu_out, size]
+    tail_args = [gridX, 1, 1, 0, 0, 0]
+    tec_num = 4
+    for i in range((gridX + tec_num - 1) // tec_num):
+        tail_args[3] = i
+        ex(*(np_args + tail_args))
+
+    testing.assert_allclose(aipu_out, a + b)
 
 
 if __name__ == "__main__":
diff --git a/third_party/aipu/test/convert_bool_arg_to_i8.mlir b/third_party/aipu/test/convert_bool_arg_to_i8.mlir
new file mode 100644
index 000000000..4764db751
--- /dev/null
+++ b/third_party/aipu/test/convert_bool_arg_to_i8.mlir
@@ -0,0 +1,9 @@
+// RUN: aipu-opt %s --convert-bool-arg-to-i8 | FileCheck %s
+
+// CHECK: func.func @test_func(%[[ARG0:.*]]: i8, %[[ARG1:.*]]: i32) -> i1
+func.func @test_func(%arg0: i1, %arg1: i32) -> i1 {
+  // CHECK: %0 = arith.trunci %arg0 : i8 to i1
+  %true = arith.constant true
+  %1 = arith.ori %arg0, %true : i1
+  return %1 : i1
+}
diff --git a/third_party/aipu/test/lit.cfg b/third_party/aipu/test/lit.cfg
new file mode 100644
index 000000000..d3a4b4831
--- /dev/null
+++ b/third_party/aipu/test/lit.cfg
@@ -0,0 +1,4 @@
+import lit.formats
+config.name = "AIPU Test"
+config.test_format = lit.formats.ShTest(True)
+config.suffixes = [".mlir"]
diff --git a/third_party/aipu/test/one-shot-bufferize-preprocessing.mlir b/third_party/aipu/test/one-shot-bufferize-preprocessing.mlir
new file mode 100644
index 000000000..dcc7b8dd0
--- /dev/null
+++ b/third_party/aipu/test/one-shot-bufferize-preprocessing.mlir
@@ -0,0 +1,26 @@
+// The following code snippet is from:
+// https://chromium.googlesource.com/external/github.com/llvm/llvm-project/+/refs/heads/upstream/users/matthias-springer/scf_bufferization_preprocessing
+// License: Apache 2.0
+// Modifications: mlir-opt -> aipu-opt
+
+// RUN: aipu-opt %s -scf-loop-bufferization-preprocessing -one-shot-bufferize="bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map" -canonicalize | FileCheck %s
+// CHECK-LABEL: func @conflict_in_loop(
+//  CHECK-SAME:     %[[A:.*]]: memref<10xf32>
+func.func @conflict_in_loop(%A: tensor<10xf32>, %f: f32, %idx: index, %lb: index, %ub: index, %step: index) -> f32 {
+  // CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
+  %r = scf.for %i = %lb to %ub step %step iter_args(%tA = %A) -> (tensor<10xf32>) {
+    // CHECK: %[[alloc:.*]] = memref.alloc()
+    // CHECK: memref.copy %[[A]], %[[alloc]]
+    // CHECK: memref.store %{{.*}}, %[[alloc]]
+    %0 = tensor.insert %f into %tA[%i] : tensor<10xf32>
+    // CHECK: %[[read:.*]] = memref.load %[[A]]
+    %read = tensor.extract %tA[%idx] : tensor<10xf32>
+    // CHECK: vector.print %[[read]]
+    vector.print %read : f32
+    // CHECK: memref.copy %[[alloc]], %[[A]]
+    scf.yield %0 : tensor<10xf32>
+  }
+  // CHECK: memref.load %[[A]]
+  %f0 = tensor.extract %r[%step] : tensor<10xf32>
+  return %f0 : f32
+}
diff --git a/third_party/aipu/tools/CMakeLists.txt b/third_party/aipu/tools/CMakeLists.txt
new file mode 100644
index 000000000..8e5c677e5
--- /dev/null
+++ b/third_party/aipu/tools/CMakeLists.txt
@@ -0,0 +1 @@
+add_subdirectory(aipu-opt)
diff --git a/third_party/aipu/tools/aipu-opt/CMakeLists.txt b/third_party/aipu/tools/aipu-opt/CMakeLists.txt
new file mode 100644
index 000000000..7d509a2ac
--- /dev/null
+++ b/third_party/aipu/tools/aipu-opt/CMakeLists.txt
@@ -0,0 +1,18 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+
+add_llvm_executable(aipu-opt aipu-opt.cpp PARTIAL_SOURCES_INTENDED)
+
+llvm_update_compile_flags(aipu-opt)
+target_link_libraries(aipu-opt PRIVATE
+  ${dialect_libs}
+  ${conversion_libs}
+  # MLIR core
+  MLIROptLib
+  MLIRPass
+  MLIRTransforms
+  # AIPU pass
+  MLIRAIPUPasses
+)
+
+mlir_check_all_link_libraries(aipu-opt)
diff --git a/third_party/aipu/tools/aipu-opt/aipu-opt.cpp b/third_party/aipu/tools/aipu-opt/aipu-opt.cpp
new file mode 100644
index 000000000..6bb9ea386
--- /dev/null
+++ b/third_party/aipu/tools/aipu-opt/aipu-opt.cpp
@@ -0,0 +1,15 @@
+#include "Passes/Passes.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/InitAllPasses.h"
+#include "mlir/Tools/mlir-opt/MlirOptMain.h"
+
+int main(int argc, char **argv) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllPasses();
+  mlir::registerAllDialects(registry);
+  mlir::aipu::registerAIPUConvertBoolArg2I8();
+  mlir::aipu::registerSCFLoopBufferizationPreprocessing();
+
+  return mlir::asMainReturnCode(
+      mlir::MlirOptMain(argc, argv, "AIPU optimizer driver\n", registry));
+}

From e0e9829518ab35277fec9648a8b7c133d57d0eaf Mon Sep 17 00:00:00 2001
From: Lucien0 <16538059+Lucien0@users.noreply.github.com>
Date: Fri, 20 Jun 2025 17:52:24 +0800
Subject: [PATCH 2/3] [Backend] Update codegen for new flir

---
 third_party/aipu/backend/codegen.py | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/third_party/aipu/backend/codegen.py b/third_party/aipu/backend/codegen.py
index 627a856b9..9a7e98e9b 100644
--- a/third_party/aipu/backend/codegen.py
+++ b/third_party/aipu/backend/codegen.py
@@ -248,17 +248,7 @@ def exit_scope(self):
 
     def dispatch(self, op, stage):
         op_name = "func.func" if isinstance(op, func.FuncOp) else op.name
-        # Extern Call
-        if "extern_call" in op.attributes:
-            func_name = op.attributes["extern_call"].value
-            if func_name == "fmod":
-                fmod = lambda x, y: T.call_extern(_get_type(op.result), "fmod", x, y)
-                self.gen_binary(op, fmod)
-            if func_name == "div_rz":
-                self.ib.emit(tir.call_extern("void", "__vset_rounding_mode_rtz"))
-                self.gen_binary(op, T.Div)
-                self.ib.emit(tir.call_extern("void", "__vset_rounding_mode_rtn"))
-            return
+
         # Memref Dialect
         if op_name == "memref.reinterpret_cast":
             self.gen_memref_reinterpret_cast(op)
@@ -356,6 +346,14 @@ def dispatch(self, op, stage):
         elif op_name == "math.trunc":
             trunc = lambda x: T.call_extern(_get_type(op.result), "trunc", x)
             self.gen_unary(op, trunc)
+        # MathExt Dialect
+        elif op_name == "mathext.fmod":
+            fmod = lambda x, y: T.call_extern(_get_type(op.result), "fmod", x, y)
+            self.gen_binary(op, fmod)
+        elif op_name == "mathext.div_rz":
+            self.ib.emit(tir.call_extern("void", "__vset_rounding_mode_rtz"))
+            self.gen_binary(op, T.Div)
+            self.ib.emit(tir.call_extern("void", "__vset_rounding_mode_rtn"))
         # Func Dialect
         elif op_name == "func.return":
             self.gen_func_return(op)

From 506904db8a93f6c02328ecafbd0ac1f453abfb5c Mon Sep 17 00:00:00 2001
From: zhengyang <zhengyang@baai.ac.cn>
Date: Fri, 20 Jun 2025 20:19:23 +0800
Subject: [PATCH 3/3] [Backend] Update codegen for new flir commit

---
 python/setup_tools/utils/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/setup_tools/utils/__init__.py b/python/setup_tools/utils/__init__.py
index dd4f98db3..103b96178 100644
--- a/python/setup_tools/utils/__init__.py
+++ b/python/setup_tools/utils/__init__.py
@@ -20,7 +20,7 @@ class FlagTreeBackend:
     FlagTreeBackend(name="cambricon", url="https://github.com/Cambricon/triton-linalg.git",
                     tag="00f51c2e48a943922f86f03d58e29f514def646d"),
     FlagTreeBackend(name="flir", url="git@github.com:FlagTree/flir.git",
-                    tag="318ed13e396d4d0ed84773975c8507c6e3f0275d"),
+                    tag="243690fb1c8b9f032c1f938271414831a6cfe406"),
     FlagTreeBackend(
         name="ascend",
         url="https://gitee.com/ascend/triton-ascend.git",