Issue

无

背景

为优化间接访存场景的兜底实现，适配 A5 环境需求，需新增专门的间接访存接口，并通过开关控制不同场景下的实现路径：A5 使能时使用高效接口，未使能时保持原有逻辑。

描述

1. TritonToUnstructure 新增compile-on-a5使能开关
2. 新增 TT_Indirect_Load 接口
3. 优化当前间接访存兜底路径：
   当 compile-on-a5=True 时，使用 func:FuncOp(IndirectLoad)。
   当 compile-on-a5=False 时，保留原有兜底方案（For 循环嵌套与 extract_slice 等操作）。

测试

测试脚本：pytest -sv ascend/examples/pytest_ut/test_ldst.py::test_ldst_indirect_05

• compile-on-a5=True 时，Adapter IR 中出现 call @triton_indirect_load 调用，直接下发接口完成间接访存

# map = affine_map<(d0) -> (d0)>
module {
  func.func private @triton_indirect_load(memref<?xf32>, tensor<8x16xi64>) -> tensor<8x16xf32>
  func.func @triton_ldst_indirect_05_kernel(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf32> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg3: memref<?xi64> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf32> {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "aiv"} {
    %c16 = arith.constant 16 : index
    %c8_i32 = arith.constant 8 : i32
    %0 = arith.muli %arg9, %c8_i32 : i32
    %1 = tensor.empty() : tensor<8xi32>
    %2 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%1 : tensor<8xi32>) {
    ^bb0(%out: i32):
      %14 = linalg.index 0 : index
      %15 = arith.index_cast %14 : index to i32
      linalg.yield %15 : i32
    } -> tensor<8xi32>
    %reinterpret_cast = memref.reinterpret_cast %arg3 to offset: [0], sizes: [16], strides: [1] : memref<?xi64> to memref<16xi64, strided<[1]>>
    %alloc = memref.alloc() : memref<16xi64>
    memref.copy %reinterpret_cast, %alloc : memref<16xi64, strided<[1]>> to memref<16xi64>
    %3 = bufferization.to_tensor %alloc restrict writable : memref<16xi64>
    %expanded = tensor.expand_shape %2 [[0, 1]] output_shape [8, 1] : tensor<8xi32> into tensor<8x1xi32>
    %4 = tensor.empty() : tensor<8x1xi32>
    %5 = linalg.fill ins(%arg5 : i32) outs(%4 : tensor<8x1xi32>) -> tensor<8x1xi32>
    %6 = arith.muli %expanded, %5 : tensor<8x1xi32>
    %7 = arith.extsi %6 : tensor<8x1xi32> to tensor<8x1xi64>
    %8 = tensor.empty() : tensor<8x16xi64>
    %collapsed = tensor.collapse_shape %7 [[0, 1]] : tensor<8x1xi64> into tensor<8xi64>
    %broadcasted = linalg.broadcast ins(%collapsed : tensor<8xi64>) outs(%8 : tensor<8x16xi64>) dimensions = [1] 
    %broadcasted_0 = linalg.broadcast ins(%3 : tensor<16xi64>) outs(%8 : tensor<8x16xi64>) dimensions = [0] 
    %9 = arith.addi %broadcasted, %broadcasted_0 : tensor<8x16xi64>
    %10 = call @triton_indirect_load(%arg4, %9) : (memref<?xf32>, tensor<8x16xi64>) -> tensor<8x16xf32>
    %11 = math.exp %10 : tensor<8x16xf32>
    %12 = arith.index_cast %0 : i32 to index
    %13 = arith.muli %12, %c16 : index
    %reinterpret_cast_1 = memref.reinterpret_cast %arg2 to offset: [%13], sizes: [8, 16], strides: [16, 1] : memref<?xf32> to memref<8x16xf32, strided<[16, 1], offset: ?>>
    bufferization.materialize_in_destination %11 in writable %reinterpret_cast_1 : (tensor<8x16xf32>, memref<8x16xf32, strided<[16, 1], offset: ?>>) -> ()
    return
  }
}

• compile-on-a5=False 时，AdapterIR 仍然使用原有的间接访存实现，通过 scf.for 循环嵌套和 tensor.extract 完成数据提取。

module {
  func.func @triton_ldst_indirect_05_kernel(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf32> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg3: memref<?xi64> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf32> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "aiv"} {
    %c16 = arith.constant 16 : index
    %c8 = arith.constant 8 : index
    %c1 = arith.constant 1 : index
    %c0 = arith.constant 0 : index
    %c8_i32 = arith.constant 8 : i32
    %0 = arith.muli %arg9, %c8_i32 : i32
    %reinterpret_cast = memref.reinterpret_cast %arg3 to offset: [0], sizes: [16], strides: [1] : memref<?xi64> to memref<16xi64, strided<[1]>>
    %alloc = memref.alloc() : memref<16xi64>
    memref.copy %reinterpret_cast, %alloc : memref<16xi64, strided<[1]>> to memref<16xi64>
    %1 = bufferization.to_tensor %alloc restrict writable : memref<16xi64>
    %2 = tensor.empty() : tensor<8x16xf32>
    %3 = scf.for %arg12 = %c0 to %c8 step %c1 iter_args(%arg13 = %2) -> (tensor<8x16xf32>) {
      %7 = scf.for %arg14 = %c0 to %c16 step %c1 iter_args(%arg15 = %arg13) -> (tensor<8x16xf32>) {
        %8 = arith.index_cast %arg12 : index to i32
        %9 = arith.muli %8, %arg5 : i32
        %10 = arith.extsi %9 : i32 to i64
        %extracted = tensor.extract %1[%arg14] {DiscreteMemAccess} : tensor<16xi64>
        %11 = arith.addi %10, %extracted : i64
        %12 = arith.index_cast %11 : i64 to index
        %reinterpret_cast_1 = memref.reinterpret_cast %arg4 to offset: [%12], sizes: [1], strides: [1] : memref<?xf32> to memref<1xf32, strided<[1], offset: ?>>
        %13 = memref.load %reinterpret_cast_1[%c0] : memref<1xf32, strided<[1], offset: ?>>
        %inserted = tensor.insert %13 into %arg15[%arg12, %arg14] : tensor<8x16xf32>
        scf.yield {DiscreteMemAccess} %inserted : tensor<8x16xf32>
      } {ExtractedLoadOrStore}
      scf.yield %7 : tensor<8x16xf32>
    } {ExtractedLoadOrStore}
    %4 = math.exp %3 : tensor<8x16xf32>
    %5 = arith.index_cast %0 : i32 to index
    %6 = arith.muli %5, %c16 : index
    %reinterpret_cast_0 = memref.reinterpret_cast %arg2 to offset: [%6], sizes: [8, 16], strides: [16, 1] : memref<?xf32> to memref<8x16xf32, strided<[16, 1], offset: ?>>
    bufferization.materialize_in_destination %4 in writable %reinterpret_cast_0 : (tensor<8x16xf32>, memref<8x16xf32, strided<[16, 1], offset: ?>>) -> ()
    return
  }
}