YQL-16067 Implement LLVM for wide ToFlow.

3 files changed, 87 insertions, 44 deletions

diff --git a/ydb/library/yql/minikql/comp_nodes/mkql_flow.cpp b/ydb/library/yql/minikql/comp_nodes/mkql_flow.cpp
index 3505e54681..0a76ee8f57 100644
--- a/ydb/library/yql/minikql/comp_nodes/mkql_flow.cpp
+++ b/ydb/library/yql/minikql/comp_nodes/mkql_flow.cpp
@@ -282,60 +282,88 @@ private:
     IComputationNode* const Flow;
 };
 
-class TToWideFlowWrapper : public TWideFlowSourceComputationNode<TToWideFlowWrapper> {
-    typedef TWideFlowSourceComputationNode<TToWideFlowWrapper> TBaseComputation;
+class TToWideFlowWrapper : public TWideFlowSourceCodegeneratorNode<TToWideFlowWrapper> {
+using TBaseComputation = TWideFlowSourceCodegeneratorNode<TToWideFlowWrapper>;
 public:
     TToWideFlowWrapper(TComputationMutables& mutables, IComputationNode* stream, ui32 width)
         : TBaseComputation(mutables, EValueRepresentation::Any)
         , Stream(stream)
         , Width(width)
+        , TempStateIndex(std::exchange(mutables.CurValueIndex, mutables.CurValueIndex + Width))
     {}
 
-    EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const
-    {
-        auto& s = GetState(state, ctx);
-        auto status = s.StreamValue.WideFetch(s.Values.data(), Width);
-        switch (status) {
+    EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
+        if (!state.HasValue()) {
+            state = Stream->GetValue(ctx);
+        }
+
+        switch (const auto status = state.WideFetch(ctx.MutableValues.get() + TempStateIndex, Width)) {
         case NUdf::EFetchStatus::Finish:
             return EFetchResult::Finish;
         case NUdf::EFetchStatus::Yield:
             return EFetchResult::Yield;
         case NUdf::EFetchStatus::Ok:
-            for (ui32 i = 0; i < Width; ++i) {
-                if (output[i]) {
-                  *output[i] = std::move(s.Values[i]);
-                }
+            break;
+        }
+
+        for (auto i = 0U; i < Width; ++i) {
+            if (const auto out = output[i]) {
+                *out = std::move(ctx.MutableValues[TempStateIndex + i]);
             }
-            return EFetchResult::One;
         }
+        return EFetchResult::One;
     }
+#ifndef MKQL_DISABLE_CODEGEN
+    TGenerateResult DoGenGetValues(const TCodegenContext& ctx, Value* statePtr, BasicBlock*& block) const {
+        auto& context = ctx.Codegen->GetContext();
+        const auto valueType = Type::getInt128Ty(context);
+        const auto indexType = Type::getInt32Ty(context);
+        const auto values = GetElementPtrInst::CreateInBounds(valueType, ctx.GetMutables(), {ConstantInt::get(indexType, TempStateIndex)}, "values", &ctx.Func->getEntryBlock().back());
 
-private:
-    struct TState : public TComputationValue<TState> {
-        NUdf::TUnboxedValue StreamValue;
-        TVector<NUdf::TUnboxedValue> Values;
-
-        TState(TMemoryUsageInfo* memInfo, NUdf::TUnboxedValue&& streamValue, ui32 width)
-            : TComputationValue(memInfo)
-            , StreamValue(std::move(streamValue))
-            , Values(width)
-        {
-        }
-    };
+        const auto init = BasicBlock::Create(context, "init", ctx.Func);
+        const auto main = BasicBlock::Create(context, "main", ctx.Func);
+
+        const auto load = new LoadInst(valueType, statePtr, "load", block);
+        const auto state = PHINode::Create(load->getType(), 2U, "state", main);
+        state->addIncoming(load, block);
 
+        BranchInst::Create(init, main, IsInvalid(load, block), block);
+
+        block = init;
+
+        GetNodeValue(statePtr, Stream, ctx, block);
+
+        const auto save = new LoadInst(valueType, statePtr, "save", block);
+        state->addIncoming(save, block);
+        BranchInst::Create(main, block);
+
+        block = main;
+
+        const auto status = CallBoxedValueVirtualMethod<NUdf::TBoxedValueAccessor::EMethod::WideFetch>(indexType, state, ctx.Codegen, block, values, ConstantInt::get(indexType, Width));
+
+        const auto ok = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, status, ConstantInt::get(indexType, static_cast<ui32>(NUdf::EFetchStatus::Ok)), "ok", block);
+        const auto yield = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, status, ConstantInt::get(indexType, static_cast<ui32>(NUdf::EFetchStatus::Yield)), "yield", block);
+        const auto special = SelectInst::Create(yield, ConstantInt::get(indexType, static_cast<i32>(EFetchResult::Yield)), ConstantInt::get(indexType, static_cast<i32>(EFetchResult::Finish)), "special", block);
+        const auto complete = SelectInst::Create(ok, ConstantInt::get(indexType, static_cast<i32>(EFetchResult::One)), special, "complete", block);
+
+        TGettersList getters(Width);
+        for (auto i = 0U; i < getters.size(); ++i) {
+            getters[i] = [idx = TempStateIndex + i, values, valueType, indexType](const TCodegenContext& ctx, BasicBlock*& block) {
+                const auto valuePtr = GetElementPtrInst::CreateInBounds(valueType, ctx.GetMutables(), {ConstantInt::get(indexType, idx)}, (TString("ptr_") += ToString(idx)).c_str(), block);
+                return new LoadInst(valueType, valuePtr, (TString("val_") += ToString(idx)).c_str(), block);
+            };
+        };
+        return {complete, std::move(getters)};
+    }
+#endif
+private:
     void RegisterDependencies() const final {
         this->DependsOn(Stream);
     }
 
-    TState& GetState(NUdf::TUnboxedValue& state, TComputationContext& ctx) const {
-        if (!state.HasValue()) {
-            state = ctx.HolderFactory.Create<TState>(Stream->GetValue(ctx), Width);
-        }
-        return *static_cast<TState*>(state.AsBoxed().Get());
-    }
-
     IComputationNode* const Stream;
     const ui32 Width;
+    const ui32 TempStateIndex;
 };
 
 class TFromWideFlowWrapper : public TMutableComputationNode<TFromWideFlowWrapper> {
@@ -412,10 +440,9 @@ IComputationNode* WrapToFlow(TCallable& callable, const TComputationNodeFactoryC
     const auto outType = AS_TYPE(TFlowType, callable.GetType()->GetReturnType())->GetItemType();
     const auto kind = GetValueRepresentation(outType);
     if (type->IsStream()) {
-        auto streamType = AS_TYPE(TStreamType, type);
-        if (streamType->GetItemType()->IsMulti()) {
-            ui32 width = GetWideComponentsCount(streamType);
-            return new TToWideFlowWrapper(ctx.Mutables, LocateNode(ctx.NodeLocator, callable, 0), width);
+        if (const auto streamType = AS_TYPE(TStreamType, type); streamType->GetItemType()->IsMulti()) {
+            const auto multiType = AS_TYPE(TMultiType, streamType->GetItemType());
+            return new TToWideFlowWrapper(ctx.Mutables, LocateNode(ctx.NodeLocator, callable, 0), multiType->GetElementsCount());
         }
         return new TToFlowWrapper<true>(ctx.Mutables, kind, LocateNode(ctx.NodeLocator, callable, 0));
     } else if (type->IsList()) {
diff --git a/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.h b/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.h
index 9ab88085c8..fe9db52682 100644
--- a/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.h
+++ b/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.h
@@ -63,6 +63,8 @@ using TPairStateFlowCodegeneratorNode = TPairStateFlowComputationNode<TDerived>;
 template <typename TDerived>
 using TFlowSourceCodegeneratorNode = TFlowSourceComputationNode<TDerived>;
 
+template <typename TDerived>
+using TWideFlowSourceCodegeneratorNode = TWideFlowSourceComputationNode<TDerived>;
 }
 }
 #else
@@ -503,6 +505,23 @@ protected:
     }
 };
 
+template <typename TDerived>
+class TWideFlowSourceCodegeneratorNode: public TWideFlowSourceComputationNode<TDerived>, public ICodegeneratorInlineWideNode
+{
+using TBase = TWideFlowSourceComputationNode<TDerived>;
+protected:
+    TWideFlowSourceCodegeneratorNode(TComputationMutables& mutables,EValueRepresentation stateKind)
+        : TBase(mutables, stateKind)
+    {}
+
+    TGenerateResult GenGetValues(const TCodegenContext& ctx, BasicBlock*& block) const final {
+        auto& context = ctx.Codegen->GetContext();
+        const auto valueType = Type::getInt128Ty(context);
+        const auto statePtr = GetElementPtrInst::CreateInBounds(valueType, ctx.GetMutables(), {ConstantInt::get(Type::getInt32Ty(context), static_cast<const IComputationNode*>(this)->GetIndex())}, "state_ptr", block);
+        return static_cast<const TDerived*>(this)->DoGenGetValues(ctx, statePtr, block);
+    }
+};
+
 template <typename TDerived, bool SerializableState = false>
 class TStatefulFlowCodegeneratorNode: public TStatefulFlowComputationNode<TDerived, SerializableState>, public ICodegeneratorInlineNode
 {
diff --git a/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h b/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h
index d2f4471e3f..731b23f3c1 100644
--- a/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h
+++ b/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h
@@ -226,9 +226,8 @@ class TFlowSourceBaseComputationNode: public TStatefulComputationNode<IFlowInter
 {
     using TBase = TStatefulComputationNode<IFlowInterface>;
 protected:
-    TFlowSourceBaseComputationNode(TComputationMutables& mutables, EValueRepresentation kind, EValueRepresentation stateKind)
+    TFlowSourceBaseComputationNode(TComputationMutables& mutables, EValueRepresentation stateKind)
         : TBase(mutables, stateKind)
-        , RepresentationKind(kind)
     {}
 
     TString DebugString() const override {
@@ -248,8 +247,6 @@ protected:
             node->SetOwner(this);
         }
     }
-
-    const EValueRepresentation RepresentationKind;
 private:
     bool IsTemporaryValue() const final {
         return true;
@@ -276,24 +273,25 @@ private:
     mutable std::unordered_set<const IComputationNode*> Sources; // TODO: remove const and mutable.
 };
 
-
 template <typename TDerived>
 class TFlowSourceComputationNode: public TFlowSourceBaseComputationNode<TDerived, IComputationNode>
 {
     using TBase = TFlowSourceBaseComputationNode<TDerived, IComputationNode>;
 protected:
     TFlowSourceComputationNode(TComputationMutables& mutables, EValueRepresentation kind, EValueRepresentation stateKind)
-        : TBase(mutables, kind, stateKind)
+        : TBase(mutables, stateKind), RepresentationKind(kind)
     {}
 
 private:
     EValueRepresentation GetRepresentation() const final {
-        return this->RepresentationKind;
+        return RepresentationKind;
     }
 
     NUdf::TUnboxedValue GetValue(TComputationContext& compCtx) const final {
         return static_cast<const TDerived*>(this)->DoCalculate(this->ValueRef(compCtx), compCtx);
     }
+private:
+    const EValueRepresentation RepresentationKind;
 };
 
 template <typename TDerived>
@@ -302,9 +300,8 @@ class TWideFlowSourceComputationNode: public TFlowSourceBaseComputationNode<TDer
     using TBase = TFlowSourceBaseComputationNode<TDerived, IComputationWideFlowNode>;
 protected:
     TWideFlowSourceComputationNode(TComputationMutables& mutables, EValueRepresentation stateKind)
-        : TBase(mutables, EValueRepresentation::Any, stateKind)
+        : TBase(mutables, stateKind)
     {}
-
 private:
     EValueRepresentation GetRepresentation() const final {
         THROW yexception() << "Failed to get representation kind.";