[P1] Stride needed (additional operands) for the GEP

[Jackcuii]

By now, we have three types of GEPs. (1 opr, 2 oprs, 3 oprs)

operations:
                - opcode: "GEP"
                  id: 119
                  time_step: 17
                  invalid_iterations: 1
                  src_operands:
                    - operand: "#16"
                      color: "RED"
                    - operand: "$17"
                      color: "RED"
                    - operand: "$2"
                      color: "RED"
                  dst_operands:
                    - operand: "SOUTH"
                      color: "RED"

For A[i][j], the GEP will be GEP(base(A), i, j), however, we need the stride/shape of A to get the correct result.

It could be either an individual operand or it is calculated out of the GEP.

[tancheng]

GEP(base(A), i, j), however, we need the stride/shape of A to get the correct result.

GEMINI says:

1. For Static Shapes (Automatic)
If the LLVM IR contains a multi-index GEP on a structured type, you simply map it to an llvm.getelementptr in MLIR. The MLIR LLVMToLLVMIRTranslation (and the inverse) handles the offset math automatically because the type carries the stride.

The GEP for 2D/3D/4D/5D needs additional instructions/operations depending on the static/dynamic shape info. I hope there is existing pass can do that for us, if no, we need to insert appropriate shift/multiply operations for it... @ShangkunLi Can you help investigate the way to support this?

[ShangkunLi]

GEP(base(A), i, j), however, we need the stride/shape of A to get the correct result.

GEMINI says:

1. For Static Shapes (Automatic)
If the LLVM IR contains a multi-index GEP on a structured type, you simply map it to an llvm.getelementptr in MLIR. The MLIR LLVMToLLVMIRTranslation (and the inverse) handles the offset math automatically because the type carries the stride.

The GEP for 2D/3D/4D/5D needs additional instructions/operations depending on the static/dynamic shape info. I hope there is existing pass can do that for us, if no, we need to insert appropriate shift/multiply operations for it... @ShangkunLi Can you help investigate the way to support this?

Hi~ @tancheng @Jackcuii @n0thingNoob, after brief investigations, the LLVM dialect uses an elem_type TypeAttr to denote the strides for the llvm.getelementptr op.

https://github.com/llvm/llvm-project/blob/9f3a9ea6ae689fa0fbd0137999834915cc908e5c/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td#L296

For example,

%16 = llvm.getelementptr inbounds %arg5[%12, %14] : (!llvm.ptr, i64, i64) -> !llvm.ptr, !llvm.array<1100 x i32>

This means this llvm.getelementptr op uses sizeof(!llvm.array<1100 x i32>) = 1100 × 4 = 4400 bytes as a stride for the first dimension (i.e., %12), and sizeof(i32) = 4 bytes for the second dimension (i.e., %14).

I plan to introduce a I32ArrayAttr:$strides attribute in the neura.gep op. Like,

// Defines a pointer computation operation.
def Neura_GEP : Op<NeuraDialect, "gep", [AttrSizedOperandSegments]> {
  let summary = "Pointer computation using offset indices";
  let description = [{
    Computes a pointer address from a base pointer and multi-dimensional indices.
    The `strides` attribute specifies the byte stride for each index level.

    Address = base + sum_i(index[i] * strides[i])

    Example:
      // 1D: A[i] where A is float*
      %0 = neura.gep %base[%i] {strides = [4]} : ...

      // 2D: A[i][j] where A is int[N][1200]
      %0 = neura.gep %base[%i, %j] {strides = [4800, 4]} : ...
  }];
  let arguments = (ins Optional<AnyType>:$base, Variadic<AnyType>:$indices,
                       I32ArrayAttr:$strides);
  let results = (outs AnyType:$result);
}

The $strides denotes the exact strides of the incremental stride in bytes for each index dimension. WDYT?

[tancheng]

Sounds good for static offset. What if the tensor shape (or stride) is not known at compilation time?

[tancheng]

One more question, if the stride is an attribute, how should it be passed in DFG or codegen? Should it be another operand?

[ShangkunLi]

Sounds good for static offset. What if the tensor shape (or stride) is not known at compilation time?

For dynamic shape, it is only allowed in the C99 standard, and clang can transform it into explicit mul + gep.

For example,

void kernel(int n, int m, int A[n][m], int B[n][m]) {
  for (int i = 0; i < n; i++) {
    B[i][0] = A[i][0] * 2;
  }
}

clang transforms it into the following ir,

define dso_local void @kernel(i32 noundef %n, i32 noundef %m, ptr nocapture noundef readonly %A, ptr nocapture noundef writeonly %B) local_unnamed_addr #0 {
entry:
  %0 = zext i32 %m to i64
  %cmp11 = icmp sgt i32 %n, 0
  br i1 %cmp11, label %for.body.preheader, label %for.cond.cleanup

for.body.preheader:                               ; preds = %entry
  %wide.trip.count = zext nneg i32 %n to i64
  br label %for.body

for.cond.cleanup:                                 ; preds = %for.body, %entry
  ret void

for.body:                                         ; preds = %for.body.preheader, %for.body
  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.body ]
  %1 = mul nuw nsw i64 %indvars.iv, %0
  %arrayidx = getelementptr inbounds nuw i32, ptr %A, i64 %1
  %2 = load i32, ptr %arrayidx, align 4, !tbaa !5
  %mul = shl nsw i32 %2, 1
  %arrayidx3 = getelementptr inbounds nuw i32, ptr %B, i64 %1
  store i32 %mul, ptr %arrayidx3, align 4, !tbaa !5
  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !9
}

[ShangkunLi]

One more question, if the stride is an attribute, how should it be passed in DFG or codegen? Should it be another operand?

The stride attribute is a configuration for the gep operation (i.e., used to generate the configuration for the gep FU in the codegen pass); it doesn't need to flow through DFG edges.

[tancheng]

Sounds good for static offset. What if the tensor shape (or stride) is not known at compilation time?

For dynamic shape, it is only allowed in the C99 standard, and clang can transform it into explicit mul + gep.

For example,

void kernel(int n, int m, int A[n][m], int B[n][m]) {
for (int i = 0; i < n; i++) {
B[i][0] = A[i][0] * 2;
}
}
clang transforms it into the following ir,

define dso_local void @kernel(i32 noundef %n, i32 noundef %m, ptr nocapture noundef readonly %A, ptr nocapture noundef writeonly %B) local_unnamed_addr #0 {
entry:
%0 = zext i32 %m to i64
%cmp11 = icmp sgt i32 %n, 0
br i1 %cmp11, label %for.body.preheader, label %for.cond.cleanup

for.body.preheader: ; preds = %entry
%wide.trip.count = zext nneg i32 %n to i64
br label %for.body

for.cond.cleanup: ; preds = %for.body, %entry
ret void

for.body: ; preds = %for.body.preheader, %for.body
%indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.body ]
%1 = mul nuw nsw i64 %indvars.iv, %0
%arrayidx = getelementptr inbounds nuw i32, ptr %A, i64 %1
%2 = load i32, ptr %arrayidx, align 4, !tbaa !5
%mul = shl nsw i32 %2, 1
%arrayidx3 = getelementptr inbounds nuw i32, ptr %B, i64 %1
store i32 %mul, ptr %arrayidx3, align 4, !tbaa !5
%indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
%exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !9
}

This llvm IR can also be smoothly raised up to MLIR llvm, right?

[tancheng]

One more question, if the stride is an attribute, how should it be passed in DFG or codegen? Should it be another operand?

The stride attribute is a configuration for the gep operation (i.e., used to generate the configuration for the gep FU in the codegen pass); it doesn't need to flow through DFG edges.

I don't actually like using cmd to configure certain FU before execution, we'd better minimize such cases as much as possible.. (because then codegen should also generate some CMD related stuff)
Can you think about is it possible to use constant as an operand (so then we just need to put some constant data into the const_queue)? The issue maybe how many constant a GEP needs.